def __getitem__(self, idx):
intrinsics, _, _, _ = util.parse_intrinsics(os.path.join(self.instance_dir, "intrinsics.txt"),
trgt_sidelength=self.img_sidelength)
intrinsics = torch.Tensor(intrinsics).float()
rgb = data_util.load_rgb(self.color_paths[idx], sidelength=self.img_sidelength)
rgb = rgb.reshape(3, -1).transpose(1, 0)
pose = data_util.load_pose(self.pose_paths[idx])
if self.has_params:
params = data_util.load_params(self.param_paths[idx])
else:
params = np.array([0])
uv = np.mgrid[0:self.img_sidelength, 0:self.img_sidelength].astype(np.int32)
uv = torch.from_numpy(np.flip(uv, axis=0).copy()).long()
uv = uv.reshape(2, -1).transpose(1, 0)
sample = {
"instance_idx": torch.Tensor([self.instance_idx]).squeeze(),
"rgb": torch.from_numpy(rgb).float(),
"pose": torch.from_numpy(pose).float(),
"uv": uv,
"param": torch.from_numpy(params).float(),
"intrinsics": intrinsics
}
return sample
def __init__(self, pose_dir):
super().__init__()
all_pose_paths = sorted(glob(os.path.join(pose_dir, '*.txt')))
self.all_poses = [
torch.from_numpy(data_util.load_pose(path))
for path in all_pose_paths
]
def read_view_tuple(self, idx):
gt_rgb = self.load_rgb(self.all_color[idx])
pose = data_util.load_pose(self.all_poses[idx])
this_view = {
'gt_rgb': torch.from_numpy(gt_rgb),
'pose': torch.from_numpy(pose)
}
return this_view
def __getitem__(self, idx):
img = data_util.load_rgb(self.rgb_paths[idx],
sidelength=self.sidelength).transpose(
2, 0, 1)
pose = data_util.load_pose(self.pose_paths[idx])
sample = {'image': img, 'pose': pose.flatten()[:12]}
return sample
def get_nearest_neighbors_pose(train_pose_dir,
test_pose_dir,
sampling_pattern='skip_2',
metric='cos'):
if sampling_pattern != 'all':
skip_val = int(sampling_pattern.split('_')[-1])
else:
skip_val = 0
train_pose_files = sorted(glob(os.path.join(train_pose_dir, '*.txt')))
idcs = list(range(len(train_pose_files)))[::skip_val + 1]
train_pose_files = train_pose_files[::skip_val + 1]
test_pose_files = sorted(glob(os.path.join(test_pose_dir, '*.txt')))
train_poses = np.stack(
[data_util.load_pose(pose)[:3, 3] for pose in train_pose_files],
axis=0)
train_poses /= np.linalg.norm(train_poses, axis=1, keepdims=True)
test_poses = np.stack(
[data_util.load_pose(pose)[:3, 3] for pose in test_pose_files], axis=0)
test_poses /= np.linalg.norm(test_poses, axis=1, keepdims=True)
if metric == 'cos':
cos_distance_mat = test_poses.dot(
train_poses.T) # nxn matrix of cosine distances
nn_idcs = [
idcs[int(val)] for val in np.argmax(cos_distance_mat, axis=1)
]
elif metric == 'l2':
l2_distance_mat = np.linalg.norm(test_poses[:, None, :] -
train_poses[None, :, :],
axis=2)
nn_idcs = [
idcs[int(val)] for val in np.argmin(l2_distance_mat, axis=1)
]
return nn_idcs
def save_circles(model, results_dir, img_paths, num_renders=50):
circles_dir = os.path.join(results_dir, 'circles')
util.cond_mkdir(circles_dir)
print('Generating circle views around scene')
for j, img_path in enumerate(img_paths):
circle_dir = os.path.join(circles_dir, f'{j:03d}')
util.cond_mkdir(circle_dir)
# Copy reference image into directory
shutil.copy(img_path, os.path.join(circle_dir, '0000_ref_image.png'))
img = torch.from_numpy(
data_util.load_rgb(img_path).transpose(2, 0,
1)).to(device).unsqueeze(0)
split = img_path.split('/')
split[-1] = split[-1].split('.')[0] + '.txt'
split[-2] = 'pose'
ref_pose = data_util.load_pose('/'.join(split))
sample_poses = data_util.gen_pose_circle(ref_pose,
n_poses=num_renders,
centre=[0., 0., 0.])
actions = np.zeros([num_renders, 12])
for i, target_pose in enumerate(sample_poses):
actions[i] = (np.linalg.inv(target_pose) @ ref_pose).flatten()[:12]
actions = torch.from_numpy(actions).float().to(device)
with torch.no_grad():
state = model.encoder(img)
state = state.repeat(num_renders, 1, 1)
state = model.transition_model(state, actions)
out = model.decoder(state.reshape(-1, model.embedding_dim))
masks, _, recs = model.compose_image(out)
print(masks.shape)
plt.imshow(masks[0, 0].cpu().detach().numpy())
plt.show()
plt.imshow(masks[0, 1].cpu().detach().numpy())
plt.show()
plt.imshow(masks[0, 2].cpu().detach().numpy())
plt.show()
for i, rec in enumerate(recs):
torchvision.utils.save_image(rec,
os.path.join(
circle_dir, f'{i+1:04d}.png'),
normalize=True,
range=(-1, 1))
print('Saved one circle view.')
def __getitem__(self, idx):
img1 = data_util.load_rgb(self.rgb_paths[idx * 2],
sidelength=self.sidelength)
img2 = data_util.load_rgb(self.rgb_paths[idx * 2 + 1],
sidelength=self.sidelength)
img1 = img1.transpose(2, 0, 1)
img2 = img2.transpose(2, 0, 1)
pose1 = data_util.load_pose(self.pose_paths[idx * 2])
pose2 = data_util.load_pose(self.pose_paths[idx * 2 + 1])
transf21 = np.linalg.inv(pose2) @ pose1
transf12 = np.linalg.inv(transf21)
sample = {
'image1': img1,
'image2': img2,
'transf21': transf21.flatten()[:12],
'transf12': transf12.flatten()[:12],
'pose1': pose1.flatten()[:12],
'pose2': pose2.flatten()[:12]
}
return sample
def __init__(self, root_dir, img_size, sampling_pattern):
super().__init__()
self.img_size = img_size
self.color_dir = os.path.join(root_dir, 'rgb')
self.pose_dir = os.path.join(root_dir, 'pose')
if not os.path.isdir(self.color_dir):
print("Error! root dir is wrong")
return
self.all_color = sorted(data_util.glob_imgs(self.color_dir))
self.all_poses = sorted(glob(os.path.join(self.pose_dir, '*.txt')))
# Subsample the trajectory for training / test set split as well as the result matrix
file_lists = [self.all_color, self.all_poses]
if sampling_pattern != 'all':
if sampling_pattern.split('_')[0] == 'skip':
skip_val = int(sampling_pattern.split('_')[-1])
for i in range(len(file_lists)):
dummy_list = deepcopy(file_lists[i])
file_lists[i].clear()
file_lists[i].extend(dummy_list[::skip_val + 1])
else:
print("Unknown sampling pattern!")
return None
# Buffer files
print("Buffering files...")
self.all_views = []
self.all_views_nearest = []
for i in range(self.__len__()):
if not i % 10:
print(i)
self.all_views.append(self.read_view_tuple(i))
self.all_views_nearest.append(self.read_view_tuple_nearest(i))
# Calculate the ranking of nearest neigbors
self.nn_idcs, _ = data_util.get_nn_ranking(
[data_util.load_pose(pose) for pose in self.all_poses])
print("*" * 100)
print("Sampling pattern ", sampling_pattern)
print("Image size ", self.img_size)
print("*" * 100)
def __init__(self, root_dir):
super().__init__()
self.pose_dir = os.path.join(root_dir, 'pose')
self.all_poses = sorted(glob(os.path.join(self.pose_dir, '*.txt')))
# Buffer files
print("Buffering files...")
self.all_views = []
for i in range(self.__len__()):
if not i % 10:
print(i)
# based on deepvoxel implmentation, pose is camera_to_world, i.e. for y = pose*x, x is camera, y is world
self.all_views.append(data_util.load_pose(self.all_poses[i]))
self.gaze = self.get_gaze_vector()
self.pos = self.get_camera_pos()
