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 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 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 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 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 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 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 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}
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()
#wrapimg2 = padder.unpad(wrapimg2[0]).permute(1, 2, 0).cpu().numpy()
# save result
subname = imfile1.split("/")
savename = os.path.join(args.result, str(scaling) + '_' + subname[-1])
diffimg = np.abs(image1 - warpimg1)
print(str(scaling) + ': ' + str(np.mean(diffimg[200:4100, 200:7480,:])))
print(str(scaling) + ': ' + str(i_loss.mean())
img_flo = 0.5*(image1 + warpimg1)
#flow = padder.unpad(flow_up[0]).permute(1, 2, 0).cpu().numpy()
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}