x, y = int(x), int(y) img = cv2.imread(img_name) label = np.array([x, y], dtype=np.float32) return img[np.newaxis, :, :, :], label.reshape((1, label.shape[0])) if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('-g', '--gpu', default=None, type=int) args = parser.parse_args() if args is not None: os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu) net = JumpModel() dataset = JumpData() img = tf.placeholder(tf.float32, [None, 640, 720, 3], name='img') label = tf.placeholder(tf.float32, [None, 2], name='label') is_training = tf.placeholder(np.bool, name='is_training') keep_prob = tf.placeholder(np.float32, name='keep_prob') lr = tf.placeholder(np.float32, name='lr') pred = net.forward(img, is_training, keep_prob, 'coarse') loss = tf.reduce_mean(tf.sqrt(tf.pow(pred - label, 2) + 1e-12)) tf.summary.scalar('loss', loss) optimizer = tf.train.AdamOptimizer(lr) update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) with tf.control_dependencies(update_ops): train_op = optimizer.minimize(loss) saver = tf.train.Saver()
x, y = name[name.index('_h_') + 3: name.index('_h_') + 6], name[name.index('_w_') + 3: name.index('_w_') + 6] x, y = int(x), int(y) img = cv2.imread(img_name) label = np.array([x, y], dtype=np.float32) return img[np.newaxis, :, :, :], label.reshape((1, label.shape[0])) if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('-g', '--gpu', default=None, type=int) args = parser.parse_args() if args is not None: os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu) net = JumpModel() dataset = JumpData() img = tf.placeholder(tf.float32, [None, 640, 720, 3], name='img') label = tf.placeholder(tf.float32, [None, 2], name='label') is_training = tf.placeholder(np.bool, name='is_training') keep_prob = tf.placeholder(np.float32, name='keep_prob') lr = tf.placeholder(np.float32, name='lr') pred = net.forward(img, is_training, keep_prob, 'coarse') loss = tf.reduce_mean(tf.sqrt(tf.pow(pred - label, 2) + 1e-12)) tf.summary.scalar('loss', loss) optimizer = tf.train.AdamOptimizer(lr) update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) with tf.control_dependencies(update_ops): train_op = optimizer.minimize(loss) saver = tf.train.Saver()
lr = tf.placeholder(np.float32, name='lr') pred = net.forward(img, is_training, keep_prob) saver = tf.train.Saver() sess = tf.Session() sess.run(tf.global_variables_initializer()) ckpt = tf.train.get_checkpoint_state('./train_logs') if ckpt and ckpt.model_checkpoint_path: saver.restore(sess, ckpt.model_checkpoint_path) print('==== successfully restored ====') val_img = get_a_img(path) feed_dict = { img: val_img, is_training: False, keep_prob: 1.0, } pred_out = sess.run(pred, feed_dict=feed_dict) return pred_out if __name__ == '__main__': dataset = JumpData() name = dataset.val_name_list[0] posi = name.index('_res') img_name = name[:posi] + '.png' a = time.time() pred = inference(img_name) print(pred, time.time() - a)