def demo(args): model = RAFT(args) model = torch.nn.DataParallel(model) model.load_state_dict(torch.load(args.model)) model.to(DEVICE) model.eval() with torch.no_grad(): # sintel images image1 = load_image('images/sintel_0.png') image2 = load_image('images/sintel_1.png') flow_predictions = model(image1, image2, iters=args.iters, upsample=False) display(image1[0], image2[0], flow_predictions[-1][0]) # kitti images image1 = load_image('images/kitti_0.png') image2 = load_image('images/kitti_1.png') flow_predictions = model(image1, image2, iters=16) display(image1[0], image2[0], flow_predictions[-1][0]) # davis images image1 = load_image('images/davis_0.jpg') image2 = load_image('images/davis_1.jpg') flow_predictions = model(image1, image2, iters=16) display(image1[0], image2[0], flow_predictions[-1][0])
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 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 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 = load_image_list(images) for i in range(images.shape[0] - 1): image1 = images[i, None] image2 = images[i + 1, None] print(image1.size()) print(image2.size()) image1 = F.interpolate(image1, (480, 720)) image2 = F.interpolate(image2, (480, 720)) flow_low, flow_up = model(image1, image2, iters=20, test_mode=True) to_save = flow_up.squeeze(0).cpu().numpy() np.save('tosave.npy', to_save) viz(image1, flow_up)
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): 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 = RAFT(args) model = torch.nn.DataParallel(model) model.load_state_dict(torch.load(args.model)) model.to(DEVICE) model.eval() with torch.no_grad(): cap = cv2.VideoCapture('video.mp4') _, left_frame = cap.read() h, w, _ = left_frame.shape left_tensor = preprocess(left_frame) while (1): _, right_frame = cap.read() right_tensor = preprocess(right_frame) start1 = time.time() flow_predictions = model(left_tensor, right_tensor, iters=args.iters, upsample=True) print(time.time() - start1) flow_image = postprocess(flow_predictions, w * 2, h * 2) cv2.imshow('frame', flow_image) k = cv2.waitKey(25) if (k == 27): break left_tensor = right_tensor.clone()
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 batch_calOpt(args): model = RAFT(args) model = torch.nn.DataParallel(model) model.load_state_dict(torch.load(args.model)) model.to(DEVICE) model.eval() # Transform transform = transforms.Compose([transforms.ToTensor()]) # ImageFolder and Loader image_dataset = OpticalFlowFolder(IMAGE_FOLDER_PATH, transform=transform) image_loader = torch.utils.data.DataLoader(image_dataset, batch_size=BATCH_SIZE) start = time.time() with torch.no_grad(): for left, right in image_loader: # Most of the time, this preprocessing is not needed # Especially if the video dimensions are multiple of 8s _, _, h, w = left.shape if ((h % 8 != 0) or (w % 8 != 0)): left = pad8(left) right = pad8(right) # Forward flow_predictions = model(left, right, iters=args.iters, upsample=False) print("Time Elapsed: ", time.time() - start)
def prepare(args): global model model = torch.nn.DataParallel(RAFT(args)) model.load_state_dict(torch.load(args.model)) model = model.module model.to(DEVICE) model.eval()
def load_model(args) : model = torch.nn.DataParallel(RAFT(args)) model.load_state_dict(torch.load(args.model)) model = model.module model.to(DEVICE) model.eval() return model
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(): for i, folder in enumerate(sorted(os.listdir(args.path))[:100]): print(f"video {i+1}") images = glob.glob( os.path.join( args.path + folder + "/images/", "*.png")) + glob.glob( os.path.join(args.path + folder + "/images/", "*.jpg")) masks = glob.glob( os.path.join( args.path + folder + "/masks/", "*.png")) + glob.glob( os.path.join(args.path + folder + "/masks/", "*.jpg")) fact_people = FactorsPeople(DEVICE) ret = [] is_ret = False images = sorted(images)[:100] masks = sorted(masks)[:100] count = 0 for i, imfiles in enumerate(zip(images[:-1], images[1:])): if count >= 100: print("count >= 100") break imfile1 = imfiles[0] imfile2 = imfiles[1] image1, _ = fact_people.get_image(imfile1, masks[i]) image2, _ = fact_people.get_image(imfile2, masks[i + 1]) image1 = image1.to(DEVICE) image2 = image2.to(DEVICE) image1 *= 255.0 image2 *= 255.0 _, flow_up = model(image1, image2, iters=20, test_mode=True) if not is_ret: is_ret = True ret = flow_up else: ret = torch.cat((ret, flow_up), 0) count += 1 np.save(f"/phoenix/S3/ab2383/data/flows/{i+1}.npy", ret.detach().cpu().numpy()) if i == 99: break
def inference(args): # get the RAFT model model = RAFT(args) # load pretrained weights pretrained_weights = torch.load(args.model) save = args.save if save: if not os.path.exists("demo_frames"): os.mkdir("demo_frames") if torch.cuda.is_available(): device = "cuda" # parallel between available GPUs model = torch.nn.DataParallel(model) # load the pretrained weights into model model.load_state_dict(pretrained_weights) model.to(device) else: device = "cpu" # change key names for CPU runtime pretrained_weights = get_cpu_model(pretrained_weights) # load the pretrained weights into model model.load_state_dict(pretrained_weights) # change model's mode to evaluation model.eval() video_path = args.video # capture the video and get the first frame cap = cv2.VideoCapture(video_path) ret, frame_1 = cap.read() # frame preprocessing frame_1 = frame_preprocess(frame_1, device) counter = 0 with torch.no_grad(): while True: # read the next frame ret, frame_2 = cap.read() if not ret: break # preprocessing frame_2 = frame_preprocess(frame_2, device) # predict the flow flow_low, flow_up = model(frame_1, frame_2, iters=20, test_mode=True) # transpose the flow output and convert it into numpy array ret = vizualize_flow(frame_1, flow_up, save, counter) if not ret: break frame_1 = frame_2 counter += 1
def save_results(args, save_path): 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')) save_list(path=f'{save_path}/image.list', list_obj=images) length = len(images) # images = load_image_list(images) image1 = load_image(images[0])[None] # Get sample points points1 = [] for i in range(15): for j in range(20): # pt1 = (160, 160) # x = pt1[0] # y = pt1[1] x = j * 30 y = i * 30 pt1 = (x, y) points1.append(pt1) for i in range(length - 1): # image1 = images[i,None] # image2 = images[i+1,None] image2 = load_image(images[i + 1])[None] flow_low, flow_up = model(image1, image2, iters=20, test_mode=True) filename = f'flow_{i}_{i+1}' np.savez(f'{save_path}/{filename}.npz', image1=images[i], image2=images[i + 1], flow=flow_up) viz(image1, image2, flow_up, uv1=points1, filename=f'{save_path}/{filename}.png') image1 = image2 print(f'Finish saving {i}/{length}')
def initModel(self): from argparse import Namespace modelname = "D:/git/RAFT/models/raft-sintel.pth" #self.model = torch.nn.DataParallel(RAFT(Namespace(small=False,mixed_precision=False))) self.model = torch.nn.DataParallel( RAFT(Namespace(small=False, mixed_precision=False))) self.model.load_state_dict(torch.load(modelname)) self.model = self.model.module self.model.to('cuda') self.model.eval()
def train(args): model = RAFT(args) model = nn.DataParallel(model) print("Parameter Count: %d" % count_parameters(model)) if args.restore_ckpt is not None: model.load_state_dict(torch.load(args.restore_ckpt)) model.cuda() model.train() if 'chairs' not in args.dataset: model.module.freeze_bn() train_loader = fetch_dataloader(args) optimizer, scheduler = fetch_optimizer(args, model) total_steps = 0 logger = Logger(model, scheduler) should_keep_training = True while should_keep_training: for i_batch, data_blob in enumerate(train_loader): image1, image2, flow, valid = [x.cuda() for x in data_blob] optimizer.zero_grad() flow_predictions = model(image1, image2, iters=args.iters) loss, metrics = sequence_loss(flow_predictions, flow, valid) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip) optimizer.step() scheduler.step() total_steps += 1 logger.push(metrics) if total_steps % VAL_FREQ == VAL_FREQ - 1: PATH = 'checkpoints/%d_%s.pth' % (total_steps + 1, args.name) torch.save(model.state_dict(), PATH) if total_steps == args.num_steps: should_keep_training = False break PATH = 'checkpoints/%s.pth' % args.name torch.save(model.state_dict(), PATH) return PATH
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() fps = None with torch.no_grad(): fps = FPS().start() try: image1 = None image2 = None while True: # Wait for a coherent pair of frames: depth and color frames = pipeline.wait_for_frames() color_frame = frames.get_color_frame() if not color_frame: continue im = np.asanyarray(color_frame.get_data())[:, :, ::-1].astype( np.uint8) img = torch.from_numpy(im).permute( 2, 0, 1).float().unsqueeze(0).to(DEVICE) if image2 is None: image1 = img else: image1 = image2 image2 = img flow_low, flow_up = model(image1, image2, iters=args.iters, test_mode=True) fps.update() fps.stop() print(f"FPS: {fps.fps():.2f}") viz(image1, flow_up) # check for escape key key = cv2.waitKey(1) if key == 27 or key == 1048603: break finally: pipeline.stop() cv2.destroyAllWindows()
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 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 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(): final_dict = {} for i, folder in enumerate(sorted(os.listdir(args.path))[:100]): images = glob.glob( os.path.join( args.path + folder + "/images/", "*.png")) + glob.glob( os.path.join(args.path + folder + "/images/", "*.jpg")) fact_people = FactorsPeople(DEVICE) ret = [] is_ret = False images = sorted(images)[:100] mask_tmp_pth = "/phoenix/S3/ab2383/data/TikTok_dataset/00267/masks/0001.png" for imfile1, imfile2 in zip(images[:-1], images[1:]): image1, _ = fact_people.get_image(imfile1, mask_tmp_pth) image2, _ = fact_people.get_image(imfile2, mask_tmp_pth) image1 *= 255.0 image2 *= 255.0 _, flow_up = model(image1, image2, iters=20, test_mode=True) if not is_ret: is_ret = True ret = flow_up else: ret = torch.cat((ret, flow_up), 0) final_dict[i + 1] = ret.detach().cpu().numpy() with open("/phoenix/S3/ab2383/data/hundred_hundreds.pickle", "wb") as f: pickle.dump( final_dict, f, protocol=pickle.HIGHEST_PROTOCOL, )
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 = load_image_list(images) for i in range(images.shape[0] - 1): image1 = images[i, None] image2 = images[i + 1, None] flow_low, flow_up = model(image1, image2, iters=20, test_mode=True) viz(image1, flow_up)
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 = load_image_list(images) # Get sample points points1 = [] for i in range(15): for j in range(20): # pt1 = (160, 160) # x = pt1[0] # y = pt1[1] x = j * 30 y = i * 30 pt1 = (x, y) points1.append(pt1) for i in range(images.shape[0] - 1): image1 = images[i, None] image2 = images[i + 1, None] flow_low, flow_up = model(image1, image2, iters=20, test_mode=True) flow_up_numpy = flow_up[0].permute(1, 2, 0).cpu().numpy() viz(image1, image2, flow_up, uv1=points1) k = cv2.waitKey(200) * 0xFF if k == 27: break
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 __init__(self, type): self.type = type if type == 'farneback': self.type = 'farneback' if type == 'raft': parser = argparse.ArgumentParser() parser.add_argument('--model', nargs='?', const='raft-models/raft-things.pth', type=str, help="restore checkpoint") parser.add_argument('--path', nargs='?', const='frames', type=int, help="dataset for evaluation") 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() args.model = 'raft-models/raft-things.pth' self.flow_model = torch.nn.DataParallel(RAFT(args)) self.flow_model.load_state_dict(torch.load(args.model)) self.flow_model = self.flow_model.module self.flow_model.to(DEVICE) self.flow_model.eval() if type == 'flownet': print('Flownet')
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)
grid = grid.cuda() vgrid = Variable(grid) + flo # scale grid to [-1,1] vgrid[:,0,:,:] = 2.0*vgrid[:,0,:,:].clone() / max(W-1,1)-1.0 vgrid[:,1,:,:] = 2.0*vgrid[:,1,:,:].clone() / max(H-1,1)-1.0 vgrid = vgrid.permute(0,2,3,1) output = nn.functional.grid_sample(x, vgrid) mask = torch.autograd.Variable(torch.ones(x.size()))#.cuda() if x.is_cuda: mask = mask.cuda() mask = nn.functional.grid_sample(mask, vgrid) mask[mask<0.9999] = 0 mask[mask>0] = 1 return output*mask model = torch.nn.DataParallel(RAFT(args)) 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)
def run_optical_flows(args): model = torch.nn.DataParallel(RAFT(args)) model.load_state_dict(torch.load(args.model)) model = model.module model.to(DEVICE) model.eval() basedir = "%s" % args.data_path # print(basedir) img_dir = glob.glob(basedir + '/images_*')[0] #basedir + '/images_*288' img_path_train = os.path.join(glob.glob(img_dir)[0], '%05d.png' % 0) img_train = cv2.imread(img_path_train) interval = 1 of_dir = os.path.join(basedir, 'flow_i%d' % interval) if not os.path.exists(of_dir): os.makedirs(of_dir) with torch.no_grad(): images = glob.glob(os.path.join(basedir, 'images/', '*.png')) + \ glob.glob(os.path.join(basedir, 'images/', '*.jpg')) images = load_image_list(images, args.input_flow_w) for i in range(images.shape[0] - 1): print(i) image1 = images[i, None] image2 = images[i + 1, None] _, flow_up_fwd = model(image1, image2, iters=20, test_mode=True) _, flow_up_bwd = model(image2, image1, iters=20, test_mode=True) flow_up_fwd = flow_up_fwd[0].cpu().numpy().transpose(1, 2, 0) flow_up_bwd = flow_up_bwd[0].cpu().numpy().transpose(1, 2, 0) img1 = cv2.resize( np.uint8( np.clip(image1[0].cpu().numpy().transpose(1, 2, 0), 0, 255)), (img_train.shape[1], img_train.shape[0]), cv2.INTER_LINEAR) img2 = cv2.resize( np.uint8( np.clip(image2[0].cpu().numpy().transpose(1, 2, 0), 0, 255)), (img_train.shape[1], img_train.shape[0]), cv2.INTER_LINEAR) fwd_flow = resize_flow(flow_up_fwd, img_train.shape[0], img_train.shape[1]) # fwd_flow = refinement_flow(fwd_flow, img1, img2) bwd_flow = resize_flow(flow_up_bwd, img_train.shape[0], img_train.shape[1]) # bwd_flow = refinement_flow(bwd_flow, img1, img2) fwd_mask, bwd_mask = compute_fwdbwd_mask(fwd_flow, bwd_flow) if VIZ: if not os.path.exists('./viz_flow'): os.makedirs('./viz_flow') if not os.path.exists('./viz_warp_imgs'): os.makedirs('./viz_warp_imgs') plt.figure(figsize=(12, 6)) plt.subplot(2, 3, 1) plt.imshow(img1) plt.subplot(2, 3, 4) plt.imshow(img2) plt.subplot(2, 3, 2) plt.imshow(flow_viz.flow_to_image(fwd_flow) / 255.) plt.subplot(2, 3, 3) plt.imshow(flow_viz.flow_to_image(bwd_flow) / 255.) plt.subplot(2, 3, 5) plt.imshow( flow_viz.flow_to_image(fwd_flow) / 255. * np.float32(fwd_mask[..., np.newaxis])) plt.subplot(2, 3, 6) plt.imshow( flow_viz.flow_to_image(bwd_flow) / 255. * np.float32(bwd_mask[..., np.newaxis])) plt.savefig('./viz_flow/%02d.jpg' % i) plt.close() warped_im2 = warp_flow(img2, fwd_flow) warped_im0 = warp_flow(img1, bwd_flow) cv2.imwrite('./viz_warp_imgs/im_%05d.jpg' % (i), img1[..., ::-1]) cv2.imwrite('./viz_warp_imgs/im_%05d_fwd.jpg' % (i), warped_im2[..., ::-1]) cv2.imwrite('./viz_warp_imgs/im_%05d_bwd.jpg' % (i + 1), warped_im0[..., ::-1]) np.savez(os.path.join(of_dir, '%05d_fwd.npz' % i), flow=fwd_flow, mask=fwd_mask) np.savez(os.path.join(of_dir, '%05d_bwd.npz' % (i + 1)), flow=bwd_flow, mask=bwd_mask)
def train(args): model = nn.DataParallel(RAFT(args), device_ids=args.gpus) print("Parameter Count: %d" % count_parameters(model)) if args.restore_ckpt is not None: model.load_state_dict(torch.load(args.restore_ckpt), strict=False) model.cuda() model.train() if args.stage != 'chairs': model.module.freeze_bn() train_loader = datasets.fetch_dataloader(args) optimizer, scheduler = fetch_optimizer(args, model) total_steps = 0 scaler = GradScaler(enabled=args.mixed_precision) logger = Logger(model, scheduler) VAL_FREQ = 5000 add_noise = True should_keep_training = True while should_keep_training: for i_batch, data_blob in enumerate(train_loader): optimizer.zero_grad() image1, image2, flow, valid = [x.cuda() for x in data_blob] # show_image(image1[0]) # show_image(image2[0]) if args.add_noise: stdv = np.random.uniform(0.0, 5.0) image1 = (image1 + stdv * torch.randn(*image1.shape).cuda()).clamp( 0.0, 255.0) image2 = (image2 + stdv * torch.randn(*image2.shape).cuda()).clamp( 0.0, 255.0) flow_predictions = model(image1, image2, iters=args.iters) loss, metrics = sequence_loss(flow_predictions, flow, valid) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip) scaler.step(optimizer) scheduler.step() scaler.update() logger.push(metrics) if total_steps % VAL_FREQ == VAL_FREQ - 1: PATH = 'checkpoints/%d_%s.pth' % (total_steps + 1, args.name) torch.save(model.state_dict(), PATH) results = {} for val_dataset in args.validation: if val_dataset == 'chairs': results.update(evaluate.validate_chairs(model.module)) elif val_dataset == 'sintel': results.update(evaluate.validate_sintel(model.module)) elif val_dataset == 'kitti': results.update(evaluate.validate_kitti(model.module)) logger.write_dict(results) model.train() if args.stage != 'chairs': model.module.freeze_bn() total_steps += 1 if total_steps > args.num_steps: should_keep_training = False break logger.close() PATH = 'checkpoints/%s.pth' % args.name torch.save(model.state_dict(), PATH) return PATH
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 train(args): model = nn.DataParallel(RAFT(args), device_ids=args.gpus) print("Parameter Count: %d" % count_parameters(model)) #if args.restore_ckpt is not None: # model.load_state_dict(torch.load(args.restore_ckpt), strict=False) model.cuda() model.train() if args.stage != 'chairs': model.module.freeze_bn() train_loader = datasets_self.fetch_dataloader(args) optimizer, scheduler = fetch_optimizer(args, model) total_steps = 0 scaler = GradScaler(enabled=args.mixed_precision) logger = Logger(model, scheduler) add_noise = True should_keep_training = True while should_keep_training: for i_batch, data_blob in enumerate(train_loader): optimizer.zero_grad() image1, image2, img1_orig, img2_orig = [ x.cuda() for x in data_blob ] if args.add_noise: stdv = np.random.uniform(0.0, 5.0) image1 = (image1 + stdv * torch.randn(*image1.shape).cuda()).clamp( 0.0, 255.0) image2 = (image2 + stdv * torch.randn(*image2.shape).cuda()).clamp( 0.0, 255.0) print(image1.shape) print(img1_orig.shape) flow_predictions = model(image1, image2, iters=args.iters) loss = sequence_loss(img1_orig, img2_orig, flow_predictions, args.gamma) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), args.clip) scaler.step(optimizer) scheduler.step() scaler.update() #logger.push(metrics) print('processing step ' + str(total_steps) + ', loss ' + str(loss.item())) if total_steps % VAL_FREQ == 0: print('i=' + str(total_steps) + ' loss=' + str(loss)) PATH = args.save_ckpt + '/%d_%s.pth' % (total_steps + 1, args.name) print(PATH) torch.save(model.state_dict(), PATH) #results = {} #for val_dataset in args.validation: # if val_dataset == 'chairs': # results.update(evaluate.validate_chairs(model.module)) # elif val_dataset == 'sintel': # results.update(evaluate.validate_sintel(model.module)) # elif val_dataset == 'kitti': # results.update(evaluate.validate_kitti(model.module)) # elif val_dataset == 'ESPRIT': # results.update(evaluate.validate_ESPRIT(model.module)) #logger.write_dict(results) model.train() if args.stage != 'chairs': model.module.freeze_bn() total_steps += 1 if total_steps > args.num_steps: should_keep_training = False break #logger.close() PATH = args.save_ckpt + '/%s.pth' % args.name torch.save(model.state_dict(), PATH) return PATH