from flow import Flow import os import sys test_dir = os.path.abspath(os.path.join(__file__, "../test")) # checkpoint = os.path.abspath(os.path.join(__file__, "../checkpoints/flownet2")) checkpoint = "flownet2" flow = Flow(test_dir, test_dir) skip = sys.argv[1] if len(sys.argv) > 1 else 5 input_files = filter( lambda name: os.path.isfile(os.path.join(test_dir, "color_flow", name)), os.listdir(os.path.join(test_dir, "color_flow"))) indices = sorted( [int(os.path.splitext(name)[0].split("_")[-1]) for name in input_files]) input_frames = list(filter(lambda i: i % skip == 0, indices)) total_frames = len(input_frames) print(input_frames) print("processing {} frames".format(total_frames)) index_pairs = list(zip(input_frames[:-1], input_frames[1:])) visualize = True flow.compute_flow(index_pairs, checkpoint, visualize)