Download UnFlow/src/eval_gui.py from doanh25032004/models_animerun: direct link, hf CLI and curl.
- Browser
- Download file 15 kB
-
https://huggingface.co/doanh25032004/models_animerun/resolve/main/UnFlow/src/eval_gui.py
- Command line
-
hf download hf://doanh25032004/models_animerun/UnFlow/src/eval_gui.py
-
curl -L -o eval_gui.py https://huggingface.co/doanh25032004/models_animerun/resolve/main/UnFlow/src/eval_gui.py
15 kB
| 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() | |