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()