def demo(args): model = torch.nn.DataParallel(RAFT(args)) model.load_state_dict(torch.load(args.model)) # , map_location=torch.device('cpu'))) model = model.module model.to(DEVICE) model.eval() with torch.no_grad(): images = glob.glob(os.path.join(args.path, '*.png')) + \ glob.glob(os.path.join(args.path, '*.jpg')) images = sorted(images) for imfile1, imfile2 in zip(images[:-1], images[1:]): image1 = load_image(imfile1) image2 = load_image(imfile2) padder = InputPadder(image1.shape) image1, image2 = padder.pad(image1, image2) flow_low, flow_up = model(image1, image2, iters=20, test_mode=True) flow = padder.unpad(flow_up[0]).permute(1, 2, 0).cpu().numpy() image1 = padder.unpad(image1[0]).permute(1, 2, 0).cpu().numpy() image2 = padder.unpad(image2[0]).permute(1, 2, 0).cpu().numpy() subname = imfile1.split("/") savename = os.path.join(args.result, subname[-1]) vizproject(savename, image1, image2, flow)
def demo(args): model = torch.nn.DataParallel(RAFT(args)) model.load_state_dict(torch.load(args.model)) model = model.module model.to(DEVICE) model.eval() with torch.no_grad(): images = glob.glob(os.path.join(args.path, '*.png')) + \ glob.glob(os.path.join(args.path, '*.jpg')) images = natsorted(images) for imfile1, imfile2 in tqdm(zip(images[:-1], images[1:]), total=len(images)): try : image1 = load_image(imfile1) image2 = load_image(imfile2) padder = InputPadder(image1.shape) image1, image2 = padder.pad(image1, image2) flow_low, flow_up = model(image1, image2, iters=20, test_mode=True) # Flow Up is the upsampled version if args.save : path = Path(args.path_save) path.mkdir(parents=True, exist_ok=True) flow = padder.unpad(flow_up[0]).permute(1, 2, 0).cpu().numpy() frame_utils.writeFlow(imfile1.replace(args.path,args.path_save).replace('.png','.flo'), flow) else : viz(image1, flow_up) except Exception as e : print(f'Error with {imfile1} : {e}')
def run(args): model = torch.nn.DataParallel(RAFT(args)) model.load_state_dict(torch.load('models/raft-things.pth')) model = model.module model.to(DEVICE) model.eval() output_dir = Path(args.output_dir) images_dir = Path(args.images_dir) images = list(images_dir.glob('*.png')) + list(images_dir.glob('*.jpg')) with torch.no_grad(): images = sorted(images) for i in range(len(images) - 1): im_f1 = str(images[i]) im_f2 = str(images[i + 1]) image1 = load_image(im_f1) image2 = load_image(im_f2) padder = InputPadder(image1.shape) image1, image2 = padder.pad(image1, image2) flow_low, flow_up = model(image1, image2, iters=20, test_mode=True) # 2.2 MB of_f_name = output_dir / f'{i}.npy' np.save(of_f_name, flow_up.cpu())
def create_sintel_submission(model, iters=32, warm_start=False, output_path='sintel_submission'): """ Create submission for the Sintel leaderboard """ model.eval() for dstype in ['clean', 'final']: test_dataset = datasets.MpiSintel(split='test', aug_params=None, dstype=dstype) flow_prev, sequence_prev = None, None for test_id in range(len(test_dataset)): image1, image2, (sequence, frame) = test_dataset[test_id] if sequence != sequence_prev: flow_prev = None padder = InputPadder(image1.shape) image1, image2 = padder.pad(image1[None].cuda(), image2[None].cuda()) flow_low, flow_pr = model(image1, image2, iters=iters, flow_init=flow_prev, test_mode=True) flow = padder.unpad(flow_pr[0]).permute(1, 2, 0).cpu().numpy() if warm_start: flow_prev = forward_interpolate(flow_low[0])[None].cuda() output_dir = os.path.join(output_path, dstype, sequence) output_file = os.path.join(output_dir, 'frame%04d.flo' % (frame+1)) if not os.path.exists(output_dir): os.makedirs(output_dir) frame_utils.writeFlow(output_file, flow) sequence_prev = sequence
def create_kitti_submission(model, iters=24, output_path='kitti_submission', write_png=False): """ Create submission for the Sintel leaderboard """ model.eval() test_dataset = datasets.KITTI(split='testing', aug_params=None) if not os.path.exists(output_path): os.makedirs(output_path) if write_png: out_path_png = output_path + '_png' if not os.path.exists(out_path_png): os.makedirs(out_path_png) for test_id in range(len(test_dataset)): image1, image2, (frame_id, ) = test_dataset[test_id] padder = InputPadder(image1.shape, mode='kitti') image1, image2 = padder.pad(image1[None].cuda(), image2[None].cuda()) _, flow_pr = model(image1, image2, iters=iters, test_mode=True) flow = padder.unpad(flow_pr[0]).permute(1, 2, 0).cpu().numpy() if write_png: output_filename_png = os.path.join(out_path_png, frame_id + '.png') cv2.imwrite(output_filename_png, flow_viz.flow_to_image(flow)) output_filename = os.path.join(output_path, frame_id) frame_utils.writeFlowKITTI(output_filename, flow)
def demo(args): model = torch.nn.DataParallel(RAFT(args)) model.load_state_dict(torch.load(args.model)) model = model.module model.to(DEVICE) model.eval() with torch.no_grad(): images = glob.glob(os.path.join(args.path, '*.png')) + \ glob.glob(os.path.join(args.path, '*.jpg')) images = sorted(images) flows = sorted(glob.glob(args.flow_dir + '*_up.flo')) idx = 0 for imfile1, imfile2, flow_fn in zip(images[:-1], images[1:], flows): idx = idx + 1 print(idx) image1 = load_image(imfile1) image2 = load_image(imfile2) padder = InputPadder(image1.shape) image1, image2 = padder.pad(image1, image2) flow_up = readFlow(flow_fn) # flow_low, flow_up = model(image1, image2, iters=20, test_mode=True) # flow_low, flow_up = model(image1, image2, iters=20, test_mode=True) out_fname = flow_fn.split('/')[-1].replace('.flo', '.jpg') viz_opt_file(image1, flow_up, os.path.join(args.outdir, out_fname))
def demo(args): out_dir = Path(args.out_dir) out_dir.mkdir(parents=True, exist_ok=True) model = torch.nn.DataParallel(RAFT(args)) model.load_state_dict(torch.load(args.model)) model = model.module model.to(DEVICE) model.eval() with torch.no_grad(): images = glob.glob(os.path.join(args.path, '*.png')) + \ glob.glob(os.path.join(args.path, '*.jpg')) images = sorted(images) for index, (imfile1, imfile2) in tqdm(enumerate(zip(images[:-1], images[1:]))): image1 = load_image(imfile1) image2 = load_image(imfile2) padder = InputPadder(image1.shape) image1, image2 = padder.pad(image1, image2) flow_low, flow_up = model(image1, image2, iters=20, test_mode=True) viz(image1, flow_up, index, args.out_dir) print(f'done processing {args.path}')
def computeRAFT(self, img1, img2, it=20): with torch.no_grad(): padder = InputPadder(img1.shape) image1, image2 = padder.pad(img1, img2) flow_low, flow_up = self.raftModel(image1, image2, iters=it, test_mode=True) return flow_up
def inference(self, image1, image2): if self.type == 'farneback': self.mask = np.zeros_like(image1) gray1 = cv2.cvtColor(image1, cv2.COLOR_BGR2GRAY) gray2 = cv2.cvtColor(image2, cv2.COLOR_BGR2GRAY) flow = cv2.calcOpticalFlowFarneback( gray1, gray2, flow=None, pyr_scale=0.5, levels=10, winsize=15, iterations=10, poly_n=7, poly_sigma=1.5, flags=cv2.OPTFLOW_FARNEBACK_GAUSSIAN) self.result = flow if self.type == 'raft': image1 = torch.from_numpy(image1).permute(2, 0, 1).float() image1 = image1[None].to(DEVICE) image2 = torch.from_numpy(image2).permute(2, 0, 1).float() image2 = image2[None].to(DEVICE) padder = InputPadder(image2.shape) image1, image2 = padder.pad(image1, image2) _, flow_up = self.flow_model(image1, image2, iters=5, test_mode=True) self.result = flow_up[0].permute(1, 2, 0).detach().cpu().numpy() if self.type == 'flownet': self.flow = True
def extract_of(model, inpath, outpath, videofolder): """ extract optical flow from frames found at inpath/videofolder/{:06d}.jpg and saves them to outpath/videofolder/{:06d}.jpg """ if not os.path.exists(os.path.join(outpath, videofolder)): os.makedirs(os.path.join(outpath, videofolder)) with torch.no_grad(): images = sorted(os.listdir(os.path.join(inpath, videofolder))) for i, [imfile1, imfile2] in enumerate(zip(images[:-1], images[1:])): image1 = load_image(os.path.join(inpath, videofolder, imfile1)) #[1, 3, W, H] image2 = load_image(os.path.join(inpath, videofolder, imfile2)) padder = InputPadder(image1.shape) image1, image2 = padder.pad(image1, image2) flow_low, flow_up = model( image1, image2, iters=20, test_mode=True) # torch.Size([1, 2, 440, 1024]) discretize_and_save_flow( flow_up[0], osp.join(outpath, videofolder, "flow_{:06d}.jpg".format(i)), osp.join(outpath, videofolder, "flow_{:06d}.pkl".format(i)), )
def validate_sintel(model, iters=32): """ Peform validation using the Sintel (train) split """ model.eval() results = {} for dstype in ['clean', 'final']: val_dataset = datasets.MpiSintel(split='training', dstype=dstype) epe_list = [] for val_id in range(len(val_dataset)): image1, image2, flow_gt, _ = val_dataset[val_id] image1 = image1[None].cuda() image2 = image2[None].cuda() padder = InputPadder(image1.shape) image1, image2 = padder.pad(image1, image2) flow_low, flow_pr = model(image1, image2, iters=iters, test_mode=True) flow = padder.unpad(flow_pr[0]).cpu() epe = torch.sum((flow - flow_gt)**2, dim=0).sqrt() epe_list.append(epe.view(-1).numpy()) epe_all = np.concatenate(epe_list) epe = np.mean(epe_all) px1 = np.mean(epe_all<1) px3 = np.mean(epe_all<3) px5 = np.mean(epe_all<5) print("Validation (%s) EPE: %f, 1px: %f, 3px: %f, 5px: %f" % (dstype, epe, px1, px3, px5)) results[dstype] = np.mean(epe_list) return results
def computeRAFT(net, img1, img2, it=20): B, C, H, W = img1.size() with torch.no_grad(): padder = InputPadder(img1.shape) image1, image2 = padder.pad(img1, img2) flow_low, flow_up = net(image1, image2, iters=it, test_mode=True) return flow_up[:, :, :H, :]
def load_image_list(image_files): images = [] for imfile in sorted(image_files): images.append(load_image(imfile)) images = torch.stack(images, dim=0) images = images.to(DEVICE) padder = InputPadder(images.shape) return padder.pad(images)[0]
def predict(args): model = torch.nn.DataParallel(RAFT(args)) model.load_state_dict(torch.load(args.model)) model = model.module model.to(DEVICE) model.eval() with torch.no_grad(): images = glob.glob(os.path.join(args.path, '*.png')) + \ glob.glob(os.path.join(args.path, '*.jpg')) folder = os.path.basename(args.path) floout = os.path.join(args.outroot, folder) rawfloout = os.path.join(args.raw_outroot, folder) os.makedirs(floout, exist_ok=True) os.makedirs(rawfloout, exist_ok=True) gap = args.gap images = sorted(images) images_ = images[:-gap] for index, imfile1 in enumerate(images_): if args.reverse: image1 = load_image(images[index+gap]) image2 = load_image(imfile1) svfile = images[index+gap] else: image1 = load_image(imfile1) image2 = load_image(images[index + gap]) svfile = imfile1 padder = InputPadder(image1.shape) image1, image2 = padder.pad(image1, image2) flow_low, flow_up = model(image1, image2, iters=20, test_mode=True) flopath = os.path.join(floout, os.path.basename(svfile)) rawflopath = os.path.join(rawfloout, os.path.basename(svfile)) flo = flow_up[0].permute(1, 2, 0).cpu().numpy() # save raw flow writeFlowFile(rawflopath[:-4]+'.flo', flo) # save image. flo = flow_viz.flow_to_image(flo) cv2.imwrite(flopath[:-4]+'.png', flo[:, :, [2, 1, 0]])
def infer(model, seq_img_dir, suffix, iters=24, backward_flow=True): if backward_flow: flow_img_dir = os.path.join(seq_img_dir, '../flow_backward_img_{}'.format(suffix)) flow_np_dir = os.path.join(seq_img_dir, '../flow_backward_np_{}'.format(suffix)) # flow_np_save_path = os.path.join(seq_img_dir, '../flow_backward_{}.npy'.format(suffix)) else: flow_img_dir = os.path.join(seq_img_dir, '../flow_forward_img_{}'.format(suffix)) flow_np_dir = os.path.join(seq_img_dir, '../flow_forward_np_{}'.format(suffix)) # flow_np_save_path = os.path.join(seq_img_dir, '../flow_forward_{}.npy'.format(suffix)) if not os.path.exists(flow_img_dir): os.makedirs(flow_img_dir) if not os.path.exists(flow_np_dir): os.makedirs(flow_np_dir) model.eval() dataset = datasets.InferVideoDataset(seq_img_dir, backward_flow=backward_flow) # flow_list, flow_img_list = [], [] for val_id in tqdm.tqdm(range(len(dataset))): image1, image2, path1, path2 = dataset[val_id] image1 = image1[None].cuda() image2 = image2[None].cuda() padder = InputPadder(image1.shape, mode='sintel') image1, image2 = padder.pad(image1, image2) flow_low, flow_pr = model(image1, image2, iters=iters, test_mode=True) flow = padder.unpad(flow_pr[0]).cpu() # map flow to rgb image # pdb.set_trace() # flow = flow[0].permute(1,2,0).cpu().numpy() flow = flow.permute(1, 2, 0).cpu().numpy() flow_img = flow_viz.flow_to_image(flow) # flow_list.append(flow) # flow_img_list.append(flow_img) imageio.imwrite(os.path.join(flow_img_dir, path1.split('/')[-1]), flow_img) np.save( os.path.join(flow_np_dir, path1.split('/')[-1].split('.')[0] + '.npy'), flow)
def demo(args): model = torch.nn.DataParallel(RAFT(args), device_ids=[0, 1], output_device=1).cuda() model.load_state_dict(torch.load(args.model)) # model = model.module # model.to(DEVICE) model.eval() with torch.no_grad(): # images = glob.glob(os.path.join(args.path, '*.png')) + \ # glob.glob(os.path.join(args.path, '*.jpg')) with open('seqheads.txt') as fs: image_seqheads = fs.read().splitlines() fileidx = 0 for img1 in tqdm(image_seqheads): city, n1, n2, suf = img1.split('/', maxsplit=1)[-1].split("_") n2 = int(n2) for idx in range(1, 11): img2 = f"leftImg8bit_sequence/{city}_{n1}_{n2+idx:06d}_{suf}" savepath = f"flow_sequence/{city}_{n1}_{n2+idx:06d}" print(f"{img1}\n{img2}") image1 = load_image(os.path.join(args.path, img1)) image2 = load_image(os.path.join(args.path, img2)) savepath = os.path.join(args.path, savepath) padder = InputPadder(image1.shape) image1, image2 = padder.pad(image1, image2) image1_b = torch.cat([image1, image2], 0) image2_b = torch.cat([image2, image1], 0) flow_low, flow_up = model(image1_b, image2_b, iters=20, test_mode=True) flow_up_fwd = flow_up[0].unsqueeze(0) flow_up_bwd = flow_up[1].unsqueeze(0) # viz(image1, image2, flow_up_fwd, flow_up_bwd, fileidx) save_flow(flow_up_fwd, f"{savepath}_flow_fwd.png") save_flow(flow_up_bwd, f"{savepath}_flow_bwd.png") fileidx += 1
def evaluate_davis(model, iters=32): """ Peform validation using the Sintel (train) split """ model.eval() val_dataset = datasets.DAVISDataset(split='train') for val_id in tqdm(range(len(val_dataset))): image1, image2, image_paths = val_dataset[val_id] image1 = image1[None].cuda() image2 = image2[None].cuda() padder = InputPadder(image1.shape) image1, image2 = padder.pad(image1, image2) _, flow_pr = model(image1, image2, iters=iters, test_mode=True) forward_flow = padder.unpad(flow_pr[0]).permute(1, 2, 0).cpu().numpy() _, flow_pr = model(image2, image1, iters=iters, test_mode=True) backward_flow = padder.unpad(flow_pr[0]).permute(1, 2, 0).cpu().numpy() # find out result storing paths fpath = image_paths[0] ind = fpath.rfind("/") name = fpath[ind + 1:fpath.rfind(".")] folder_path = fpath[:ind] flow_folder = folder_path.replace("JPEGImages", "Flows") flowviz_folder = folder_path.replace("JPEGImages", "FlowVizs") flow_path = os.path.join(flow_folder, f"forward_{name}.flo") flowviz_path = os.path.join(flowviz_folder, f"forward_{name}.png") if not os.path.exists(flow_folder): os.makedirs(flow_folder) if not os.path.exists(flowviz_folder): os.makedirs(flowviz_folder) frame_utils.writeFlow(flow_path, forward_flow) Image.fromarray(flow_viz.flow_to_image(forward_flow)).save( open(flowviz_path, "wb"), format="PNG") flow_path = os.path.join(flow_folder, f"backward_{name}.flo") flowviz_path = os.path.join(flowviz_folder, f"backward_{name}.png") frame_utils.writeFlow(flow_path, backward_flow) Image.fromarray(flow_viz.flow_to_image(backward_flow)).save( open(flowviz_path, "wb"), format="PNG")
def compute_flow_dir(model, dirpath, dirpathsave, resize=None) : images = glob.glob(os.path.join(dirpath, '*.png')) + \ glob.glob(os.path.join(dirpath, '*.jpg')) images = natsorted(images) for imfile1, imfile2 in tqdm(zip(images[:-1], images[1:]), total=len(images)): image1 = load_image(imfile1) image2 = load_image(imfile2) extension=imfile1.split('.')[-1] padder = InputPadder(image1.shape) image1, image2 = padder.pad(image1, image2) flow_low, flow_up = model(image1, image2, iters=20, test_mode=True) # Flow Up is the upsampled version if resize is not None : flow_up = nn.functional.interpolate(flow_up, size=resize, mode='bilinear', align_corners=False) path = Path(dirpathsave) path.mkdir(parents=True, exist_ok=True) flow = padder.unpad(flow_up[0]).permute(1, 2, 0).cpu().numpy() frame_utils.writeFlow(imfile1.replace(dirpath, dirpathsave).replace(extension,'flo'), flow)
def demo(args): model = torch.nn.DataParallel(RAFT(args)) model.load_state_dict(torch.load(args.model)) model = model.module model.to(DEVICE) model.eval() with torch.no_grad(): images = glob.glob(os.path.join(args.path, '*.png')) + \ glob.glob(os.path.join(args.path, '*.jpg')) images = sorted(images) for imfile1, imfile2 in zip(images[:-1], images[1:]): image1 = load_image(imfile1) image2 = load_image(imfile2) padder = InputPadder(image1.shape) image1, image2 = padder.pad(image1, image2) flow_low, flow_up = model(image1, image2, iters=20, test_mode=True) viz(image1, flow_up)
def run(img1_path, img2_path, iters=20): img1 = torch.from_numpy(np.array(Image.open(img1_path)).astype(np.uint8)).permute(2, 0, 1).float() img2 = torch.from_numpy(np.array(Image.open(img2_path)).astype(np.uint8)).permute(2, 0, 1).float() images = torch.stack([img1, img2], dim=0).to(DEVICE) images = InputPadder(images.shape).pad(images)[0] with torch.no_grad(): image1 = images[0, None] image2 = images[1, None] flow_low, flow_up = model(image1, image2, iters=iters, test_mode=True) flow = flow_up[0].permute(1, 2, 0).cpu().numpy() return flow
def validate_kitti(model, args, iters=24): """ Peform validation using the KITTI-2015 (train) split """ model.eval() val_dataset = datasets.KITTI(split='training', root=args.dataset) from tqdm import tqdm out_list, epe_list = [], [] for _, val_id in enumerate(tqdm(list(range(len(val_dataset))))): image1, image2, flow_gt, valid_gt = val_dataset[val_id] image1 = image1[None].cuda() image2 = image2[None].cuda() padder = InputPadder(image1.shape, mode='kitti') image1, image2 = padder.pad(image1, image2) flow_low, flow_pr = model(image1, image2, iters=iters, test_mode=True) flow = padder.unpad(flow_pr[0]).cpu() epe = torch.sum((flow - flow_gt)**2, dim=0).sqrt() mag = torch.sum(flow_gt**2, dim=0).sqrt() epe = epe.view(-1) mag = mag.view(-1) val = valid_gt.view(-1) >= 0.5 out = ((epe > 3.0) & ((epe / mag) > 0.05)).float() epe_list.append(epe[val].mean().item()) out_list.append(out[val].cpu().numpy()) epe_list = np.array(epe_list) out_list = np.concatenate(out_list) epe = np.mean(epe_list) f1 = 100 * np.mean(out_list) print("Validation KITTI: %f, %f" % (epe, f1)) return {'kitti-epe': epe, 'kitti-f1': f1}
def raft_flow_a_to_b(model, imfile1, imfile2): image1_shape = Image.open(imfile1).size image2_shape = Image.open(imfile2).size image1 = load_image(imfile1) image2 = load_image(imfile2) padder = InputPadder(image1.shape) image1, image2 = padder.pad(image1, image2) flow_low, flow_up = model(image1, image2, iters=20, test_mode=True) # viz(image1, flow_up) x, y = np.meshgrid(np.linspace(0, 1024 - 1, 1024), np.linspace(0, 1024 - 1, 1024)) base_grid = np.stack([x, y], axis=-1) flow_up_np = flow_up[0].cpu().permute(1, 2, 0).numpy() + base_grid scale_x = image2_shape[0] / 1024 scale_y = image2_shape[1] / 1024 flow_up_np[..., 0] *= scale_x flow_up_np[..., 1] *= scale_y flow_up_np = float_image_resize(flow_up_np, image1_shape[::-1]) # return flow_up_np # import IPython # IPython.embed() # assert 0 # img_a = imageio.imread(imfile1) # img_b = imageio.imread(imfile2) # x, y = np.meshgrid(np.linspace(0, 1024 - 1, 1024), # np.linspace(0, 1024 - 1, 1024)) # base_grid = np.stack([x, y], axis=-1) # flow_up_np = flow_up[0].cpu().permute(1,2,0).numpy() + base_grid # warped = cv2.remap(img_b, flow_up_np[..., 0].astype(np.float32), flow_up_np[..., 1].astype(np.float32), interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT) return flow_up_np, None import IPython IPython.embed() assert 0
def demo(args): model = torch.nn.DataParallel(RAFT(args)) model.load_state_dict(torch.load(args.model)) model = model.module model.to(DEVICE) model.eval() with torch.no_grad(): images = glob.glob(os.path.join(args.path, '*.png')) + \ glob.glob(os.path.join(args.path, '*.jpg')) images = sorted(images) img = images[0] for i in range(len(images)-1): image1 = load_image(images[i]) image2 = load_image(images[i+1]) padder = InputPadder(image1.shape) image1, image2 = padder.pad(image1, image2) flow_low, flow_up = model(image1, image2, iters=20, test_mode=True) viz(image1, flow_up, i)
def opt_flow_estimation(args): """ args.path: path to the directory of the dataset that contains the images. - base_dir/ - ArgoVerse/ - video1/ - frame1.jpg - frame2.jpg - video2/ - BDD/ - Charades/ - LaSOT/ - YFCC100M/ """ model = torch.nn.DataParallel(RAFT(args)) model.load_state_dict(torch.load(args.model)) model = model.module model.to(DEVICE) model.eval() with torch.no_grad(): base_dir = args.path data_srcs = [ fn.split('/')[-1] for fn in sorted(glob.glob(os.path.join(base_dir, '*'))) ] if args.datasrc: data_srcs = [args.datasrc] for data_src in data_srcs: print("Processing", data_src) videos = [ fn.split('/')[-1] for fn in sorted( glob.glob(os.path.join(base_dir, data_src, '*'))) ] for idx, video in enumerate(tqdm.tqdm(videos)): fpath = os.path.join(base_dir, data_src, video) images = glob.glob(os.path.join(fpath, '*.png')) + \ glob.glob(os.path.join(fpath, '*.jpg')) images = sorted(images) for imfile1, imfile2 in zip(images[:-1], images[1:]): image1 = load_image(imfile1) image2 = load_image(imfile2) padder = InputPadder(image1.shape) image1, image2 = padder.pad(image1, image2) flow_low, flow_up = model(image1, image2, iters=20, test_mode=True) # Store the flow vector flow_fname = imfile1.split("/")[-1].replace(".jpg", ".flo") flow_up_fname = flow_fname.split(".") flow_up_fname[0] = flow_up_fname[0] + "_up" flow_up_fname = ".".join(flow_up_fname) flow_low_fname = flow_fname.split(".") flow_low_fname[0] = flow_low_fname[0] + "_low" flow_low_fname = ".".join(flow_low_fname) up_fname = os.path.join(args.outdir, data_src, video, flow_up_fname) # low_fname = os.path.join(args.outdir, data_src, video, flow_low_fname) if not os.path.exists( os.path.join(args.outdir, data_src, video)): os.makedirs(os.path.join(args.outdir, data_src, video)) flow_up = flow_up[0].permute(1, 2, 0).cpu().numpy() # flow_low = flow_low[0].permute(1, 2, 0).cpu().numpy() writeFlow(up_fname, flow_up)
def validate_kitti_colorjitter(model, args, iters=24): """ Peform validation using the KITTI-2015 (train) split """ from torchvision.transforms import ColorJitter from tqdm import tqdm model.eval() val_dataset = datasets.KITTI(split='training', root=args.dataset) jitterparam = 0.86 photo_aug = ColorJitter(brightness=jitterparam, contrast=jitterparam, saturation=jitterparam, hue=jitterparam / 3.14) def color_transform(img1, img2, photo_aug): torch.manual_seed(1234) np.random.seed(1234) img1 = img1.permute([1, 2, 0]).numpy().astype(np.uint8) img2 = img2.permute([1, 2, 0]).numpy().astype(np.uint8) image_stack = np.concatenate([img1, img2], axis=0) image_stack = np.array(photo_aug(Image.fromarray(image_stack)), dtype=np.uint8) img1, img2 = np.split(image_stack, 2, axis=0) img1 = torch.from_numpy(img1).permute([2, 0, 1]).float() img2 = torch.from_numpy(img2).permute([2, 0, 1]).float() return img1, img2 out_list, epe_list = [], [] for _, val_id in enumerate(tqdm(list(range(len(val_dataset))))): image1, image2, flow_gt, valid_gt = val_dataset[val_id] image1, image2 = color_transform(image1, image2, photo_aug) image1 = image1[None].cuda() image2 = image2[None].cuda() padder = InputPadder(image1.shape, mode='kitti') image1, image2 = padder.pad(image1, image2) flow_low, flow_pr = model(image1, image2, iters=iters, test_mode=True) flow = padder.unpad(flow_pr[0]).cpu() epe = torch.sum((flow - flow_gt)**2, dim=0).sqrt() mag = torch.sum(flow_gt**2, dim=0).sqrt() epe = epe.view(-1) mag = mag.view(-1) val = valid_gt.view(-1) >= 0.5 print("Index: %d, valnum: %d" % (val_id, torch.sum(valid_gt).item())) out = ((epe > 3.0) & ((epe / mag) > 0.05)).float() epe_list.append(epe[val].mean().item()) out_list.append(out[val].cpu().numpy()) epe_list = np.array(epe_list) out_list = np.concatenate(out_list) epe = np.mean(epe_list) f1 = 100 * np.mean(out_list) print("jitterparam:%f, Validation KITTI: %f, %f" % (jitterparam, epe, f1)) return {'kitti-epe': epe, 'kitti-f1': f1}
def validate_kitti_customized(model, iters=24): """ Peform validation using the KITTI-2015 (train) split """ model.eval() val_dataset = datasets.KITTI( split='training', root='/home/shengjie/Documents/Data/Kitti/kitti_stereo/stereo15') out_list, epe_list = [], [] for val_id in range(len(val_dataset)): image1, image2, flow_gt, valid_gt = val_dataset[val_id] image1 = image1[None].cuda() image2 = image2[None].cuda() padder = InputPadder(image1.shape, mode='kitti') image1, image2 = padder.pad(image1, image2) flow_low, flow_pr = model(image1, image2, iters=iters, test_mode=True) flow = padder.unpad(flow_pr[0]).cpu() flowT = flow.cpu() flownp = flowT.numpy() image1_vls = padder.unpad(image1[0]).cpu() image2_vls = padder.unpad(image2[0]).cpu() image1_vlsnp = image1_vls.permute([1, 2, 0]).cpu().numpy().astype(np.uint8) image2_vlsnp = image2_vls.permute([1, 2, 0]).cpu().numpy().astype(np.uint8) flow_gt_vls_np = flow_gt.cpu().numpy() valid_gt_vls_np = valid_gt.cpu().numpy() _, h, w = flowT.shape xx, yy = np.meshgrid(range(w), range(h), indexing='xy') resampledxx = xx + flowT[0].cpu().numpy() resampledyy = yy + flowT[1].cpu().numpy() epipole_vote(xx, yy, flownp, image1_vlsnp, image2_vlsnp, flow_gt_vls_np, valid_gt_vls_np) resampledxx = ((resampledxx / (w - 1)) - 0.5) * 2 resampledyy = ((resampledyy / (h - 1)) - 0.5) * 2 resamplegrid = torch.stack( [torch.from_numpy(resampledxx), torch.from_numpy(resampledyy)], dim=2).unsqueeze(0).float() image1_recon_vls = torch.nn.functional.grid_sample( input=image2_vls.unsqueeze(0), grid=resamplegrid, mode='bilinear', padding_mode='reflection') # rndx = np.random.randint(0, w) # rndy = np.random.randint(0, h) rndx = 215 rndy = 278 tarx = rndx + flownp[0, int(rndy), int(rndx)] tary = rndy + flownp[1, int(rndy), int(rndx)] plt.figure() plt.imshow(image1.squeeze().permute([1, 2, 0 ]).cpu().numpy().astype(np.uint8)) plt.scatter(rndx, rndy, 1, 'r') plt.figure() plt.imshow(image2.squeeze().permute([1, 2, 0 ]).cpu().numpy().astype(np.uint8)) plt.scatter(tarx, tary, 1, 'r') plt.figure() plt.imshow(image1_recon_vls.squeeze().permute( [1, 2, 0]).cpu().numpy().astype(np.uint8)) import PIL.Image as Image from core.utils.flow_viz import flow_to_image flowimg = flow_to_image(flow.permute([1, 2, 0]).cpu().numpy()) Image.fromarray(flowimg).show() epe = torch.sum((flow - flow_gt)**2, dim=0).sqrt() mag = torch.sum(flow_gt**2, dim=0).sqrt() epe = epe.view(-1) mag = mag.view(-1) val = valid_gt.view(-1) >= 0.5 out = ((epe > 3.0) & ((epe / mag) > 0.05)).float() epe_list.append(epe[val].mean().item()) out_list.append(out[val].cpu().numpy()) epe_list = np.array(epe_list) out_list = np.concatenate(out_list) epe = np.mean(epe_list) f1 = 100 * np.mean(out_list) print("Validation KITTI: %f, %f" % (epe, f1)) return {'kitti-epe': epe, 'kitti-f1': f1}
def run(args): in_path_left = os.path.join(args.path, "image_left") in_path_right = os.path.join(args.path, "image_right") out_path_foward = os.path.join(args.path, "flow_forward") out_path_backward = os.path.join(args.path, "flow_backward") print("Computing forward and backwrd optical flow between:") print(f" {in_path_left}") print(f" {in_path_right}\n") create_dir(out_path_foward) create_dir(out_path_backward) num_gpus = torch.cuda.device_count() batch_size = num_gpus data_loader = DataLoader(dataset=FlowForwardBackwardDataset( in_path_left, in_path_right), batch_size=batch_size, shuffle=False, num_workers=2) model = torch.nn.DataParallel(RAFT(args)) model.load_state_dict(torch.load(args.model)) # model = model.module ## -> No idea why this was needed model.to(DEVICE) model.eval() print("\nStart ...") print(f"Running on {num_gpus} GPUs\n") with torch.no_grad(): for img_str, img_left, img_right in tqdm(data_loader): padder = InputPadder(img_left.shape) img_left, img_right = padder.pad(img_left, img_right) # get foward flow - left -> right _, forward_flow_up = model(img_left, img_right, iters=20, test_mode=True) forward_flow = forward_flow_up.permute(0, 2, 3, 1).cpu().numpy() # get backward flow - right -> left _, backward_flow_up = model(img_right, img_left, iters=20, test_mode=True) backward_flow = backward_flow_up.permute(0, 2, 3, 1).cpu().numpy() # Loop over the batches for i in range(batch_size): save_flow(forward_flow[i], os.path.join(out_path_foward, img_str[i] + "_flo")) save_flow(backward_flow[i], os.path.join(out_path_backward, img_str[i] + "_flo")) print("\nDone!")
model.load_state_dict(torch.load(args.model , map_location=torch.device(DEVICE))) model = model.module model.to(DEVICE) model.eval() with torch.no_grad(): images = glob.glob(os.path.join(args.path, '*.png')) + \ glob.glob(os.path.join(args.path, '*.jpg')) images = sorted(images) for imfile1, imfile2 in zip(images[100:101], images[101:102]):#zip(images[:-1], images[1:]): print(imfile1 + ' ' + imfile2) for scaling in range(8,9,1): image1, img_orig1 = load_image(imfile1,scaling=scaling) image2, img_orig2 = load_image(imfile2,scaling=scaling) padder = InputPadder(image1.shape) image1, image2 = padder.pad(image1, image2) flow_low, flow_up = model(image1, image2, iters=20, test_mode=True) B,C,W,H = img_orig1.shape Bf,Cf,Wf,Hf = flow_up.shape flow_up = F.interpolate(flow_up,(W,H),mode='bicubic') flow_up[:,0,:,:] = flow_up[:,0,:,:]*W/Wf flow_up[:,1,:,:] = flow_up[:,1,:,:]*H/Hf warpimg1 = warp(img_orig2,flow_up) #wrapimg2 = warp(image1,flow_up) image1 = padder.unpad(img_orig1[0]).permute(1, 2, 0).cpu() #image2 = padder.unpad(image2[0]).permute(1, 2, 0).cpu().numpy() warpimg1 = padder.unpad(warpimg1[0]).permute(1, 2, 0).cpu() i_loss = (image1 - warpimg1).abs() image1 = image1.numpy() warpimg1 = warpimg1.numpy()
def convert2of(video_path, model_path): parser = argparse.ArgumentParser() parser.add_argument('--small', action='store_true', help='use small model') parser.add_argument('--mixed_precision', action='store_true', help='use mixed precision') parser.add_argument('--alternate_corr', action='store_true', help='use efficent correlation implementation') args = parser.parse_args() model = torch.nn.DataParallel(RAFT(args)) model.load_state_dict(torch.load(model_path)) model = model.module model.to("cuda") model.eval() # path = os.path.join(video_path) subdirs = os.listdir(video_path) for subdir in subdirs: images = [] # compute the path to the subdir subpath = os.path.join(video_path, subdir) # elems = os.listdir(subpath) # for elem in elems: # # name = elem[:-4] # path = os.path.join(subpath, elem) # sample = os.listdir(path) # for name in sample: # print(name) # ssspath = os.path.join(path, name) cap = cv2.VideoCapture(subpath) fps = cap.get(cv2.CAP_PROP_FPS) while True: ret, frame = cap.read() if ret: frame = frame[int(720 * 0.15):int(720 * 0.85), int(1280 * 0.15):int(1280 * 0.85)] frame = cv2.resize(frame, (int(455), int(256))) images.append(torch.from_numpy(frame).permute(2, 0, 1).float()) else: break # cap = de.VideoReader(video_path, width = 455, height= 256) # fps = len(cap) # print(fps) # shape = cap[0].shape # print(shape) # i = 0 # for i in cap: # frame = cap[i].asnumpy() # i = i + 1 print("Read frames finished") images = torch.stack(images, dim=0) images = images.to("cuda") padder = InputPadder(images.shape) images = padder.pad(images)[0] fourcc = cv2.VideoWriter_fourcc(*'MP42') image1 = images[0, None] image2 = images[1, None] start_t = time.time() with torch.no_grad(): flow_low, flow_up = model(image1, image2, iters=20, test_mode=True) print("Each prediction cost {}s".format(time.time() - start_t)) output_image = viz(image1, flow_up) # out = cv2.VideoWriter(dst_path, fourcc, # fps, (output_image.shape[1], output_image.shape[0])) dst_path = os.path.join(video_path, 'move_opt') if not os.path.exists(dst_path): os.makedirs(dst_path) dst_path = os.path.join(dst_path, subdir) out = imageio.get_writer(dst_path, format='mp4', mode='I', fps=fps) #print(output_image.shape) with torch.no_grad(): for i in range(images.shape[0] - 1): image1 = images[i, None] image2 = images[i + 1, None] _, flow_up = model(image1, image2, iters=20, test_mode=True) tmp = viz(image1, flow_up) # tmp = cv2.cvtColor(tmp, cv2.COLOR_RGB2BGR) # gray = cv2.cvtColor(tmp, cv2.COLOR_RGB2GRAY) # tmp = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR) # cv2.imshow('', tmp) # cv2.waitKey(1) # out.write(tmp) out.append_data(tmp) cap.release()