import os import sys import shutil import tensorflow as tf import numpy as np import png from e2eflow.core.flow_util import flow_to_color, flow_error_avg, outlier_pct from e2eflow.core.flow_util import flow_error_image from e2eflow.util import config_dict from e2eflow.core.image_warp import image_warp from e2eflow.kitti.input import KITTIInput from e2eflow.kitti.data import KITTIData from e2eflow.chairs.data import ChairsData from e2eflow.chairs.input import ChairsInput from e2eflow.sintel.data import SintelData from e2eflow.sintel.input import SintelInput from e2eflow.middlebury.input import MiddleburyInput from e2eflow.middlebury.data import MiddleburyData from e2eflow.core.unsupervised import unsupervised_loss from e2eflow.core.input import resize_input, resize_output_crop, resize_output, resize_output_flow from e2eflow.core.train import restore_networks from e2eflow.ops import forward_warp from e2eflow.gui import display from e2eflow.core.losses import DISOCC_THRESH, occlusion, create_outgoing_mask from e2eflow.util import convert_input_strings tf.app.flags.DEFINE_string('dataset', 'kitti', 'Name of dataset to evaluate on. One of {kitti, sintel, chairs, mdb}.') tf.app.flags.DEFINE_string('variant', 'train_2012', 'Name of variant to evaluate on.' 'If dataset = kitti, one of {train_2012, train_2015, test_2012, test_2015}.' 'If dataset = sintel, one of {train_clean, train_final}.' 'If dataset = mdb, one of {train, test}.') tf.app.flags.DEFINE_string('ex', '', 'Experiment name(s) (can be comma separated list).') tf.app.flags.DEFINE_integer('num', 10, 'Number of examples to evaluate. Set to -1 to evaluate all.') tf.app.flags.DEFINE_integer('num_vis', 100, 'Number of evalutations to visualize. Set to -1 to visualize all.') tf.app.flags.DEFINE_string('gpu', '0', 'GPU device to evaluate on.') tf.app.flags.DEFINE_boolean('output_benchmark', False, 'Output raw flow files.') tf.app.flags.DEFINE_boolean('output_visual', False, 'Output flow visualization files.') tf.app.flags.DEFINE_boolean('output_backward', False, 'Output backward flow files.') tf.app.flags.DEFINE_boolean('output_png', True, # TODO finish .flo output 'Raw output format to use with output_benchmark.' 'Outputs .png flow files if true, output .flo otherwise.') FLAGS = tf.app.flags.FLAGS NUM_EXAMPLES_PER_PAGE = 4 def write_rgb_png(z, path, bitdepth=8): z = z[0, :, :, :] with open(path, 'wb') as f: writer = png.Writer(width=z.shape[1], height=z.shape[0], bitdepth=bitdepth) z2list = z.reshape(-1, z.shape[1]*z.shape[2]).tolist() writer.write(f, z2list) def flow_to_int16(flow): _, h, w, _ = tf.unstack(tf.shape(flow)) u, v = tf.unstack(flow, num=2, axis=3) r = tf.cast(tf.maximum(0.0, tf.minimum(u * 64.0 + 32768.0, 65535.0)), tf.uint16) g = tf.cast(tf.maximum(0.0, tf.minimum(v * 64.0 + 32768.0, 65535.0)), tf.uint16) b = tf.ones([1, h, w], tf.uint16) return tf.stack([r, g, b], axis=3) def write_flo(flow, filename): """ write optical flow in Middlebury .flo format :param flow: optical flow map :param filename: optical flow file path to be saved :return: None """ flow = flow[0, :, :, :] f = open(filename, 'wb') magic = np.array([202021.25], dtype=np.float32) height, width = flow.shape[:2] magic.tofile(f) np.int32(width).tofile(f) np.int32(height).tofile(f) data = np.float32(flow).flatten() data.tofile(f) f.close() def _evaluate_experiment(name, input_fn, data_input): normalize_fn = data_input._normalize_image resized_h = data_input.dims[0] resized_w = data_input.dims[1] current_config = config_dict('../config.ini') exp_dir = os.path.join(current_config['dirs']['log'], 'ex', name) config_path = os.path.join(exp_dir, 'config.ini') if not os.path.isfile(config_path): config_path = '../config.ini' if not os.path.isdir(exp_dir) or not tf.train.get_checkpoint_state(exp_dir): exp_dir = os.path.join(current_config['dirs']['checkpoints'], name) config = config_dict(config_path) params = config['train'] convert_input_strings(params, config_dict('../config.ini')['dirs']) dataset_params_name = 'train_' + FLAGS.dataset if dataset_params_name in config: params.update(config[dataset_params_name]) ckpt = tf.train.get_checkpoint_state(exp_dir) if not ckpt: raise RuntimeError("Error: experiment must contain a checkpoint") ckpt_path = exp_dir + "/" + os.path.basename(ckpt.model_checkpoint_path) with tf.Graph().as_default(): #, tf.device('gpu:' + FLAGS.gpu): inputs = input_fn() im1, im2, input_shape = inputs[:3] truth = inputs[3:] height, width, _ = tf.unstack(tf.squeeze(input_shape), num=3, axis=0) im1 = resize_input(im1, height, width, resized_h, resized_w) im2 = resize_input(im2, height, width, resized_h, resized_w) # TODO adapt train.py _, flow, flow_bw = unsupervised_loss( (im1, im2), normalization=data_input.get_normalization(), params=params, augment=False, return_flow=True) im1 = resize_output(im1, height, width, 3) im2 = resize_output(im2, height, width, 3) flow = resize_output_flow(flow, height, width, 2) flow_bw = resize_output_flow(flow_bw, height, width, 2) flow_fw_int16 = flow_to_int16(flow) flow_bw_int16 = flow_to_int16(flow_bw) im1_pred = image_warp(im2, flow) im1_diff = tf.abs(im1 - im1_pred) #im2_diff = tf.abs(im1 - im2) #flow_bw_warped = image_warp(flow_bw, flow) if len(truth) == 4: flow_occ, mask_occ, flow_noc, mask_noc = truth flow_occ = resize_output_crop(flow_occ, height, width, 2) flow_noc = resize_output_crop(flow_noc, height, width, 2) mask_occ = resize_output_crop(mask_occ, height, width, 1) mask_noc = resize_output_crop(mask_noc, height, width, 1) #div = divergence(flow_occ) #div_bw = divergence(flow_bw) occ_pred = 1 - (1 - occlusion(flow, flow_bw)[0]) def_pred = 1 - (1 - occlusion(flow, flow_bw)[1]) disocc_pred = forward_warp(flow_bw) < DISOCC_THRESH disocc_fw_pred = forward_warp(flow) < DISOCC_THRESH image_slots = [((im1 * 0.5 + im2 * 0.5) / 255, 'overlay'), (im1_diff / 255, 'brightness error'), #(im1 / 255, 'first image', 1, 0), #(im2 / 255, 'second image', 1, 0), #(im2_diff / 255, '|first - second|', 1, 2), (flow_to_color(flow), 'flow'), #(flow_to_color(flow_bw), 'flow bw prediction'), #(tf.image.rgb_to_grayscale(im1_diff) > 20, 'diff'), #(occ_pred, 'occ'), #(def_pred, 'disocc'), #(disocc_pred, 'reverse disocc'), #(disocc_fw_pred, 'forward disocc prediction'), #(div, 'div'), #(div < -2, 'neg div'), #(div > 5, 'pos div'), #(flow_to_color(flow_occ, mask_occ), 'flow truth'), (flow_error_image(flow, flow_occ, mask_occ, mask_noc), 'flow error') # (blue: correct, red: wrong, dark: occluded) ] # list of (scalar_op, title) scalar_slots = [(flow_error_avg(flow_noc, flow, mask_noc), 'EPE_noc'), (flow_error_avg(flow_occ, flow, mask_occ), 'EPE_all'), (outlier_pct(flow_noc, flow, mask_noc), 'outliers_noc'), (outlier_pct(flow_occ, flow, mask_occ), 'outliers_all')] elif len(truth) == 2: flow_gt, mask = truth flow_gt = resize_output_crop(flow_gt, height, width, 2) mask = resize_output_crop(mask, height, width, 1) image_slots = [((im1 * 0.5 + im2 * 0.5) / 255, 'overlay'), (im1_diff / 255, 'brightness error'), (flow_to_color(flow), 'flow'), (flow_to_color(flow_gt, mask), 'gt'), ] # list of (scalar_op, title) scalar_slots = [(flow_error_avg(flow_gt, flow, mask), 'EPE_all')] else: image_slots = [(im1 / 255, 'first image'), #(im1_pred / 255, 'warped second image', 0, 1), (im1_diff / 255, 'warp error'), #(im2 / 255, 'second image', 1, 0), #(im2_diff / 255, '|first - second|', 1, 2), (flow_to_color(flow), 'flow prediction')] scalar_slots = [] num_ims = len(image_slots) image_ops = [t[0] for t in image_slots] scalar_ops = [t[0] for t in scalar_slots] image_names = [t[1] for t in image_slots] scalar_names = [t[1] for t in scalar_slots] all_ops = image_ops + scalar_ops image_lists = [] averages = np.zeros(len(scalar_ops)) sess_config = tf.ConfigProto(allow_soft_placement=True) exp_out_dir = os.path.join('../out', name) if FLAGS.output_visual or FLAGS.output_benchmark: if os.path.isdir(exp_out_dir): shutil.rmtree(exp_out_dir) os.makedirs(exp_out_dir) shutil.copyfile(config_path, os.path.join(exp_out_dir, 'config.ini')) with tf.Session(config=sess_config) as sess: saver = tf.train.Saver(tf.global_variables()) sess.run(tf.global_variables_initializer()) sess.run(tf.local_variables_initializer()) restore_networks(sess, params, ckpt, ckpt_path) coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(sess=sess, coord=coord) # TODO adjust for batch_size > 1 (also need to change image_lists appending) max_iter = FLAGS.num if FLAGS.num > 0 else None try: num_iters = 0 while not coord.should_stop() and (max_iter is None or num_iters != max_iter): all_results = sess.run([flow, flow_bw, flow_fw_int16, flow_bw_int16] + all_ops) flow_fw_res, flow_bw_res, flow_fw_int16_res, flow_bw_int16_res = all_results[:4] all_results = all_results[4:] image_results = all_results[:num_ims] scalar_results = all_results[num_ims:] iterstr = str(num_iters).zfill(6) if FLAGS.output_visual: path_col = os.path.join(exp_out_dir, iterstr + '_flow.png') path_overlay = os.path.join(exp_out_dir, iterstr + '_img.png') path_error = os.path.join(exp_out_dir, iterstr + '_err.png') write_rgb_png(image_results[0] * 255, path_overlay) write_rgb_png(image_results[1] * 255, path_col) write_rgb_png(image_results[2] * 255, path_error) if FLAGS.output_benchmark: path_fw = os.path.join(exp_out_dir, iterstr) if FLAGS.output_png: write_rgb_png(flow_fw_int16_res, path_fw + '_10.png', bitdepth=16) else: write_flo(flow_fw_res, path_fw + '_10.flo') if FLAGS.output_backward: path_fw = os.path.join(exp_out_dir, iterstr + '_01.png') write_rgb_png(flow_bw_int16_res, path_bw, bitdepth=16) if num_iters < FLAGS.num_vis: image_lists.append(image_results) averages += scalar_results if num_iters > 0: sys.stdout.write('\r') num_iters += 1 sys.stdout.write("-- evaluating '{}': {}/{}" .format(name, num_iters, max_iter)) sys.stdout.flush() print() except tf.errors.OutOfRangeError: pass averages /= num_iters coord.request_stop() coord.join(threads) for t, avg in zip(scalar_slots, averages): _, scalar_name = t print("({}) {} = {}".format(name, scalar_name, avg)) return image_lists, image_names def main(argv=None): os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu print("-- evaluating: on {} pairs from {}/{}" .format(FLAGS.num, FLAGS.dataset, FLAGS.variant)) default_config = config_dict() dirs = default_config['dirs'] if FLAGS.dataset == 'kitti': data = KITTIData(dirs['data'], development=True) data_input = KITTIInput(data, batch_size=1, normalize=False, dims=(384,1280)) inputs = getattr(data_input, 'input_' + FLAGS.variant)() elif FLAGS.dataset == 'chairs': data = ChairsData(dirs['data'], development=True) data_input = ChairsInput(data, batch_size=1, normalize=False, dims=(384,512)) if FLAGS.variant == 'test_2015' and FLAGS.num == -1: FLAGS.num = 200 elif FLAGS.variant == 'test_2012' and FLAGS.num == -1: FLAGS.num = 195 elif FLAGS.dataset == 'sintel': data = SintelData(dirs['data'], development=True) data_input = SintelInput(data, batch_size=1, normalize=False, dims=(512,1024)) if FLAGS.variant in ['test_clean', 'test_final'] and FLAGS.num == -1: FLAGS.num = 552 elif FLAGS.dataset == 'mdb': data = MiddleburyData(dirs['data'], development=True) data_input = MiddleburyInput(data, batch_size=1, normalize=False, dims=(512,640)) if FLAGS.variant == 'test' and FLAGS.num == -1: FLAGS.num = 12 input_fn = getattr(data_input, 'input_' + FLAGS.variant) results = [] for name in FLAGS.ex.split(','): result, image_names = _evaluate_experiment(name, input_fn, data_input) results.append(result) display(results, image_names) if __name__ == '__main__': tf.app.run()