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

update inferencer; add load_from_preprocessed_pts

Browse Source
This commit is contained in:
hofee
2024-09-20 19:00:08 +08:00
parent a621749cc9
commit 8517255245
7 changed files with 69 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

48
core/nbv_dataset.py
View File

@@ -7,7 +7,7 @@ from PytorchBoot.utils.log_util import Log
import torch
import os
import sys
sys.path.append(r"/home/data/hofee/project/nbv_rec/nbv_reconstruction")
sys.path.append(r"/media/hofee/data/project/python/nbv_reconstruction/nbv_reconstruction")
from utils.data_load import DataLoadUtil
from utils.pose import PoseUtil
@@ -28,6 +28,7 @@ class NBVReconstructionDataset(BaseDataset):
self.pts_num = config["pts_num"]
self.type = config["type"]
self.cache = config.get("cache")
self.load_from_preprocess = config.get("load_from_preprocess", False)
if self.type == namespace.Mode.TEST:
self.model_dir = config["model_dir"]
self.filter_degree = config["filter_degree"]
@@ -111,24 +112,28 @@ class NBVReconstructionDataset(BaseDataset):
cam_info = DataLoadUtil.load_cam_info(view_path, binocular=True)
n_to_world_pose = cam_info["cam_to_world"]
nR_to_world_pose = cam_info["cam_to_world_R"]
cached_data = None
if self.cache:
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)
point_cloud_L = DataLoadUtil.get_point_cloud(depth_L, cam_info['cam_intrinsic'], n_to_world_pose)['points_world']
point_cloud_R = DataLoadUtil.get_point_cloud(depth_R, cam_info['cam_intrinsic'], nR_to_world_pose)['points_world']
point_cloud_L = PtsUtil.random_downsample_point_cloud(point_cloud_L, 65536)
point_cloud_R = PtsUtil.random_downsample_point_cloud(point_cloud_R, 65536)
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, frame_idx, downsampled_target_point_cloud)
if self.load_from_preprocess:
downsampled_target_point_cloud = DataLoadUtil.load_from_preprocessed_pts(view_path)
else:
downsampled_target_point_cloud = cached_data
cached_data = None
if self.cache:
cached_data = self.load_from_cache(scene_name, frame_idx)
if cached_data is None:
print("load depth")
depth_L, depth_R = DataLoadUtil.load_depth(view_path, cam_info['near_plane'], cam_info['far_plane'], binocular=True)
point_cloud_L = DataLoadUtil.get_point_cloud(depth_L, cam_info['cam_intrinsic'], n_to_world_pose)['points_world']
point_cloud_R = DataLoadUtil.get_point_cloud(depth_R, cam_info['cam_intrinsic'], nR_to_world_pose)['points_world']
point_cloud_L = PtsUtil.random_downsample_point_cloud(point_cloud_L, 65536)
point_cloud_R = PtsUtil.random_downsample_point_cloud(point_cloud_R, 65536)
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, frame_idx, downsampled_target_point_cloud)
else:
downsampled_target_point_cloud = cached_data
scanned_views_pts.append(downsampled_target_point_cloud)
scanned_coverages_rate.append(coverage_rate)
@@ -205,10 +210,11 @@ if __name__ == "__main__":
torch.manual_seed(seed)
np.random.seed(seed)
config = {
"root_dir": "../data/sample_for_training/scenes",
"model_dir": "../data/scaled_object_meshes",
"root_dir": "/media/hofee/data/project/python/nbv_reconstruction/sample_for_training/preprocessed_scenes/",
"model_dir": "/media/hofee/data/data/scaled_object_meshes",
"source": "nbv_reconstruction_dataset",
"split_file": "../data/sample_for_training/OmniObject3d_train.txt",
"split_file": "/media/hofee/data/project/python/nbv_reconstruction/sample_for_training/OmniObject3d_train.txt",
"load_from_preprocess": True,
"ratio": 0.5,
"batch_size": 2,
"filter_degree": 75,

5
core/seq_dataset.py
View File

@@ -46,10 +46,12 @@ class SeqNBVReconstructionDataset(BaseDataset):
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
"max_coverage_rate": max_coverage_rate,
"best_seq_len": best_seq_len,
})
return datalist[5:]
@@ -98,6 +100,7 @@ class SeqNBVReconstructionDataset(BaseDataset):
"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"],
}
return data_item