-
Notifications
You must be signed in to change notification settings - Fork 0
/
data.py
123 lines (108 loc) · 5.2 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
from __future__ import print_function
from keras.preprocessing.image import ImageDataGenerator
import numpy as np
import os
import glob
import skimage.io as io
import skimage.transform as trans
import cv2 as cv
def adjustData(img,mask,flag_multi_class,num_class):
if(flag_multi_class):
img = img / 255
mask = mask[:,:,:,0] if(len(mask.shape) == 4) else mask[:,:,0]
new_mask = np.zeros(mask.shape + (num_class,))
for i in range(num_class):
#for one pixel in the image, find the class in mask and convert it into one-hot vector
#index = np.where(mask == i)
#index_mask = (index[0],index[1],index[2],np.zeros(len(index[0]),dtype = np.int64) + i) if (len(mask.shape) == 4) else (index[0],index[1],np.zeros(len(index[0]),dtype = np.int64) + i)
#new_mask[index_mask] = 1
new_mask[mask == i,i] = 1
new_mask = np.reshape(new_mask,(new_mask.shape[0],new_mask.shape[1]*new_mask.shape[2],new_mask.shape[3])) if flag_multi_class else np.reshape(new_mask,(new_mask.shape[0]*new_mask.shape[1],new_mask.shape[2]))
mask = new_mask
elif(np.max(img) > 1):
img = img / 255
mask = mask /255
mask[mask > 0.5] = 1
mask[mask <= 0.5] = 0
return (img,mask)
# for training
def trainGenerator(batch_size,train_path,image_folder,mask_folder,aug_dict,image_color_mode = "grayscale",
mask_color_mode = "grayscale",image_save_prefix = "image",mask_save_prefix = "mask",
flag_multi_class = False,num_class = 2,save_to_dir = None,target_size = (256,256),seed = 1):
image_datagen = ImageDataGenerator(**aug_dict)
mask_datagen = ImageDataGenerator(**aug_dict)
image_generator = image_datagen.flow_from_directory(train_path,classes = [image_folder],class_mode = None,color_mode = image_color_mode,
target_size = target_size,
batch_size = batch_size,
save_to_dir = save_to_dir,
save_prefix = image_save_prefix,
)
mask_generator = mask_datagen.flow_from_directory(
train_path,
classes = [mask_folder],
class_mode = None,
color_mode = mask_color_mode,
target_size = target_size,
batch_size = batch_size,
save_to_dir = save_to_dir,
save_prefix = mask_save_prefix,
)
train_generator = zip(image_generator, mask_generator)
for (img,mask) in train_generator:
img,mask = adjustData(img,mask,flag_multi_class,num_class)
yield (img,mask)
def testGenerator(test_path,num_image = 40,target_size = (256,256),flag_multi_class = False,as_gray = True):
for i in os.listdir(test_path):
# img =io.imread(os.path.join(test_path,i),as_gray = as_gray)
img = cv.imread(os.path.join(test_path, i), cv.IMREAD_GRAYSCALE)
# step-1 intensity contrast enhancement
clahe = cv.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
img = clahe.apply(img)
ret, thresh = cv.threshold(img, 0, 255, cv.THRESH_BINARY_INV+ cv.THRESH_OTSU)
kernel = np.ones((3, 3), np.uint8)
median = cv.medianBlur(thresh, 15)
closing = cv.morphologyEx(median, cv.MORPH_CLOSE, kernel, iterations=2)
# closing =closing / 255
img1 = trans.resize(closing,target_size)
img1 = np.reshape(img1,img1.shape+(1,)) if (not flag_multi_class) else img1
img1 = np.reshape(img1,(1,)+img1.shape)
yield img1
# def geneTrainNpy(image_path,mask_path,flag_multi_class = False,num_class = 2,image_prefix = "image",mask_prefix = "mask",image_as_gray = True,mask_as_gray = True):
# image_name_arr = glob.glob(os.path.join(image_path,"%s*.png"%image_prefix))
# image_arr = []
# mask_arr = []
# for index,item in enumerate(image_name_arr):
# img = io.imread(item,as_gray = image_as_gray)
# img = np.reshape(img,img.shape + (1,)) if image_as_gray else img
# mask = io.imread(item.replace(image_path,mask_path).replace(image_prefix,mask_prefix),as_gray = mask_as_gray)
# mask = np.reshape(mask,mask.shape + (1,)) if mask_as_gray else mask
# img,mask = adjustData(img,mask,flag_multi_class,num_class)
# image_arr.append(img)
# mask_arr.append(mask)
# image_arr = np.array(image_arr)
# mask_arr = np.array(mask_arr)
# return image_arr,mask_arr
S= [128,128,128]
B = [128,0,0]
P = [192,192,128]
R= [128,64,128]
P = [60,40,222]
T = [128,128,0]
S = [192,128,128]
F = [64,64,128]
C = [64,0,128]
P = [64,64,0]
B = [0,128,192]
U = [0,0,0]
COLOR_DICT = np.array([S, B, P, R, P,T, S, F, C, P, B, U])
def labelVisualize(num_class,color_dict,img):
img = img[:,:,0] if len(img.shape) == 3 else img
img_out = np.zeros(img.shape + (3,))
for i in range(num_class):
img_out[img == i,:] = color_dict[i]
return img_out / 255
def saveResult(save_path, npyfile, flag_multi_class=False, num_class=1):
for i, item in enumerate(npyfile):
img = labelVisualize(num_class, COLOR_DICT, item) if flag_multi_class else item[:, :, 0]
# img=np.resize((img,np.uint8()))
io.imsave(os.path.join(save_path, "%d_predict.png" % i), img)