def label_convertion(self, label_image):
# convert from raw images to pixel-wise labels
label_image = transform.resize(label_image, [self.hwc[0], self.hwc[1]])
label_image = img_as_ubyte(label_image)
tool0 = utils.toolbox()
label = tool0.convert_label(label_image, self.classnum)
return label
def train(max_epoch=10000, batch_size=2, restore_iter=0):
model = resnet50.resnet50(using_weights=True)
loader = data_loader.VOC_loader()
optimizer = tf.train.AdamOptimizer(1e-4)
with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)):
train_step = optimizer.minimize(model.loss)
saver = tf.train.Saver()
init = tf.global_variables_initializer()
sess = tf.Session()
sess.run(init)
batch_per_epoch = len(loader.trainlist) // batch_size
tool0 = utils.toolbox()
if restore_iter != 0:
saver.restore(sess, "./models/__iter__" + str(restore_iter) + ".ckpt")
print("restoring")
for e in range(restore_iter, max_epoch):
train_iterator = loader.train_next_batch(batch_size)
loss_avg = 0
for b in range(batch_per_epoch):
batch_x, batch_y = next(train_iterator)
# scaling
batch_x = scaling(batch_x)
sess.run(train_step,
feed_dict={
model.x: batch_x,
model.y: batch_y
})
batch_loss = sess.run(model.loss,
feed_dict={
model.x: batch_x,
model.y: batch_y
})
loss_avg += batch_loss
loss_avg /= batch_per_epoch
print("epoch" + str(e + 1) + ":" + str(loss_avg))
if not os.path.exists("./models"):
os.system("mkdir ./models")
if (e + 1) % 50 == 0:
saver.save(sess, "./models/__iter__" + str(e + 1) + ".ckpt")
pred_imgs, pred_labels = loader.prepare_val_batch([0, 10])
pred_imgs = scaling(pred_imgs)
for i in range(len(pred_imgs)):
pred_i = sess.run(model.logits,
feed_dict={model.x: [pred_imgs[i]]})
pred_i = pred_i[0]
prediction = np.argmax(pred_i, axis=-1)
pred_i = tool0.convert_image(pred_i, 22)
label_i = tool0.convert_image(pred_labels[i], 22)
if i < 6:
cv2.imshow("win0", pred_i)
cv2.imshow("win1", label_i.astype(np.uint8))
cv2.waitKey()
cv2.destroyAllWindows()
def label_convertion(file_nums):
# input the label image numbers and return converted one-hot label np.array
tool0 = utils.toolbox()
converted_labels = []
for label_num in file_nums:
label = cv2.imread(dataset_dir + label_dir + label_num + ".png")
label = cv2.resize(label, hw)
label = tool0.convert_label(label, 22)
converted_labels.append(label)
return converted_labels
def test_func(self):
img = skimage_read_BGR(self.root + "SegmentationClass/" +
self.trainlist[0] + ".png")
img = cv2.imread(self.root + "SegmentationClass/" + self.trainlist[0] +
".png")
img = self.preprocess(img)
tool0 = utils.toolbox()
label = tool0.convert_label(img, self.classnum)
#print(label)
label_image = tool0.convert_image(label, self.classnum)
cv2.imshow("win1", label_image.astype(np.uint8))
cv2.waitKey()
cv2.destroyAllWindows()
label_image = self.label_convertion(label_image)
batch_x.append(input_image)
batch_y.append(label_image)
return batch_x, batch_y
def prepare_train_batch(self, index=[0, 10]):
# input an [lower,upper] index range to get val images and labels from self.vallist
batch_x = []
batch_y = []
for name in self.trainlist[index[0]:index[1]]:
input_image = cv2.imread(self.root + "JPEGImages/" + name + ".jpg")
label_image = cv2.imread(self.root + "SegmentationClass/" + name +
".png")
input_image = self.preprocess(input_image)
label_image = self.label_convertion(label_image)
batch_x.append(input_image)
batch_y.append(label_image)
return batch_x, batch_y
loader = VOC_loader()
train = loader.train_next_batch(batch_size=2)
tool0 = utils.toolbox()
for i in range(3):
batch_x, batch_y = next(train)
label_img = tool0.convert_image(batch_y[0], 22)
cv2.imshow("image", batch_x[0])
cv2.imshow("label", label_img.astype(np.uint8))
cv2.waitKey()
cv2.destroyAllWindows()