def train(self, train_gen, valid_gen, cb):
if (self.cf.train_model):
losses = {"g": [], "mIoU":[]}
for i in range(self.cf.dataset.n_classes):
losses["mIoU"].append([])
n_epochs = self.cf.n_epochs
batch_size = self.cf.batch_size_train
#img_shape = self.model.img_shape
#n_classes = self.model.n_classes
img_shape = (256, 256, 3)
n_classes = self.cf.dataset.n_classes
save_path = self.cf.savepath
tag = 'train'
color_map = self.cf.dataset.color_map
saveResults = True
display_every_batches = 200
save_every_batches = 200
void_label = self.cf.dataset.void_class
train_it = Two_Image_Iterator('/datatmp/Datasets/segmentation/cityscapes/train/',batch_size=batch_size,
target_size=img_shape[:-1])
valid_it = Two_Image_Iterator('/datatmp/Datasets/segmentation/cityscapes/valid/',batch_size=batch_size,
target_size=img_shape[:-1])
generator_model = self.generator
while train_it.epochs_completed() < n_epochs:
input_image, gt_one_hot, y_gen = get_batch_for_generator(train_it, n_classes)
lg = generator_model.train_on_batch(input_image, gt_one_hot)
if train_it.total_batches_seen % display_every_batches == 0:
input_image_valid, gt_one_hot_valid, y_gen_valid = get_batch_for_generator(valid_it, n_classes)
pred_valid = generator_model.predict(input_image_valid)
mIoU = self.get_mIoU(pred_valid, gt_one_hot_valid)
losses["g"].append(lg)
for i in range(self.cf.dataset.n_classes):
losses["mIoU"][i].append(mIoU[i])
#np.append(losses["mIoU"], mIoU, axis=1)
print('epoch {}, batch {}, loss generator {}, mIoU {}'.\
format(train_it.epochs_completed(), train_it.total_batches_seen, lg, mIoU.tolist()))
self.plot_loss(losses)
if saveResults and train_it.total_batches_seen % save_every_batches == 0:
pred = generator_model.predict(input_image)
save_img3(input_image*255, gt_one_hot, pred, save_path, train_it.epochs_completed(),
color_map, n_classes, tag+str(train_it.epochs_completed()),
void_label)
def predict(self, test_gen, tag='pred'):
if self.cf.pred_model:
print('Predict method not implemented.')
return
# TODO fix model predict method
print('\n > Predicting the model...')
# Load best trained model
# self.model.load_weights(os.path.join(self.cf.savepath, "weights.hdf5"))
self.model.load_weights(self.cf.weights_file)
# Create a data generator
data_gen_queue, _stop, _generator_threads = GeneratorEnqueuer(
self.test_gen, max_q_size=1)
# Process the dataset
start_time = time.time()
for _ in range(
int(
math.ceil(self.cf.dataset.n_images_train /
float(self.cf.batch_size_test)))):
# Get data for this minibatch
data = data_gen_queue.get()
x_true = data[0]
y_true = data[1].astype('int32')
# Get prediction for this minibatch
y_pred = self.model.predict(x_true)
# Compute the argmax
y_pred = np.argmax(y_pred, axis=1)
# Reshape y_true
y_true = np.reshape(
y_true,
(y_true.shape[0], y_true.shape[2], y_true.shape[3]))
save_img3(x_true, y_true, y_pred, self.cf.savepath, 0,
self.cf.dataset.color_map, self.cf.dataset.classes,
tag + str(_), self.cf.dataset.void_class)
# Stop data generator
_stop.set()
total_time = time.time() - start_time
fps = float(self.cf.dataset.n_images_test) / total_time
s_p_f = total_time / float(self.cf.dataset.n_images_test)
print(' Predicting time: {}. FPS: {}. Seconds per Frame: {}'.
format(total_time, fps, s_p_f))
def predict(self, test_gen, tag='pred'):
if self.cf.pred_model:
# TODO model predict method for other tasks
print('\n > Predicting the model...')
if self.cf.problem_type == 'segmentation':
# Load best trained model
self.model.load_weights(self.cf.weights_file)
# Create output directory
if not os.path.exists(
os.path.join(self.cf.savepath, 'Predictions')):
os.makedirs(os.path.join(self.cf.savepath, 'Predictions'))
# Create a data generator
enqueuer = GeneratorEnqueuer(test_gen, wait_time=0.05)
enqueuer.start(workers=1, max_queue_size=1)
output_generator = enqueuer.get()
# Process the dataset
start_time = time.time()
for i in range(
int(
math.ceil(self.cf.dataset.n_images_test /
float(self.cf.batch_size_test)))):
# Get data for this minibatch
data = next(output_generator)
x_true = data[0]
y_true = data[1].astype('int32')
# Get prediction for this minibatch
y_pred = self.model.predict(x_true)
# Reshape y_true and compute the y_pred argmax
if K.image_dim_ordering() == 'th':
print(
'Predict method not implemented for th dim ordering.'
)
return
else:
# Find the most probable class of each pixel
y_pred = np.argmax(y_pred, axis=2)
y_true = np.argmax(y_true, axis=2)
# Reshape from (?,172800) to (?, 360, 480)
# print('Acc with shapes y_pred={} and y_true{}: {}'.format(y_pred.shape,y_true.shape,np.mean(np.ravel(y_pred==y_true))))
y_true = np.reshape(y_true,
(x_true.shape[0], x_true.shape[1],
x_true.shape[2]))
y_pred = np.reshape(y_pred,
(x_true.shape[0], x_true.shape[1],
x_true.shape[2]))
# print('Acc with shapes y_pred={} and y_true{}: {}'.format(y_pred.shape,y_true.shape,np.mean(np.ravel(y_pred==y_true))))
# Save output images
save_img3(x_true, y_true, y_pred,
os.path.join(self.cf.savepath, 'Predictions'), 0,
self.cf.dataset.color_map,
self.cf.dataset.classes, tag + str(i),
self.cf.dataset.void_class)
# Stop data generator
if enqueuer is not None:
enqueuer.stop()
else:
print('Predict method not implemented.')
return
total_time = time.time() - start_time
fps = float(self.cf.dataset.n_images_test) / total_time
s_p_f = total_time / float(self.cf.dataset.n_images_test)
print(' Predicting time: {}. FPS: {}. Seconds per Frame: {}'.
format(total_time, fps, s_p_f))
