def create_kitti_submission(model, iters=24, output_path='kitti_submission'):
""" 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)
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()
output_filename = os.path.join(output_path, frame_id)
frame_utils.writeFlowKITTI(output_filename, flow)
def fetch_dataloader(args):
""" Create the data loader for the corresponding training set """
if args.dataset == 'chairs':
train_dataset = datasets.FlyingChairs(args, image_size=args.image_size)
elif args.dataset == 'things':
clean_dataset = datasets.SceneFlow(args,
image_size=args.image_size,
dstype='frames_cleanpass')
final_dataset = datasets.SceneFlow(args,
image_size=args.image_size,
dstype='frames_finalpass')
train_dataset = clean_dataset + final_dataset
elif args.dataset == 'sintel':
clean_dataset = datasets.MpiSintel(args,
image_size=args.image_size,
dstype='clean')
final_dataset = datasets.MpiSintel(args,
image_size=args.image_size,
dstype='final')
train_dataset = clean_dataset + final_dataset
elif args.dataset == 'kitti':
train_dataset = datasets.KITTI(args,
image_size=args.image_size,
is_val=False)
gpuargs = {'num_workers': 4, 'drop_last': True}
train_loader = DataLoader(train_dataset,
batch_size=args.batch_size,
pin_memory=True,
shuffle=True,
**gpuargs)
print('Training with %d image pairs' % len(train_dataset))
return train_loader
def validate_kitti(args, model, iters=32):
""" Evaluate trained model on KITTI (train) """
model.eval()
val_dataset = datasets.KITTI(args,
do_augument=False,
is_val=True,
do_pad=True)
with torch.no_grad():
epe_list, out_list = [], []
for i in range(len(val_dataset)):
image1, image2, flow_gt, valid_gt = val_dataset[i]
image1 = image1[None].cuda()
image2 = image2[None].cuda()
flow_gt = flow_gt.cuda()
valid_gt = valid_gt.cuda()
flow_predictions = model.module(image1, image2, iters=iters)
flow_pr = flow_predictions[-1][0]
epe = torch.sum((flow_pr - 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)
print("Validation KITTI: %f, %f" %
(np.mean(epe_list), 100 * np.mean(out_list)))
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 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}