class TrainCheck(Callback):
def __init__(self, output_path, model_name, img_shape, nb_class):
self.epoch = 0
self.output_path = output_path
self.model_name = model_name
self.img_shape = img_shape
self.nb_class = nb_class
self.palette = VOCPalette(nb_class=nb_class)
def result_map_to_img(self, res_map):
res_map = np.squeeze(res_map)
argmax_idx = np.argmax(res_map, axis=2).astype('uint8')
return argmax_idx
def on_epoch_end(self, epoch, logs={}):
self.epoch = epoch + 1
# self.visualize(os.path.join(self.output_path,'test.png'))
def visualize(self, path):
imgorg = Image.open(path).convert('RGB')
img = imgorg.resize((self.img_shape[1], self.img_shape[0]),
Image.ANTIALIAS)
img_arr = np.array(img)
img_arr = img_arr / 127.5 - 1
img_arr = np.expand_dims(img_arr, 0)
pred = self.model.predict(img_arr)
res_img = self.result_map_to_img(pred[0])
PIL_img_pal = self.palette.genlabelpal(res_img)
PIL_img_pal = PIL_img_pal.resize((imgorg.size[0], imgorg.size[1]),
Image.ANTIALIAS)
PIL_img_pal.save(
os.path.join(
self.output_path,
self.model_name + '_epoch_' + str(self.epoch) + '.png'))
# resize the input image to the input layer's size
img = imgorg.resize((img_width, img_height), Image.ANTIALIAS)
# convert to numpy array
img_arr = np.array(img)
# Centering helps normalization image (-1 ~ 1 value)
img_arr = img_arr / 127.5 - 1
# batch size is set to one
img_arr = np.expand_dims(img_arr, 0)
# img_arr = np.expand_dims(img_arr, 3)
# go forward of the network
pred = model.predict(img_arr)
# compute the maximum axis number in probability tensor, which represents the label assigned
# res is the final label map
res = result_map_to_img(pred[0])
# print color on label map use palette
PIL_img_pal = palette.genlabelpal(res)
# resize the result to the original size
PIL_img_pal = PIL_img_pal.resize((imgorg.size[0], imgorg.size[1]), Image.ANTIALIAS)
obj = findObject(np.array(PIL_img_pal))
# plot figures
plt.ion()
plt.figure('Unet test')
plt.suptitle(img_path)
plt.subplot(1, 3, 1), plt.title('org')
plt.imshow(imgorg), plt.axis('off')
plt.subplot(1, 3, 2), plt.title(obj)
plt.imshow(PIL_img_pal), plt.axis('off')
plt.subplot(1, 3, 3), plt.title('label')
plt.imshow(imglab), plt.axis('off')
_ = np.zeros(B_SIZE)
seed = random.randrange(1, 1000)
x_tmp_gen = x_data_gen.flow(np.array(x), _, batch_size=B_SIZE, seed=seed)
y_tmp_gen = y_data_gen.flow(np.array(y), _, batch_size=B_SIZE, seed=seed)
# Finally, yield x, y data.
for j in range(GEN_NUM):
x_result, _ = next(x_tmp_gen)
y_result, _ = next(y_tmp_gen)
x_res = x_result[0]
y_res = y_result[0]
y_res = np.squeeze(y_res, axis=2)
img = image.array_to_img(x_res)
lab_arr = y_res.astype('uint8')
lab = palette.genlabelpal(lab_arr)
X_res_path = gen_img_path + "{}".format(
namesimg[i]) + '_' + str(j) + ".png"
Y_res_path = gen_lab_path + "{}".format(
namesimg[i]) + '_' + str(j) + ".png"
img.save(X_res_path)
lab.save(Y_res_path)
if i < val_num:
f_val.writelines(namesimg[i] + '_' + str(j) + "\n")
else:
f_train.writelines(namesimg[i] + '_' + str(j) + "\n")
# img.show()
# lab.show()
f_train.close()
f_val.close()