Skip to content

Instantly share code, notes, and snippets.

@ClementPinard
Last active September 2, 2020 16:57
Show Gist options
  • Save ClementPinard/e7353dee56faf3f4c62f0753a2703568 to your computer and use it in GitHub Desktop.
Save ClementPinard/e7353dee56faf3f4c62f0753a2703568 to your computer and use it in GitHub Desktop.
import torch, kornia, h5py, imageio
import matplotlib.pyplot as plt
import numpy as np
# This is a failure mode, with different aspect ratios
fname1 = '47698078_3766965066'
fname2 = '18698491_4586522698'
# This works
# fname1 = '271147142_778c4e7999_o'
# fname2 = '275191466_a33f8c30b7_o'
def get_K_Rt(K_: [3, 3], R: [3, 3], T: [3]):
B = 1
# add batch dimension
K = K_[None] # [B, 3, 3]
R = R[None] # [B, 3, 3]
T = T[None, :, None] # [B, 3, 1]
Rt_3x4 = kornia.projection_from_Rt(R, T)
Rt = torch.zeros(B, 4, 4)
Rt[:, :3, :] = Rt_3x4
Rt[:, 3, 3] = 1.
K_4x4 = torch.eye(4)[None]
K_4x4[:, :3, :3] = K
return K_4x4, Rt
def read_data(fname):
''' read the files and return a
* image tensor
* depth tensor
* kornia.PinholeCamera instance
'''
img = imageio.imread(f'data/images/{fname}.jpg')
img = torch.from_numpy(img).to(torch.float32) / 255
with h5py.File(f'data/depth_maps/{fname}.h5', 'r') as hdf:
depth = torch.from_numpy(hdf['depth'][()]).to(torch.float32)
with h5py.File(f'data/calibration/calibration_{fname}.h5', 'r') as hdf:
K = torch.from_numpy(hdf['K'][()]).to(torch.float32)
R = torch.from_numpy(hdf['R'][()]).to(torch.float32)
T = torch.from_numpy(hdf['T'][()]).to(torch.float32)
K, Rt = get_K_Rt(K, R, T)
return img.unsqueeze(0), depth.unsqueeze(0), K, Rt
img1, dep1, K1, Rt1 = read_data(fname1)
img2, dep2, K2, Rt2 = read_data(fname2)
height1 = torch.Tensor(img1.shape[1])
width1 = torch.Tensor(img1.shape[2])
height2 = torch.Tensor(img2.shape[1])
width2 = torch.Tensor(img2.shape[2])
cam1 = kornia.geometry.camera.PinholeCamera(K1, Rt1, height1, width1)
cam2 = kornia.geometry.camera.PinholeCamera(K2, Rt2, height2, width2)
img1_bchw = img1.permute(0, 3, 1, 2)
img2_bchw = img2.permute(0, 3, 1, 2)
warp_12 = kornia.geometry.warp.depth_warp(cam2, cam1, dep1[:, None], img2_bchw, dep1.shape[1], dep1.shape[2])
warp_21 = kornia.geometry.warp.depth_warp(cam1, cam2, dep2[:, None], img1_bchw, dep2.shape[1], dep2.shape[2])
fig, axes = plt.subplots(2, 3, constrained_layout=True)
axes[0, 0].imshow(img1[0].numpy())
axes[1, 0].imshow(img2[0].numpy())
axes[0, 1].imshow(dep1[0].numpy())
axes[1, 1].imshow(dep2[0].numpy())
axes[0, 2].imshow(warp_12[0].permute(1, 2, 0).numpy())
axes[1, 2].imshow(warp_21[0].permute(1, 2, 0).numpy())
plt.show()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment