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()
