add multi seq training
This commit is contained in:
@@ -28,6 +28,7 @@ class SeqNBVReconstructionDataset(BaseDataset):
|
||||
|
||||
self.model_dir = config["model_dir"]
|
||||
self.filter_degree = config["filter_degree"]
|
||||
self.load_from_preprocess = config.get("load_from_preprocess", False)
|
||||
|
||||
|
||||
def load_scene_name_list(self):
|
||||
@@ -38,10 +39,30 @@ class SeqNBVReconstructionDataset(BaseDataset):
|
||||
scene_name_list.append(scene_name)
|
||||
return scene_name_list
|
||||
|
||||
def get_datalist_new(self):
|
||||
datalist = []
|
||||
for scene_name in self.scene_name_list:
|
||||
label_num = DataLoadUtil.get_label_num(self.root_dir, scene_name)
|
||||
for i in range(label_num):
|
||||
label_path = DataLoadUtil.get_label_path(self.root_dir, scene_name, i)
|
||||
label_data = DataLoadUtil.load_label(label_path)
|
||||
best_seq = label_data["best_sequence"]
|
||||
max_coverage_rate = label_data["max_coverage_rate"]
|
||||
first_frame = best_seq[0]
|
||||
best_seq_len = len(best_seq)
|
||||
datalist.append({
|
||||
"scene_name": scene_name,
|
||||
"first_frame": first_frame,
|
||||
"max_coverage_rate": max_coverage_rate,
|
||||
"best_seq_len": best_seq_len,
|
||||
"label_idx": i,
|
||||
})
|
||||
return datalist
|
||||
|
||||
def get_datalist(self):
|
||||
datalist = []
|
||||
for scene_name in self.scene_name_list:
|
||||
label_path = DataLoadUtil.get_label_path(self.root_dir, scene_name)
|
||||
label_path = DataLoadUtil.get_label_path_old(self.root_dir, scene_name)
|
||||
label_data = DataLoadUtil.load_label(label_path)
|
||||
best_seq = label_data["best_sequence"]
|
||||
max_coverage_rate = label_data["max_coverage_rate"]
|
||||
@@ -52,8 +73,9 @@ class SeqNBVReconstructionDataset(BaseDataset):
|
||||
"first_frame": first_frame,
|
||||
"max_coverage_rate": max_coverage_rate,
|
||||
"best_seq_len": best_seq_len,
|
||||
"best_seq": best_seq,
|
||||
})
|
||||
return datalist[5:]
|
||||
return datalist
|
||||
|
||||
def __getitem__(self, index):
|
||||
data_item_info = self.datalist[index]
|
||||
@@ -62,27 +84,27 @@ class SeqNBVReconstructionDataset(BaseDataset):
|
||||
max_coverage_rate = data_item_info["max_coverage_rate"]
|
||||
scene_name = data_item_info["scene_name"]
|
||||
first_cam_info = DataLoadUtil.load_cam_info(DataLoadUtil.get_path(self.root_dir, scene_name, first_frame_idx), binocular=True)
|
||||
first_frame_to_world = first_cam_info["cam_to_world"]
|
||||
|
||||
first_view_path = DataLoadUtil.get_path(self.root_dir, scene_name, first_frame_idx)
|
||||
first_left_cam_pose = first_cam_info["cam_to_world"]
|
||||
first_right_cam_pose = first_cam_info["cam_to_world_R"]
|
||||
first_center_cam_pose = first_cam_info["cam_to_world_O"]
|
||||
|
||||
first_depth_L, first_depth_R = DataLoadUtil.load_depth(first_view_path, first_cam_info['near_plane'], first_cam_info['far_plane'], binocular=True)
|
||||
first_L_to_L_pose = np.eye(4)
|
||||
first_R_to_L_pose = np.dot(np.linalg.inv(first_left_cam_pose), first_right_cam_pose)
|
||||
first_point_cloud_L = DataLoadUtil.get_point_cloud(first_depth_L, first_cam_info['cam_intrinsic'], first_L_to_L_pose)['points_world']
|
||||
first_point_cloud_R = DataLoadUtil.get_point_cloud(first_depth_R, first_cam_info['cam_intrinsic'], first_R_to_L_pose)['points_world']
|
||||
|
||||
first_point_cloud_L = PtsUtil.random_downsample_point_cloud(first_point_cloud_L, 65536)
|
||||
first_point_cloud_R = PtsUtil.random_downsample_point_cloud(first_point_cloud_R, 65536)
|
||||
first_overlap_points = DataLoadUtil.get_overlapping_points(first_point_cloud_L, first_point_cloud_R)
|
||||
first_downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(first_overlap_points, self.pts_num)
|
||||
first_to_first_pose = np.eye(4)
|
||||
first_to_first_rot_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(first_to_first_pose[:3,:3]))
|
||||
first_to_first_trans = first_to_first_pose[:3,3]
|
||||
first_to_first_9d = np.concatenate([first_to_first_rot_6d, first_to_first_trans], axis=0)
|
||||
if self.load_from_preprocess:
|
||||
first_downsampled_target_point_cloud = DataLoadUtil.load_from_preprocessed_pts(first_view_path)
|
||||
else:
|
||||
first_depth_L, first_depth_R = DataLoadUtil.load_depth(first_view_path, first_cam_info['near_plane'], first_cam_info['far_plane'], binocular=True)
|
||||
|
||||
first_point_cloud_L = DataLoadUtil.get_point_cloud(first_depth_L, first_cam_info['cam_intrinsic'], first_left_cam_pose)['points_world']
|
||||
first_point_cloud_R = DataLoadUtil.get_point_cloud(first_depth_R, first_cam_info['cam_intrinsic'], first_right_cam_pose)['points_world']
|
||||
|
||||
first_point_cloud_L = PtsUtil.random_downsample_point_cloud(first_point_cloud_L, 65536)
|
||||
first_point_cloud_R = PtsUtil.random_downsample_point_cloud(first_point_cloud_R, 65536)
|
||||
first_overlap_points = DataLoadUtil.get_overlapping_points(first_point_cloud_L, first_point_cloud_R)
|
||||
first_downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(first_overlap_points, self.pts_num)
|
||||
|
||||
first_to_world_rot_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(first_left_cam_pose[:3,:3]))
|
||||
first_to_world_trans = first_left_cam_pose[:3,3]
|
||||
first_to_world_9d = np.concatenate([first_to_world_rot_6d, first_to_world_trans], axis=0)
|
||||
diag = DataLoadUtil.get_bbox_diag(self.model_dir, scene_name)
|
||||
voxel_threshold = diag*0.02
|
||||
first_O_to_first_L_pose = np.dot(np.linalg.inv(first_left_cam_pose), first_center_cam_pose)
|
||||
@@ -90,17 +112,17 @@ class SeqNBVReconstructionDataset(BaseDataset):
|
||||
model_points_normals = DataLoadUtil.load_points_normals(self.root_dir, scene_name)
|
||||
data_item = {
|
||||
"first_pts": np.asarray([first_downsampled_target_point_cloud],dtype=np.float32),
|
||||
"first_to_first_9d": np.asarray([first_to_first_9d],dtype=np.float32),
|
||||
"first_to_world_9d": np.asarray([first_to_world_9d],dtype=np.float32),
|
||||
"scene_name": scene_name,
|
||||
"max_coverage_rate": max_coverage_rate,
|
||||
"voxel_threshold": voxel_threshold,
|
||||
"filter_degree": self.filter_degree,
|
||||
"first_frame_to_world": np.asarray(first_frame_to_world, dtype=np.float32),
|
||||
"O_to_L_pose": first_O_to_first_L_pose,
|
||||
"first_frame_coverage": first_frame_coverage,
|
||||
"scene_path": scene_path,
|
||||
"model_points_normals": model_points_normals,
|
||||
"best_seq_len": data_item_info["best_seq_len"],
|
||||
"first_frame_id": first_frame_idx,
|
||||
}
|
||||
return data_item
|
||||
|
||||
@@ -111,10 +133,9 @@ class SeqNBVReconstructionDataset(BaseDataset):
|
||||
def collate_fn(batch):
|
||||
collate_data = {}
|
||||
collate_data["first_pts"] = [torch.tensor(item['first_pts']) for item in batch]
|
||||
collate_data["first_to_first_9d"] = [torch.tensor(item['first_to_first_9d']) for item in batch]
|
||||
collate_data["first_frame_to_world"] = torch.stack([torch.tensor(item["first_frame_to_world"]) for item in batch])
|
||||
collate_data["first_to_world_9d"] = [torch.tensor(item['first_to_world_9d']) for item in batch]
|
||||
for key in batch[0].keys():
|
||||
if key not in ["first_pts", "first_to_first_9d", "first_frame_to_world"]:
|
||||
if key not in ["first_pts", "first_to_world_9d"]:
|
||||
collate_data[key] = [item[key] for item in batch]
|
||||
return collate_data
|
||||
return collate_fn
|
||||
|
||||
Reference in New Issue
Block a user