def __call__(self, img1, img2, flow):
if self.padding:
img1, img2, flow = self.reflection_padding(img1, img2, flow)
if self.show_aug:
plt.subplot(221)
plt.title('1 (Input)')
show_image(img1, subplote=True)
if self.show_aug:
plt.subplot(222)
plt.title('flow (Input)')
show_image(flow_to_image(flow), subplote=True)
img1, img2 = self.color_transform(img1, img2)
img1, img2 = self.eraser_transform(img1, img2)
img1, img2, flow = self.spatial_transform(img1, img2, flow)
if self.show_aug:
plt.subplot(223)
plt.title('1 (spatial)')
show_image(img1, subplote=True)
if self.show_aug:
plt.subplot(224)
plt.title('flow (spatial)')
show_image(flow_to_image(flow))
img1 = np.ascontiguousarray(img1)
img2 = np.ascontiguousarray(img2)
flow = np.ascontiguousarray(flow)
return img1, img2, flow
def run_pair_tensor(model, image1, image2, tresh, force, strength, smooth,
intepolation):
with torch.no_grad():
padder = InputPadder(image1.shape)
image1, image2 = padder.pad(image1, image2)
flow_low, flow_up = model(image1, image2, iters=20, test_mode=True)
image1 = padder.unpad(image1)
flow_up = padder.unpad(flow_up)
blurImage, normalizedFlow = MotionBlur(image1, flow_up, tresh, force,
strength, smooth, intepolation)
#return image1, blurImage, flow_up, normalizedFlow
#img = image1[0].permute(1,2,0).cpu().numpy()
blur = blurImage[0].permute(1, 2, 0).cpu().numpy()
flo1 = flow_up[0].permute(1, 2, 0).cpu().numpy()
flo2 = normalizedFlow.permute(1, 2, 0).cpu().numpy()
flo1 = flow_viz.flow_to_image(flo1)
flo2 = flow_viz.flow_to_image(flo2)
del image1, image2, padder, flow_low, flow_up, blurImage, normalizedFlow
return blur, flo1, flo2
def write_vls_eval(self, data_blob, flowpred, selector, evalidx, step):
img1 = data_blob['img1'][0].permute([1, 2, 0]).numpy().astype(np.uint8)
insmap = data_blob['insmap'][0].squeeze().numpy()
figmask = tensor2disp(selector, vmax=1, viewind=0)
insvls = Image.fromarray(vls_ins(img1, insmap))
flowgtvls = Image.fromarray(
flow_to_image(data_blob['flowmap'][0].permute([1, 2, 0]).numpy(),
rad_max=10))
flowpredvls = Image.fromarray(
flow_to_image(flowpred[0].detach().cpu().permute([1, 2,
0]).numpy(),
rad_max=10))
img_val_up = np.concatenate(
[np.array(figmask), np.array(insvls)], axis=1)
img_val_down = np.concatenate(
[np.array(flowgtvls), np.array(flowpredvls)], axis=1)
img_val = np.concatenate(
[np.array(img_val_up),
np.array(img_val_down)], axis=0)
self.writer.add_image('predvls_eval_{}'.format(str(evalidx).zfill(2)),
(torch.from_numpy(img_val).float() /
255).permute([2, 0, 1]), step)
def write_vls(self, image1, image2, flowgt, flow_predictions, valid, depth):
img1 = image1[0].detach().cpu().detach().permute([1, 2, 0]).numpy().astype(np.uint8)
img2 = image2[0].detach().cpu().detach().permute([1, 2, 0]).numpy().astype(np.uint8)
validnp = valid[0].detach().cpu().numpy() == 1
depthnp = depth[0].squeeze().numpy()
flow_pred = flow_to_image(flow_predictions[-1][0].permute([1, 2, 0]).detach().cpu().numpy(), rad_max=150)
flow_gt = flow_to_image(flowgt[0].permute([1, 2, 0]).detach().cpu().numpy(), rad_max=150)
h, w = image1.shape[2::]
xx, yy = np.meshgrid(range(w), range(h), indexing='xy')
pixelloc = np.stack([xx, yy], axis=0)
pixelloc = torch.from_numpy(pixelloc).float() + flow_predictions[-1][0].detach().cpu()
pixelloc[0, :, :] = ((pixelloc[0, :, :] / w) - 0.5) * 2
pixelloc[1, :, :] = ((pixelloc[1, :, :] / h) - 0.5) * 2
pixelloc = pixelloc.permute([1, 2, 0])
xxf = xx[validnp]
yyf = yy[validnp]
depthf = depthnp[validnp]
xxf_oview = flowgt[0, 0].detach().cpu().numpy()[validnp] + xxf
yyf_oview = flowgt[0, 1].detach().cpu().numpy()[validnp] + yyf
vmax = 0.15
tnp = 1 / depthf / vmax
tnp = self.cm(tnp)
fig = plt.figure(figsize=(16, 2.5))
canvas = FigureCanvasAgg(fig)
fig.add_subplot(1, 2, 1)
plt.scatter(xxf, yyf, 1, tnp)
plt.imshow(img1)
fig.add_subplot(1, 2, 2)
plt.scatter(xxf_oview, yyf_oview, 1, tnp)
plt.imshow(img2)
plt.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=0, hspace=0)
canvas.draw()
buf = canvas.buffer_rgba()
plt.close()
X = np.asarray(buf)
X = np.array(Image.fromarray(X).resize([w * 2, h], Image.BILINEAR))
image1_recon = torch.nn.functional.grid_sample(image2[0, :, :, :].detach().unsqueeze(0).cpu(), pixelloc.unsqueeze(0), mode='bilinear', align_corners=False)
image1_recon = image1_recon[0].permute([1, 2, 0]).numpy().astype(np.uint8)
img_mid = np.concatenate([flow_pred, flow_gt], axis=1)
img_down = np.concatenate([img1, image1_recon], axis=1)
img_vls = np.concatenate([X[:, :, 0:3], img_mid, img_down], axis=0)
self.writer.add_image('img_vls', (torch.from_numpy(img_vls).float() / 255).permute([2, 0, 1]), self.total_steps)
def write_vls(self,
data_blob,
flowpred,
mDdoutputs,
selector,
step,
X=None):
img1 = data_blob['img1'][0].permute([1, 2, 0]).numpy().astype(np.uint8)
insmap = data_blob['insmap'][0].squeeze().numpy()
figmask = tensor2disp(selector, vmax=1, viewind=0)
insvls = Image.fromarray(vls_ins(img1, insmap))
flowgtvls = Image.fromarray(
flow_to_image(data_blob['flowmap'][0].permute([1, 2, 0]).numpy(),
rad_max=10))
flowpredvls = Image.fromarray(
flow_to_image(flowpred[-1][0].detach().cpu().permute([1, 2,
0]).numpy(),
rad_max=10))
depthgtvls = tensor2disp(1 / data_blob['depthmap'],
vmax=0.15,
viewind=0)
depthpredvls = tensor2disp(1 / mDdoutputs[('mDepth', 0)],
vmax=0.15,
viewind=0)
img_val_up = np.concatenate(
[np.array(figmask), np.array(insvls)], axis=1)
img_val_down = np.concatenate(
[np.array(flowgtvls), np.array(flowpredvls)], axis=1)
img_val_mid = np.concatenate(
[np.array(depthgtvls),
np.array(depthpredvls)], axis=1)
img_val = np.concatenate([
np.array(img_val_up),
np.array(img_val_mid),
np.array(img_val_down)
],
axis=0)
self.writer.add_image('predvls', (torch.from_numpy(img_val).float() /
255).permute([2, 0, 1]), step)
if X is not None:
X = X[:, :, 0:3]
self.writer.add_image('Sampling', (torch.from_numpy(X).float() /
255).permute([2, 0, 1]), step)
def write_vls_eval(self, data_blob, outputs, tagname, step):
img1 = data_blob['img1'][0].permute([1, 2, 0]).numpy().astype(np.uint8)
img2 = data_blob['img2'][0].permute([1, 2, 0]).numpy().astype(np.uint8)
inputdepth = tensor2disp(1 / data_blob['depthmap'], vmax=1, viewind=0)
depthpredvls = tensor2disp(1 / outputs[('depth', 2)],
vmax=1,
viewind=0)
flowvls = flow_to_image(
outputs[('flowpred', 2)][0].detach().cpu().permute([1, 2,
0]).numpy(),
rad_max=10)
imgrecon = tensor2rgb(outputs[('reconImg', 2)], viewind=0)
img_val_up = np.concatenate([np.array(img1), np.array(img2)], axis=1)
img_val_mid2 = np.concatenate(
[np.array(inputdepth),
np.array(depthpredvls)], axis=1)
img_val_mid3 = np.concatenate(
[np.array(imgrecon), np.array(flowvls)], axis=1)
img_val = np.concatenate([
np.array(img_val_up),
np.array(img_val_mid2),
np.array(img_val_mid3)
],
axis=0)
self.writer.add_image('{}_predvls'.format(tagname),
(torch.from_numpy(img_val).float() /
255).permute([2, 0, 1]), step)
X = self.vls_sampling(img1, img2, data_blob['depthmap'],
data_blob['flowpred'], outputs)
self.writer.add_image('{}_X'.format(tagname),
(torch.from_numpy(X).float() / 255).permute(
[2, 0, 1]), step)
def write_vls(self, data_blob, outputs, flowselector, step):
img1 = data_blob['img1'][0].permute([1, 2, 0]).numpy().astype(np.uint8)
img2 = data_blob['img2'][0].permute([1, 2, 0]).numpy().astype(np.uint8)
insmap = data_blob['insmap'][0].squeeze().numpy()
figmask_flow = tensor2disp(flowselector, vmax=1, viewind=0)
insvls = vls_ins(img1, insmap)
depthpredvls = tensor2disp(1 / outputs[('depth', 2)],
vmax=0.15,
viewind=0)
flowvls = flow_to_image(
outputs[('flowpred', 2)][0].detach().cpu().permute([1, 2,
0]).numpy(),
rad_max=10)
imgrecon = tensor2rgb(outputs[('reconImg', 2)], viewind=0)
img_val_up = np.concatenate([np.array(insvls), np.array(img2)], axis=1)
img_val_mid2 = np.concatenate(
[np.array(depthpredvls),
np.array(figmask_flow)], axis=1)
img_val_mid3 = np.concatenate(
[np.array(imgrecon), np.array(flowvls)], axis=1)
img_val = np.concatenate([
np.array(img_val_up),
np.array(img_val_mid2),
np.array(img_val_mid3)
],
axis=0)
self.writer.add_image('predvls', (torch.from_numpy(img_val).float() /
255).permute([2, 0, 1]), step)
def write_vls(self, image1, image2, flowgt, flow_predictions):
img1 = image1[0].cpu().detach().permute([1, 2,
0]).numpy().astype(np.uint8)
img2 = image2[0].cpu().detach().permute([1, 2,
0]).numpy().astype(np.uint8)
flow_pred = flow_to_image(flow_predictions[-1][0].permute(
[1, 2, 0]).detach().cpu().numpy())
flow_gt = flow_to_image(flowgt[0].permute([1, 2,
0]).detach().cpu().numpy())
img_up = np.concatenate([img1, img2], axis=1)
img_down = np.concatenate([flow_pred, flow_gt], axis=1)
img_vls = np.concatenate([img_up, img_down], axis=0)
self.writer.add_image('img_vls', (torch.from_numpy(img_vls).float() /
255).permute([2, 0, 1]),
self.total_steps)
def write_vls(self, image1, image2, flowgt, flow_predictions):
image1 = image1.detach()
image1 = torch.stack(
[image1[:, 2, :, :], image1[:, 1, :, :], image1[:, 0, :, :]],
dim=1)
image2 = image2.detach()
image2 = torch.stack(
[image2[:, 2, :, :], image2[:, 1, :, :], image2[:, 0, :, :]],
dim=1)
img1 = image1[0].cpu().detach().permute([1, 2,
0]).numpy().astype(np.uint8)
img2 = image2[0].cpu().detach().permute([1, 2,
0]).numpy().astype(np.uint8)
flow_pred = flow_to_image(flow_predictions[-1][0].permute(
[1, 2, 0]).detach().cpu().numpy())
flow_gt = flow_to_image(flowgt[0].permute([1, 2,
0]).detach().cpu().numpy())
h, w = image1.shape[2::]
xx, yy = np.meshgrid(range(w), range(h), indexing='xy')
pixelloc = np.stack([xx, yy], axis=0)
pixelloc = torch.from_numpy(
pixelloc).float() + flow_predictions[-1][0].detach().cpu()
pixelloc[0, :, :] = ((pixelloc[0, :, :] / w) - 0.5) * 2
pixelloc[1, :, :] = ((pixelloc[1, :, :] / h) - 0.5) * 2
pixelloc = pixelloc.permute([1, 2, 0])
image1_recon = torch.nn.functional.grid_sample(
image2[0, :, :, :].detach().unsqueeze(0).cpu(),
pixelloc.unsqueeze(0),
mode='bilinear',
align_corners=False)
image1_recon = image1_recon[0].permute([1, 2,
0]).numpy().astype(np.uint8)
img_up = np.concatenate([img1, img2], axis=1)
img_mid = np.concatenate([flow_pred, flow_gt], axis=1)
img_down = np.concatenate([img1, image1_recon], axis=1)
img_vls = np.concatenate([img_up, img_mid, img_down], axis=0)
self.writer.add_image('img_vls', (torch.from_numpy(img_vls).float() /
255).permute([2, 0, 1]),
self.total_steps)
def write_vls_eval(self, data_blob, outputs, tagname, step):
img1 = data_blob['img1'][0].permute([1, 2, 0]).numpy().astype(np.uint8)
insmap = data_blob['insmap'][0].squeeze().numpy()
insvls = vls_ins(img1, insmap)
flowvls = flow_to_image(outputs[-1][0].detach().cpu().permute([1, 2, 0]).numpy(), rad_max=50)
img_val_up = np.concatenate([np.array(insvls), np.array(flowvls)], axis=1)
self.writer.add_image('{}_predvls'.format(tagname), (torch.from_numpy(img_val_up).float() / 255).permute([2, 0, 1]), step)
def write_vls(self, data_blob, outputs, step):
img1 = data_blob['img1'][0].permute([1, 2, 0]).numpy().astype(np.uint8)
img2 = data_blob['img2'][0].permute([1, 2, 0]).numpy().astype(np.uint8)
insmap = data_blob['insmap'][0].squeeze().numpy()
insvls = vls_ins(img1, insmap)
depthpredvls = tensor2disp(1 / data_blob['mdDepth_pred'],
vmax=0.15,
viewind=0)
imgrecon = tensor2rgb(outputs[('img1_recon', 2)][:, -1], viewind=0)
flow_RAFT = Image.fromarray(
flow_to_image(data_blob['flowpred'][0].permute([1, 2,
0]).cpu().numpy()))
flow_pred = Image.fromarray(
flow_to_image(outputs[('flowpred',
2)][0].permute([1, 2,
0]).detach().cpu().numpy()))
img_val_up = np.concatenate([np.array(insvls), np.array(img2)], axis=1)
img_val_mid = np.concatenate(
[np.array(depthpredvls),
np.array(imgrecon)], axis=1)
img_val_flow = np.concatenate(
[np.array(flow_RAFT), np.array(flow_pred)], axis=1)
img_val = np.concatenate([
np.array(img_val_up),
np.array(img_val_mid),
np.array(img_val_flow)
],
axis=0)
self.writer.add_image('predvls', (torch.from_numpy(img_val).float() /
255).permute([2, 0, 1]), step)
X = self.vls_sampling(np.array(insvls), img2, data_blob['depthvls'],
outputs)
self.writer.add_image('X', (torch.from_numpy(X).float() / 255).permute(
[2, 0, 1]), step)
X = self.vls_objmvment(np.array(insvls), data_blob['insmap'],
data_blob['posepred'])
self.writer.add_image('objmvment', (torch.from_numpy(X).float() /
255).permute([2, 0, 1]), step)
def viz(img, flo):
img = img[0].permute(1, 2, 0).cpu().numpy()
flo = flo[0].permute(1, 2, 0).cpu().numpy()
# map flow to rgb image
flo = flow_viz.flow_to_image(flo)
img_flo = np.concatenate([img, flo], axis=0)
cv2.imshow('image', img_flo[:, :, [2, 1, 0]] / 255.0)
cv2.waitKey()
def display(image1, image2, flow):
image1 = image1.permute(1, 2, 0).cpu().numpy() / 255.0
image2 = image2.permute(1, 2, 0).cpu().numpy() / 255.0
flow = flow.permute(1, 2, 0).cpu().numpy()
flow_image = flow_viz.flow_to_image(flow)
flow_image = cv2.resize(flow_image, (image1.shape[1], image1.shape[0]))
cv2.imshow('image1', image1[..., ::-1])
cv2.imshow('image2', image2[..., ::-1])
cv2.imshow('flow', flow_image[..., ::-1])
cv2.waitKey()
def viz(img, flo):
img = img[0].permute(1, 2, 0).cpu().numpy()
flo = flo[0].permute(1, 2, 0).cpu().numpy()
# map flow to rgb image
flo = flow_viz.flow_to_image(flo)
img_flo = np.concatenate([img, flo], axis=0)
# import matplotlib.pyplot as plt
# plt.imshow(img_flo / 255.0)
# plt.show()
cv2.imshow('image', img_flo[:, :, [2, 1, 0]] / 255.0)
cv2.waitKey(1)
def write_vls(self, data_blob, est_objpose, estflow, selector):
img1 = data_blob['img1'][0].permute([1, 2, 0]).numpy().astype(np.uint8)
img2 = data_blob['img2'][0].permute([1, 2, 0]).numpy().astype(np.uint8)
depthmap = data_blob['depthmap'][0].squeeze().numpy()
insmap = data_blob['insmap'][0].squeeze().numpy()
intrinsic = data_blob['intrinsic'][0].squeeze().numpy()
est_objpose = est_objpose[0].detach().cpu().numpy()
flowmap = data_blob['flowmap'][0].squeeze().numpy()
objpose_gt = data_blob['poses'][0].squeeze().numpy()
figmask = tensor2disp(selector, vmax=1, viewind=0)
insvls = Image.fromarray(vls_ins(img1, insmap))
posevls = self.plot_scattervls(img1, img2, depthmap, insmap, intrinsic,
objpose_gt[0], objpose_gt, est_objpose,
flowmap)
flowvls = flow_to_image(data_blob['flowmap'][0].cpu().permute(
[1, 2, 0]).numpy(),
rad_max=15)
estflowvls = flow_to_image(estflow[0].cpu().permute([1, 2, 0]).numpy(),
rad_max=15)
img_val_up = np.concatenate(
[np.array(figmask), np.array(insvls)], axis=1)
img_val_down = np.concatenate(
[np.array(flowvls), np.array(estflowvls)], axis=1)
img_val = np.concatenate(
[np.array(img_val_up),
np.array(img_val_down)], axis=0)
posevls = np.array(posevls[:, :, 0:3])
self.writer.add_image('img_val', (torch.from_numpy(img_val).float() /
255).permute([2, 0, 1]),
self.total_steps)
self.writer.add_image('posevls', (torch.from_numpy(posevls).float() /
255).permute([2, 0, 1]),
self.total_steps)
def write_vls_eval(self, data_blob, outputs, tagname, step):
img1 = data_blob['img1'][0].permute([1, 2, 0]).numpy().astype(np.uint8)
img2 = data_blob['img2'][0].permute([1, 2, 0]).numpy().astype(np.uint8)
insmap = data_blob['insmap'][0].squeeze().numpy()
insvls = vls_ins(img1, insmap)
depthpredvls = tensor2disp(1 / outputs['depth_predictions'][-1],
vmax=0.15,
viewind=0)
depthgtvls = tensor2disp(1 / data_blob['depthmap'],
vmax=0.15,
viewind=0)
flowvls = flow_to_image(outputs['flowpred'][0].detach().cpu().permute(
[1, 2, 0]).numpy(),
rad_max=50)
imgrecon = tensor2rgb(outputs['img1_recon'], viewind=0)
img_val_up = np.concatenate([np.array(insvls), np.array(img2)], axis=1)
img_val_mid2 = np.concatenate(
[np.array(depthpredvls),
np.array(depthgtvls)], axis=1)
img_val_mid3 = np.concatenate(
[np.array(imgrecon), np.array(flowvls)], axis=1)
img_val = np.concatenate([
np.array(img_val_up),
np.array(img_val_mid2),
np.array(img_val_mid3)
],
axis=0)
self.writer.add_image('{}_predvls'.format(tagname),
(torch.from_numpy(img_val).float() /
255).permute([2, 0, 1]), step)
X = self.vls_sampling(np.array(insvls), img2, data_blob['depthvls'],
data_blob['flowmap'], data_blob['insmap'],
outputs)
self.writer.add_image('{}_X'.format(tagname),
(torch.from_numpy(X).float() / 255).permute(
[2, 0, 1]), step)
def validate_kitti_colorjitter(gpu,
model,
args,
ngpus_per_node,
eval_entries,
iters=24):
""" Peform validation using the KITTI-2015 (train) split """
interval = np.floor(len(eval_entries) / ngpus_per_node).astype(
np.int).item()
if gpu == ngpus_per_node - 1:
stidx = int(interval * gpu)
edidx = len(eval_entries)
else:
stidx = int(interval * gpu)
edidx = int(interval * (gpu + 1))
print("Initialize Instance on Gpu %d, from %d to %d, total %d" %
(gpu, stidx, edidx, len(eval_entries)))
from tqdm import tqdm
model.eval()
model.cuda(gpu)
with torch.no_grad():
for val_id, entry in enumerate(
tqdm(remove_dup(eval_entries[stidx:edidx]))):
seq, index = entry.split(' ')
index = int(index)
img1path = os.path.join(args.dataset, seq,
'rgb_{}.png'.format(str(index).zfill(5)))
img2path = os.path.join(
args.dataset, seq,
'rgb_{}.png'.format(str(index + 1).zfill(5)))
if not os.path.exists(img2path) and not os.path.exists(
img2path.replace('.png', '.jpg')):
img2path = img1path
if not os.path.exists(img1path):
img1path = img1path.replace('.png', '.jpg')
img2path = img2path.replace('.png', '.jpg')
image1 = frame_utils.read_gen(img1path)
image2 = frame_utils.read_gen(img2path)
image1 = np.array(image1).astype(np.uint8)
image2 = np.array(image2).astype(np.uint8)
image1 = torch.from_numpy(image1).permute([2, 0, 1]).float()
image2 = torch.from_numpy(image2).permute([2, 0, 1]).float()
svfold = os.path.join(
args.exportroot,
seq,
)
svpath = os.path.join(args.exportroot, seq,
'{}.png'.format(str(index).zfill(5)))
os.makedirs(svfold, exist_ok=True)
image1 = image1[None].cuda(gpu)
image2 = image2[None].cuda(gpu)
flow_low, flow_pr = model(image1,
image2,
iters=iters,
test_mode=True)
flow = flow_pr.squeeze(0)
flow_numpy = flow.cpu().permute(1, 2, 0).numpy()
frame_utils.writeFlowKITTI(svpath, flow_numpy)
if args.dovls:
vlsflow = Image.fromarray(flow_viz.flow_to_image(flow_numpy))
vlsrgb1 = tensor2rgb(image1 / 255.0, viewind=0)
vlsrgb2 = tensor2rgb(image2 / 255.0, viewind=0)
w, h = vlsrgb2.size
xx, yy = np.meshgrid(range(w), range(h), indexing='xy')
xxf = xx.flatten()
yyf = yy.flatten()
rndidx = np.random.randint(0, xxf.shape[0], 100)
xxf_rnd = xxf[rndidx]
yyf_rnd = yyf[rndidx]
flowx = flow_numpy[yyf_rnd, xxf_rnd, 0]
flowy = flow_numpy[yyf_rnd, xxf_rnd, 1]
fig = plt.figure(figsize=(12, 9))
plt.subplot(2, 1, 1)
plt.scatter(xxf_rnd, yyf_rnd, 1, 'r')
plt.imshow(vlsrgb1)
plt.subplot(2, 1, 2)
plt.scatter(flowx + xxf_rnd, flowy + yyf_rnd, 1, 'r')
plt.imshow(vlsrgb2)
plt.savefig(
os.path.join(
"/media/shengjie/disk1/Prediction/nyuv2_flow_vls",
"{}_{}.jpg".format(seq,
str(index).zfill(5))))
plt.close()
# combined_left = np.concatenate([np.array(vlsrgb1), np.array(vlsrgb2)], axis=0)
# combined_right = np.concatenate([np.array(vlsflow), np.array(vlsflow)], axis=0)
# combined = np.concatenate([combined_left, combined_right], axis=1)
# vls_sv_root = os.makedirs("/media/shengjie/disk1/Prediction/nyuv2_flow_vls", exist_ok=True)
# Image.fromarray(combined).save(os.path.join("/media/shengjie/disk1/Prediction/nyuv2_flow_vls", "{}_{}.jpg".format(seq, str(index).zfill(5))))
return
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 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}
