124 lines
5.7 KiB
Python
124 lines
5.7 KiB
Python
import os
|
|
import numpy as np
|
|
from PytorchBoot.dataset import BaseDataset
|
|
import PytorchBoot.stereotype as stereotype
|
|
|
|
import sys
|
|
sys.path.append(r"/media/hofee/data/project/python/nbv_reconstruction/nbv_reconstruction")
|
|
|
|
from utils.data_load import DataLoadUtil
|
|
from utils.pose import PoseUtil
|
|
from utils.pts import PtsUtil
|
|
|
|
@stereotype.dataset("nbv_reconstruction_dataset")
|
|
class NBVReconstructionDataset(BaseDataset):
|
|
def __init__(self, config):
|
|
super(NBVReconstructionDataset, self).__init__(config)
|
|
self.config = config
|
|
self.root_dir = config["root_dir"]
|
|
self.datalist = self.get_datalist()
|
|
self.pts_num = 1024
|
|
|
|
def get_datalist(self):
|
|
datalist = []
|
|
scene_name_list = os.listdir(self.root_dir)
|
|
for scene_name in scene_name_list:
|
|
label_path = DataLoadUtil.get_label_path(self.root_dir, scene_name)
|
|
label_data = DataLoadUtil.load_label(label_path)
|
|
for data_pair in label_data["data_pairs"]:
|
|
scanned_views = data_pair[0]
|
|
next_best_view = data_pair[1]
|
|
max_coverage_rate = label_data["max_coverage_rate"]
|
|
datalist.append(
|
|
{
|
|
"scanned_views": scanned_views,
|
|
"next_best_view": next_best_view,
|
|
"max_coverage_rate": max_coverage_rate,
|
|
"scene_name": scene_name,
|
|
}
|
|
)
|
|
return datalist
|
|
|
|
def __getitem__(self, index):
|
|
data_item_info = self.datalist[index]
|
|
scanned_views = data_item_info["scanned_views"]
|
|
nbv = data_item_info["next_best_view"]
|
|
max_coverage_rate = data_item_info["max_coverage_rate"]
|
|
scene_name = data_item_info["scene_name"]
|
|
scanned_views_pts, scanned_coverages_rate, scanned_n_to_1_pose = [], [], []
|
|
first_frame_idx = scanned_views[0][0]
|
|
first_frame_to_world = DataLoadUtil.load_cam_info(DataLoadUtil.get_path(self.root_dir, scene_name, first_frame_idx))["cam_to_world"]
|
|
for view in scanned_views:
|
|
frame_idx = view[0]
|
|
coverage_rate = view[1]
|
|
view_path = DataLoadUtil.get_path(self.root_dir, scene_name, frame_idx)
|
|
depth = DataLoadUtil.load_depth(view_path)
|
|
cam_info = DataLoadUtil.load_cam_info(view_path)
|
|
mask = DataLoadUtil.load_seg(view_path)
|
|
frame_curr_to_world = cam_info["cam_to_world"]
|
|
n_to_1_pose = np.dot(np.linalg.inv(first_frame_to_world), frame_curr_to_world)
|
|
target_point_cloud = DataLoadUtil.get_target_point_cloud(depth, cam_info["cam_intrinsic"], n_to_1_pose, mask)["points_world"]
|
|
downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(target_point_cloud, self.pts_num)
|
|
scanned_views_pts.append(downsampled_target_point_cloud)
|
|
scanned_coverages_rate.append(coverage_rate)
|
|
n_to_1_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(n_to_1_pose[:3,:3]))
|
|
n_to_1_trans = n_to_1_pose[:3,3]
|
|
n_to_1_9d = np.concatenate([n_to_1_6d, n_to_1_trans], axis=0)
|
|
scanned_n_to_1_pose.append(n_to_1_9d)
|
|
|
|
nbv_idx, nbv_coverage_rate = nbv[0], nbv[1]
|
|
nbv_path = DataLoadUtil.get_path(self.root_dir, scene_name, nbv_idx)
|
|
nbv_depth = DataLoadUtil.load_depth(nbv_path)
|
|
cam_info = DataLoadUtil.load_cam_info(nbv_path)
|
|
nbv_mask = DataLoadUtil.load_seg(nbv_path)
|
|
best_frame_to_world = cam_info["cam_to_world"]
|
|
best_to_1_pose = np.dot(np.linalg.inv(first_frame_to_world), best_frame_to_world)
|
|
best_target_point_cloud = DataLoadUtil.get_target_point_cloud(nbv_depth, cam_info["cam_intrinsic"], best_to_1_pose, nbv_mask)["points_world"]
|
|
downsampled_best_target_point_cloud = PtsUtil.random_downsample_point_cloud(best_target_point_cloud, self.pts_num)
|
|
best_to_1_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(best_to_1_pose[:3,:3]))
|
|
best_to_1_trans = best_to_1_pose[:3,3]
|
|
best_to_1_9d = np.concatenate([best_to_1_6d, best_to_1_trans], axis=0)
|
|
|
|
data_item = {
|
|
"scanned_pts": np.asarray(scanned_views_pts,dtype=np.float32),
|
|
"scanned_coverage_rate": np.asarray(scanned_coverages_rate,dtype=np.float32),
|
|
"scanned_n_to_1_pose_9d": np.asarray(scanned_n_to_1_pose,dtype=np.float32),
|
|
"best_pts": np.asarray(downsampled_best_target_point_cloud,dtype=np.float32),
|
|
"best_coverage_rate": nbv_coverage_rate,
|
|
"best_to_1_pose_9d": best_to_1_9d,
|
|
"max_coverage_rate": max_coverage_rate,
|
|
"scene_name": scene_name
|
|
}
|
|
|
|
return data_item
|
|
|
|
def __len__(self):
|
|
return len(self.datalist)
|
|
|
|
if __name__ == "__main__":
|
|
import torch
|
|
config = {
|
|
"root_dir": "/media/hofee/data/data/nbv_rec/sample",
|
|
"ratio": 0.05,
|
|
"batch_size": 1,
|
|
"num_workers": 0,
|
|
}
|
|
ds = NBVReconstructionDataset(config)
|
|
print(len(ds))
|
|
dl = ds.get_loader(shuffle=True)
|
|
for idx, data in enumerate(dl):
|
|
cnt=0
|
|
print(data["scene_name"])
|
|
print(data["scanned_coverage_rate"])
|
|
print(data["best_coverage_rate"])
|
|
for pts in data["scanned_pts"][0]:
|
|
#np.savetxt(f"pts_{cnt}.txt", pts)
|
|
cnt+=1
|
|
best_pts = data["best_pts"][0]
|
|
#np.savetxt("best_pts.txt", best_pts)
|
|
for key, value in data.items():
|
|
if isinstance(value, torch.Tensor):
|
|
print(key, ":" ,value.shape)
|
|
|
|
|
|
print() |