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

change DataLoadUtil and Dataset to blender version

Browse Source
This commit is contained in:
hofee
2024-08-29 13:54:13 -05:00
parent eceedd5c15
commit 80cd4aba9d
6 changed files with 135 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
core/dataset.py
View File

@@ -1,10 +1,15 @@
import os
import numpy as np
from PytorchBoot.dataset import BaseDataset
import PytorchBoot.stereotype as stereotype
import sys
sys.path.append(r"C:\Document\Local Project\nbv_rec\nbv_reconstruction")
from utils.data_load import DataLoadUtil
from utils.pose import PoseUtil
@stereotype.dataset("nbv_reconstruction_dataset")
@stereotype.dataset("nbv_reconstruction_dataset", comment="to be modified")
class NBVReconstructionDataset(BaseDataset):
def __init__(self, config):
super(NBVReconstructionDataset, self).__init__(config)
@@ -15,9 +20,9 @@ class NBVReconstructionDataset(BaseDataset):
def get_datalist(self):
datalist = []
scene_idx_list = DataLoadUtil.get_scene_idx_list(self.root_dir)
for scene_idx in scene_idx_list:
label_path = DataLoadUtil.get_label_path(self.label_dir, scene_idx)
scene_name_list = os.listdir(self.root_dir)
for scene_name in scene_name_list:
label_path = DataLoadUtil.get_label_path(self.label_dir, scene_name)
label_data = DataLoadUtil.load_label(label_path)
for data_pair in label_data["data_pairs"]:
scanned_views = data_pair[0]
@@ -28,7 +33,7 @@ class NBVReconstructionDataset(BaseDataset):
"scanned_views": scanned_views,
"next_best_view": next_best_view,
"max_coverage_rate": max_coverage_rate,
"scene_idx": scene_idx,
"scene_name": scene_name,
}
)
return datalist
@@ -38,32 +43,39 @@ class NBVReconstructionDataset(BaseDataset):
scanned_views = data_item_info["scanned_views"]
nbv = data_item_info["next_best_view"]
max_coverage_rate = data_item_info["max_coverage_rate"]
scene_idx = data_item_info["scene_idx"]
scene_name = data_item_info["scene_name"]
scanned_views_pts, scanned_coverages_rate, scanned_cam_pose = [], [], []
for view in scanned_views:
frame_idx = view[0]
coverage_rate = view[1]
view_path = DataLoadUtil.get_path(self.root_dir, scene_idx, frame_idx)
view_path = DataLoadUtil.get_path(self.root_dir, scene_name, frame_idx)
pts = DataLoadUtil.load_depth(view_path)
scanned_views_pts.append(pts)
scanned_coverages_rate.append(coverage_rate)
cam_pose = DataLoadUtil.load_cam_info(view_path)["cam_to_world"]
scanned_cam_pose.append(cam_pose)
cam_pose_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(cam_pose[:3,:3]))
translation = cam_pose[:3,3]
cam_pose_9d = np.concatenate([cam_pose_6d, translation], axis=0)
scanned_cam_pose.append(cam_pose_9d)
nbv_idx, nbv_coverage_rate = nbv[0], nbv[1]
nbv_path = DataLoadUtil.get_path(self.root_dir, scene_idx, nbv_idx)
nbv_path = DataLoadUtil.get_path(self.root_dir, scene_name, nbv_idx)
nbv_pts = DataLoadUtil.load_depth(nbv_path)
cam_info = DataLoadUtil.load_cam_info(nbv_path)
nbv_cam_pose = cam_info["cam_to_world"]
nbv_cam_pose_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(nbv_cam_pose[:3,:3]))
translation = nbv_cam_pose[:3,3]
nbv_cam_pose_9d = np.concatenate([nbv_cam_pose_6d, translation], axis=0)
data_item = {
"scanned_views_pts": np.asarray(scanned_views_pts,dtype=np.float32),
"scanned_coverages_rate": np.asarray(scanned_coverages_rate,dtype=np.float32),
"scanned_cam_pose": np.asarray(scanned_cam_pose,dtype=np.float32),
"nbv_pts": np.asarray(nbv_pts,dtype=np.float32),
"nbv_coverage_rate": nbv_coverage_rate,
"nbv_cam_pose": nbv_cam_pose,
"nbv_cam_pose": nbv_cam_pose_9d,
"max_coverage_rate": max_coverage_rate,
"scene_name": scene_name
}
return data_item