def process(args):
N = len(args.im1)
assert N == len(args.im2) and N == len(args.out)
device = torch.device("cuda:0")
Flownet = FlowNet2(args)
print(f"Loading pretrained model from '{args.pretrained_model_flownet2}'.")
flownet2_ckpt = torch.load(args.pretrained_model_flownet2)
Flownet.load_state_dict(flownet2_ckpt["state_dict"])
Flownet.to(device)
Flownet.eval()
for im1, im2, out in zip(args.im1, args.im2, args.out):
if os.path.isfile(out):
continue
flow = infer(args, Flownet, device, im1, im2)
if args.size:
flow = resize_flow(flow, args.size)
print(flow.shape)
os.makedirs(os.path.dirname(out), exist_ok=True)
save_raw_float32_image(out, flow)
if args.visualize:
vis = flow_to_image(flow)
cv2.imwrite(os.path.splitext(out)[0] + ".png", vis)
def visualize_flow(self, warp=False):
flow_fmt = pjoin(self.path, "flow", "flow_{:06d}_{:06d}.raw")
mask_fmt = pjoin(self.path, "mask", "mask_{:06d}_{:06d}.png")
color_fmt = pjoin(self.path, "color_down", "frame_{:06d}.raw")
vis_fmt = pjoin(self.path, "vis_flow", "frame_{:06d}_{:06d}.png")
warp_fmt = pjoin(self.path, "vis_flow_warped",
"frame_{:06d}_{:06d}_warped.png")
def get_indices(name):
strs = os.path.splitext(name)[0].split("_")[1:]
return sorted((int(s) for s in strs))
for fmt in (vis_fmt, warp_fmt):
os.makedirs(os.path.dirname(fmt), exist_ok=True)
flow_names = os.listdir(os.path.dirname(flow_fmt))
for flow_name in flow_names:
indices = get_indices(flow_name)
if os.path.isfile(vis_fmt.format(*indices)) and (
not warp or os.path.isfile(warp_fmt.format(*indices))):
continue
indices_pair = [indices, indices[::-1]]
flow_fns = [flow_fmt.format(*idxs) for idxs in indices_pair]
mask_fns = [mask_fmt.format(*idxs) for idxs in indices_pair]
color_fns = [color_fmt.format(idx) for idx in indices]
flows = [image_io.load_raw_float32_image(fn) for fn in flow_fns]
flow_ims = [flowlib.flow_to_image(np.copy(flow)) for flow in flows]
colors = [
image_io.load_raw_float32_image(fn) * 255 for fn in color_fns
]
masks = [cv2.imread(fn, 0) for fn in mask_fns]
masked_colors = [
visualization.apply_mask(im, mask)
for im, mask in zip(colors, masks)
]
masked_flows = [
visualization.apply_mask(im, mask)
for im, mask in zip(flow_ims, masks)
]
masked = np.hstack(masked_colors + masked_flows)
original = np.hstack(colors + flow_ims)
visual = np.vstack((original, masked))
cv2.imwrite(vis_fmt.format(*indices), visual)
if warp:
warped = [
warp_by_flow(color, flow)
for color, flow in zip(colors[::-1], flows)
]
for idxs, im in zip([indices, indices[::-1]], warped):
cv2.imwrite(warp_fmt.format(*idxs), im)