262 lines
12 KiB
Python
262 lines
12 KiB
Python
import numpy as np
|
|
from PytorchBoot.dataset import BaseDataset
|
|
import PytorchBoot.namespace as namespace
|
|
import PytorchBoot.stereotype as stereotype
|
|
from PytorchBoot.config import ConfigManager
|
|
from PytorchBoot.utils.log_util import Log
|
|
import torch
|
|
import os
|
|
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
|
|
from utils.reconstruction import ReconstructionUtil
|
|
|
|
|
|
@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.split_file_path = config["split_file"]
|
|
self.scene_name_list = self.load_scene_name_list()
|
|
self.datalist = self.get_datalist()
|
|
|
|
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"]
|
|
if self.type == namespace.Mode.TRAIN:
|
|
self.datalist = self.datalist*100
|
|
if self.cache:
|
|
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):
|
|
scene_name_list = []
|
|
with open(self.split_file_path, "r") as f:
|
|
for line in f:
|
|
scene_name = line.strip()
|
|
scene_name_list.append(scene_name)
|
|
return scene_name_list
|
|
|
|
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_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 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, 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)
|
|
return data
|
|
else:
|
|
return None
|
|
|
|
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)
|
|
except Exception as e:
|
|
Log.error(f"Save cache failed: {e}")
|
|
# ----- Debug Trace ----- #
|
|
import ipdb; ipdb.set_trace()
|
|
# ------------------------ #
|
|
|
|
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_world_pose = [], [], []
|
|
|
|
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)
|
|
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"]
|
|
|
|
if self.load_from_preprocess:
|
|
downsampled_target_point_cloud = DataLoadUtil.load_from_preprocessed_pts(view_path)
|
|
else:
|
|
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)
|
|
n_to_world_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(n_to_world_pose[:3,:3]))
|
|
n_to_world_trans = n_to_world_pose[:3,3]
|
|
n_to_world_9d = np.concatenate([n_to_world_6d, n_to_world_trans], axis=0)
|
|
scanned_n_to_world_pose.append(n_to_world_9d)
|
|
|
|
nbv_idx, nbv_coverage_rate = nbv[0], nbv[1]
|
|
nbv_path = DataLoadUtil.get_path(self.root_dir, scene_name, nbv_idx)
|
|
cam_info = DataLoadUtil.load_cam_info(nbv_path)
|
|
best_frame_to_world = cam_info["cam_to_world"]
|
|
|
|
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,
|
|
"scanned_n_to_world_pose_9d": np.asarray(scanned_n_to_world_pose,dtype=np.float32),
|
|
"best_coverage_rate": nbv_coverage_rate,
|
|
"best_to_world_pose_9d": np.asarray(best_to_world_9d,dtype=np.float32),
|
|
"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
|
|
# model_points_normals = DataLoadUtil.load_points_normals(self.root_dir, scene_name)
|
|
# pts_list = []
|
|
# for view in scanned_views:
|
|
# frame_idx = view[0]
|
|
# view_path = DataLoadUtil.get_path(self.root_dir, scene_name, frame_idx)
|
|
# point_cloud = DataLoadUtil.get_target_point_cloud_world_from_path(view_path, binocular=True)
|
|
# cam_params = DataLoadUtil.load_cam_info(view_path, binocular=True)
|
|
# sampled_point_cloud = ReconstructionUtil.filter_points(point_cloud, model_points_normals, cam_pose=cam_params["cam_to_world"], voxel_size=voxel_threshold, theta=self.filter_degree)
|
|
# pts_list.append(sampled_point_cloud)
|
|
# nL_to_world_pose = cam_params["cam_to_world"]
|
|
# nO_to_world_pose = cam_params["cam_to_world_O"]
|
|
# nO_to_nL_pose = np.dot(np.linalg.inv(nL_to_world_pose), nO_to_world_pose)
|
|
# data_item["scanned_target_pts_list"] = pts_list
|
|
# data_item["model_points_normals"] = model_points_normals
|
|
# data_item["voxel_threshold"] = voxel_threshold
|
|
# data_item["filter_degree"] = self.filter_degree
|
|
# data_item["scene_path"] = os.path.join(self.root_dir, scene_name)
|
|
# data_item["first_frame_to_world"] = np.asarray(first_frame_to_world, dtype=np.float32)
|
|
# data_item["nO_to_nL_pose"] = np.asarray(nO_to_nL_pose, dtype=np.float32)
|
|
return data_item
|
|
|
|
def __len__(self):
|
|
return len(self.datalist)
|
|
|
|
def get_collate_fn(self):
|
|
def collate_fn(batch):
|
|
collate_data = {}
|
|
collate_data["scanned_pts"] = [torch.tensor(item['scanned_pts']) for item in batch]
|
|
collate_data["scanned_n_to_world_pose_9d"] = [torch.tensor(item['scanned_n_to_world_pose_9d']) for item in batch]
|
|
collate_data["best_to_world_pose_9d"] = torch.stack([torch.tensor(item['best_to_world_pose_9d']) for item in batch])
|
|
if "first_frame_to_world" in batch[0]:
|
|
collate_data["first_frame_to_world"] = torch.stack([torch.tensor(item["first_frame_to_world"]) for item in batch])
|
|
for key in batch[0].keys():
|
|
if key not in ["scanned_pts", "scanned_n_to_world_pose_9d", "best_to_world_pose_9d", "first_frame_to_world"]:
|
|
collate_data[key] = [item[key] for item in batch]
|
|
return collate_data
|
|
return collate_fn
|
|
|
|
|
|
# -------------- Debug ---------------- #
|
|
if __name__ == "__main__":
|
|
import torch
|
|
seed = 0
|
|
torch.manual_seed(seed)
|
|
np.random.seed(seed)
|
|
config = {
|
|
"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": "/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,
|
|
"num_workers": 0,
|
|
"pts_num": 4096,
|
|
"type": namespace.Mode.TRAIN,
|
|
}
|
|
ds = NBVReconstructionDataset(config)
|
|
print(len(ds))
|
|
#ds.__getitem__(10)
|
|
dl = ds.get_loader(shuffle=True)
|
|
for idx, data in enumerate(dl):
|
|
data = ds.process_batch(data, "cuda:0")
|
|
print(data)
|
|
# ------ Debug Start ------
|
|
import ipdb;ipdb.set_trace()
|
|
# ------ Debug End ------
|
|
#
|
|
# 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
|
|
# #np.savetxt("best_pts.txt", best_pts)
|
|
# for key, value in data.items():
|
|
# if isinstance(value, torch.Tensor):
|
|
# print(key, ":" ,value.shape)
|
|
# else:
|
|
# print(key, ":" ,len(value))
|
|
# if key == "scanned_n_to_world_pose_9d":
|
|
# for val in value:
|
|
# print(val.shape)
|
|
# if key == "scanned_pts":
|
|
# print("scanned_pts")
|
|
# for val in value:
|
|
# print(val.shape)
|
|
# cnt = 0
|
|
# for v in val:
|
|
# import ipdb;ipdb.set_trace()
|
|
# np.savetxt(f"pts_{cnt}.txt", v)
|
|
# cnt+=1
|
|
|
|
|
|
# print() |