-
Notifications
You must be signed in to change notification settings - Fork 0
/
data.py
138 lines (105 loc) · 4.15 KB
/
data.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
from keras.preprocessing.image import img_to_array, load_img, array_to_img,ImageDataGenerator
import numpy as np
import os
train_num = 30
test_num = 30
x_train = np.ndarray((train_num,512,512,3), dtype=np.uint8)
y_train = np.ndarray((train_num,512,512,1), dtype=np.uint8)
x_test = np.ndarray((test_num,512,512,1), dtype=np.uint8)
y_test = np.ndarray((test_num,512,512,1), dtype=np.uint8)
count = 0
print('-'*30)
print("loading training data...")
dir = "D:\dataset\data-train\img"
for root, dirs, files in os.walk(dir):
for file in files:
name = os.path.join(root,file).split("D:\\dataset\\data-train\\img\\")[1].split(".")[0]
img = load_img(dir+"\\"+name+".tif",target_size=(512,512))
label = load_img("D:\dataset\data-train\label\\"+name+".tif",grayscale = True,target_size=(512,512))
img = img_to_array(img)
label = img_to_array(label)
x_train[count] = img
y_train[count] = label
#print(x_train)
count += 1
print('complete')
print('-'*30)
print('start augmenting...')
#将标签作为训练集的一个通道
x_train[:,:,:,2] = y_train[:,:,:,0]
# print(x_train[:,:,:,2])
#np.save('x_train.npy',x_train)
#np.save('y_train.npy',y_train)
for count in range(train_num):
img = array_to_img(x_train[count])
img.save("D:\dataset\data-augmentation\\raw\\"+str(count)+".tif")
img_generator = ImageDataGenerator(rotation_range=0.2,
width_shift_range=0.05,
height_shift_range=0.05,
shear_range=0.05,
zoom_range=0.05,
horizontal_flip=True,
fill_mode='nearest')
#x_train = np.load("x_train.npy")
#print(x_train[:,:,2])
i = 0 # 数据增强
for num in range(train_num):
for batch in img_generator.flow(x_train,batch_size=1,
save_to_dir="D:\dataset\data-augmentation\\unsplited",
save_prefix="train",
save_format='tif'
):
i += 1
if i >= 50:
i = 0
break
dir ="D:\dataset\data-augmentation\\unsplited"
img_num = 0
for num in os.listdir(dir): #fn 表示的是文件名
img_num += 1
new_img = np.ndarray((img_num,512,512,1), dtype=np.uint8)
new_label = np.ndarray((img_num,512,512,1), dtype=np.uint8)
count = 0
for root, dirs, files in os.walk(dir):
for file in files:
# name = os.path.join(root,file).split("D:\\dataset\\data-train\\img\\")[1].split(".")[0]
# img = load_img(dir+"\\"+name+".bmp",target_size=(128,128))
# label = load_img("D:\dataset\data-train\label\\"+name+"_anno.bmp",grayscale = True,target_size=(128,128))
img = load_img(os.path.join(root,file))
img = img_to_array(img)
label = img[:,:,2]
label = label.reshape((512,512,1))
new_label[count] = label
label[label != 0] = 255
label = array_to_img(label)
label.save("D:\dataset\data-augmentation\label\\"+str(count)+".tif")
img[:,:,2] = 0
img = array_to_img(img)
img.save("D:\dataset\data-augmentation\img\\"+str(count)+".tif")
img = load_img("D:\dataset\data-augmentation\img\\"+str(count)+".tif",grayscale=True)
img = img_to_array(img)
new_img[count] = img
#print(new_img[count].shape)
count += 1
#np.save('x_train.npy',new_img)
#np.save('y_train.npy',new_label)
np.save('x_train.npy',new_img)
np.save('y_train.npy',new_label)
count = 0
print("loading test data...")
dir = "D:\dataset\data-test\img"
for root, dirs, files in os.walk(dir):
for file in files:
name = os.path.join(root,file).split("D:\\dataset\\data-test\\img\\")[1].split(".")[0]
img = load_img(dir+"\\"+name+".tif",grayscale = True,target_size=(512,512))
#label = load_img("D:\dataset\data-test\label\\"+name+"_anno.bmp",grayscale = True,target_size=(512,512))
img = img_to_array(img)
#label = img_to_array(label)
#label[label != 0] = 255
x_test[count] = img
#y_test[count] = label
count += 1
np.save('x_test.npy',x_test)
#np.save('y_test.npy',y_test)
print('complete')
print('-'*30)