hofee/nbv_reconstruction
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

define display table as world space origin

Browse Source
This commit is contained in:
hofee
2024-09-19 11:29:43 +00:00
parent 8d5d6d5df4
commit bb75372f7e
4 changed files with 34 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
core/dataset.py
View File

@@ -17,6 +17,7 @@ from utils.reconstruction import ReconstructionUtil
@stereotype.dataset("nbv_reconstruction_dataset")
class NBVReconstructionDataset(BaseDataset):
DISPLAY_TABLE_POSITION = np.asarray([0,0,0.85])
def __init__(self, config):
super(NBVReconstructionDataset, self).__init__(config)
self.config = config
@@ -37,6 +38,8 @@ class NBVReconstructionDataset(BaseDataset):
expr_root = ConfigManager.get("runner", "experiment", "root_dir")
expr_name = ConfigManager.get("runner", "experiment", "name")
self.cache_dir = os.path.join(expr_root, expr_name, "cache")
#self.preprocess_cache()
def load_scene_name_list(self):
@@ -65,9 +68,15 @@ class NBVReconstructionDataset(BaseDataset):
}
)
return datalist
def preprocess_cache(self):
Log.info("preprocessing cache...")
for item_idx in range(len(self.datalist)):
self.__getitem__(item_idx)
Log.success("finish preprocessing cache.")
def load_from_cache(self, scene_name, first_frame_idx, curr_frame_idx):
cache_name = f"{scene_name}_{first_frame_idx}_{curr_frame_idx}.txt"
def load_from_cache(self, scene_name, curr_frame_idx):
cache_name = f"{scene_name}_{curr_frame_idx}.txt"
cache_path = os.path.join(self.cache_dir, cache_name)
if os.path.exists(cache_path):
data = np.loadtxt(cache_path)
@@ -75,8 +84,8 @@ class NBVReconstructionDataset(BaseDataset):
else:
return None
def save_to_cache(self, scene_name, first_frame_idx, curr_frame_idx, data):
cache_name = f"{scene_name}_{first_frame_idx}_{curr_frame_idx}.txt"
def save_to_cache(self, scene_name, curr_frame_idx, data):
cache_name = f"{scene_name}_{curr_frame_idx}.txt"
cache_path = os.path.join(self.cache_dir, cache_name)
try:
np.savetxt(cache_path, data)
@@ -106,7 +115,7 @@ class NBVReconstructionDataset(BaseDataset):
cached_data = None
if self.cache:
cached_data = self.load_from_cache(scene_name, first_frame_idx, frame_idx)
cached_data = self.load_from_cache(scene_name, frame_idx)
if cached_data is None:
depth_L, depth_R = DataLoadUtil.load_depth(view_path, cam_info['near_plane'], cam_info['far_plane'], binocular=True)
@@ -118,7 +127,7 @@ class NBVReconstructionDataset(BaseDataset):
overlap_points = DataLoadUtil.get_overlapping_points(point_cloud_L, point_cloud_R)
downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(overlap_points, self.pts_num)
if self.cache:
self.save_to_cache(scene_name, first_frame_idx, frame_idx, downsampled_target_point_cloud)
self.save_to_cache(scene_name, frame_idx, downsampled_target_point_cloud)
else:
downsampled_target_point_cloud = cached_data
@@ -137,7 +146,6 @@ class NBVReconstructionDataset(BaseDataset):
best_to_world_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(best_frame_to_world[:3,:3]))
best_to_world_trans = best_frame_to_world[:3,3]
best_to_world_9d = np.concatenate([best_to_world_6d, best_to_world_trans], axis=0)
data_item = {
"scanned_pts": np.asarray(scanned_views_pts,dtype=np.float32),
"scanned_coverage_rate": scanned_coverages_rate,
@@ -147,6 +155,8 @@ class NBVReconstructionDataset(BaseDataset):
"max_coverage_rate": max_coverage_rate,
"scene_name": scene_name
}
if self.type == namespace.Mode.TEST:
diag = DataLoadUtil.get_bbox_diag(self.model_dir, scene_name)
voxel_threshold = diag*0.02