def main(argv=None): keep_probability = tf.placeholder(tf.float32, name="keep_probabilty") image = tf.placeholder(tf.float32, shape=[None, IMAGE_SIZE_HEIGHT, IMAGE_SIZE_WIDTH, 3], name="input_image") annotation = tf.placeholder(tf.float32, shape=[None, IMAGE_SIZE_HEIGHT, IMAGE_SIZE_WIDTH, 2], name="annotation") pred_annotation, logits = inference(image, keep_probability) # build the FCN graph prob = tf.nn.softmax(logits) logits_ = tf.reshape(logits, [1, IMAGE_SIZE_HEIGHT*IMAGE_SIZE_WIDTH, 2]) annotation_ = tf.reshape(annotation, [1, IMAGE_SIZE_HEIGHT*IMAGE_SIZE_WIDTH, 2]) loss = utils.focal_loss(logits_, annotation_) #loss = tf.reduce_mean((tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, # labels=tf.squeeze(annotation, axis=3), # name="entropy"))) #tf.summary.image('image', image) #tf.summary.image('water_gt', tf.cast(annotation, tf.float32)) #tf.summary.image('water_pred', tf.expand_dims(tf.cast(pred_annotation, tf.float32), axis=3)) #tf.summary.scalar('loss', loss) #merged = tf.summary.merge_all() #summary_writer = tf.summary.FileWriter(FLAGS.logs_dir, graph=tf.get_default_graph()) trainable_var = tf.trainable_variables() train_op = train(loss, trainable_var) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True sess = tf.Session(config = config) print("Setting up Saver...") saver = tf.train.Saver() sess.run(tf.global_variables_initializer()) ckpt = tf.train.get_checkpoint_state(FLAGS.logs_dir) if ckpt and ckpt.model_checkpoint_path and (FLAGS.mode=="visualize" or FLAGS.mode == "train-continue"): saver.restore(sess, ckpt.model_checkpoint_path) print("Model restored...") if FLAGS.mode == "train" or FLAGS.mode == "train-continue" : print("Loading data") global p; p = load_training_dataset_path() #train_imgs, gt_imgs = load_data(FLAGS.data_dir) #gt_imgs = np.expand_dims(gt_imgs, axis=3) print("Start Training...") for itr in xrange(MAX_ITERATION): print('Setp: %d'%(itr)) p1 = next_batch(1, itr) # print("p1:\n",p1) x_train, y_train, y_train2 = load_data(p1, itr) y_train = np.expand_dims(y_train, axis=3) y_train2 = np.expand_dims(y_train2, axis=3) label_in = np.concatenate((y_train2, y_train), axis=3) #x_train, y_train = next_batch(train_imgs, gt_imgs, FLAGS.batch_size) feed_dict = {image: x_train, annotation: label_in, keep_probability: 0.85} print("train feed_dict done!") start_time = time.time() sess.run(train_op, feed_dict = feed_dict) dur = time.time() - start_time print("dur : %f"% dur) if itr % 1 == 0: train_loss = sess.run(loss, feed_dict = feed_dict) print('KITTI Step: %d, Train_loss:%g'%(itr, train_loss)) #if itr % 100 == 0: # summary = sess.run(merged, feed_dict = feed_dict) # summary_writer.add_summary(summary, itr) if itr % 1000 == 0: print('Save Net Model...') saver.save(sess, FLAGS.logs_dir + "model.ckpt", itr) if itr % 1000 == 0 and itr >= 20000: FLAGS.learning_rate = FLAGS.learning_rate / 2 elif FLAGS.mode == "visualize": p = np.genfromtxt('../Dataset/on_road_test.txt', dtype='str') for idx in range(0,p.shape[0]): test_images1 = load_test_data(p[idx,0]) start_time = time.time() likelyhood, pred = sess.run([prob, pred_annotation], feed_dict={image:test_images1, keep_probability: 1.0}) dur = time.time() - start_time print("dur = ",dur) #np.save('./likelyhood/test/likelyhood_%06d'%(idx), likelyhood) #print(pred.shape) for itr in range(pred.shape[0]): utils.save_image(pred[itr].astype(np.float32), FLAGS.output_dir, name="pred_test_bin_%06d"%(idx)) print("Saved image :%d"%(idx))
def main(argv=None): keep_probability = tf.compat.v1.placeholder(tf.float32, name="keep_probabilty") image = tf.compat.v1.placeholder( tf.float32, shape=[None, IMAGE_SIZE_HEIGHT, IMAGE_SIZE_WIDTH, 3], name="input_image") annotation = tf.compat.v1.placeholder( tf.float32, shape=[None, IMAGE_SIZE_HEIGHT, IMAGE_SIZE_WIDTH, 2], name="annotation") pred_annotation, logits = inference( image, keep_probability) # build the FCN graph prob = tf.nn.softmax(logits) logits_ = tf.reshape(logits, [1, IMAGE_SIZE_HEIGHT * IMAGE_SIZE_WIDTH, 2]) annotation_ = tf.reshape(annotation, [1, IMAGE_SIZE_HEIGHT * IMAGE_SIZE_WIDTH, 2]) loss = utils.focal_loss(logits_, annotation_) #loss = tf.reduce_mean((tf.nn.sparse_softmax_cross_entropy_with_logits(logits=logits, # labels=tf.squeeze(annotation, axis=3), # name="entropy"))) #tf.summary.image('image', image) #tf.summary.image('water_gt', tf.cast(annotation, tf.float32)) #tf.summary.image('water_pred', tf.expand_dims(tf.cast(pred_annotation, tf.float32), axis=3)) #tf.summary.scalar('loss', loss) #merged = tf.summary.merge_all() #summary_writer = tf.summary.FileWriter(FLAGS.logs_dir, graph=tf.get_default_graph()) trainable_var = tf.trainable_variables() train_op = train(loss, trainable_var) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True sess = tf.Session(config=config) print("Setting up Saver...") saver = tf.train.Saver() sess.run(tf.global_variables_initializer()) print("TungNV_FLAGS.logs_dir: ", FLAGS.logs_dir) ckpt = load_model(saver, sess) if FLAGS.mode == "train": print("Loading data") global p p = load_training_dataset_path() #train_imgs, gt_imgs = load_data(FLAGS.data_dir) #gt_imgs = np.expand_dims(gt_imgs, axis=3) step = check_step_train(ckpt.model_checkpoint_path) print("Start Training...") for itr in xrange(step, MAX_ITERATION): print('Step: %d' % (itr)) p1 = next_batch(1, itr) print("p1:\n", p1) x_train, y_train, y_train2 = load_data(p1, itr) y_train = np.expand_dims(y_train, axis=3) y_train2 = np.expand_dims(y_train2, axis=3) label_in = np.concatenate((y_train2, y_train), axis=3) #x_train, y_train = next_batch(train_imgs, gt_imgs, FLAGS.batch_size) feed_dict = { image: x_train, annotation: label_in, keep_probability: 0.85 } print("train feed_dict done!") start_time = time.time() print("Run...") sess.run(train_op, feed_dict=feed_dict) dur = time.time() - start_time print("Time to run : %f" % dur) if itr % 1 == 0: train_loss = sess.run(loss, feed_dict=feed_dict) print('KITTI Step: %d, Train_loss:%g' % (itr, train_loss)) #if itr % 100 == 0: # summary = sess.run(merged, feed_dict = feed_dict) # summary_writer.add_summary(summary, itr) if itr % 5000 == 0: print('Save Net Model...') saver.save(sess, FLAGS.logs_dir + "model.ckpt", itr) if itr % 5000 == 0 and itr >= 20000: FLAGS.learning_rate = FLAGS.learning_rate / 2 elif FLAGS.mode == "test": # test_image_from_dataset(PATH_DATASET_TEST, keep_probability, image, pred_annotation, prob, sess) test_image_from_folder(PATH_IMAGE, keep_probability, image, pred_annotation, prob, sess)