283 lines
12 KiB
Python
283 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
|
|
import time
|
|
|
|
sys.path.append(r"/data/hofee/project/nbv_rec/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.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:
|
|
scale_ratio = 1
|
|
self.datalist = self.datalist*scale_ratio
|
|
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:
|
|
seq_num = DataLoadUtil.get_label_num(self.root_dir, scene_name)
|
|
scene_max_coverage_rate = 0
|
|
max_coverage_rate_list = []
|
|
|
|
for seq_idx in range(seq_num):
|
|
label_path = DataLoadUtil.get_label_path(
|
|
self.root_dir, scene_name, seq_idx
|
|
)
|
|
label_data = DataLoadUtil.load_label(label_path)
|
|
max_coverage_rate = label_data["max_coverage_rate"]
|
|
if max_coverage_rate > scene_max_coverage_rate:
|
|
scene_max_coverage_rate = max_coverage_rate
|
|
max_coverage_rate_list.append(max_coverage_rate)
|
|
|
|
if max_coverage_rate_list:
|
|
mean_coverage_rate = np.mean(max_coverage_rate_list)
|
|
|
|
for seq_idx in range(seq_num):
|
|
label_path = DataLoadUtil.get_label_path(
|
|
self.root_dir, scene_name, seq_idx
|
|
)
|
|
label_data = DataLoadUtil.load_label(label_path)
|
|
if max_coverage_rate_list[seq_idx] > mean_coverage_rate - 0.1:
|
|
for data_pair in label_data["data_pairs"]:
|
|
scanned_views = data_pair[0]
|
|
next_best_view = data_pair[1]
|
|
datalist.append(
|
|
{
|
|
"scanned_views": scanned_views,
|
|
"next_best_view": next_best_view,
|
|
"seq_max_coverage_rate": max_coverage_rate,
|
|
"scene_name": scene_name,
|
|
"label_idx": seq_idx,
|
|
"scene_max_coverage_rate": scene_max_coverage_rate,
|
|
}
|
|
)
|
|
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}")
|
|
|
|
def voxel_downsample_with_mapping(self, point_cloud, voxel_size=0.003):
|
|
voxel_indices = np.floor(point_cloud / voxel_size).astype(np.int32)
|
|
unique_voxels, inverse, counts = np.unique(voxel_indices, axis=0, return_inverse=True, return_counts=True)
|
|
idx_sort = np.argsort(inverse)
|
|
idx_unique = idx_sort[np.cumsum(counts)-counts]
|
|
downsampled_points = point_cloud[idx_unique]
|
|
return downsampled_points, inverse
|
|
|
|
|
|
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["seq_max_coverage_rate"]
|
|
scene_name = data_item_info["scene_name"]
|
|
(
|
|
scanned_views_pts,
|
|
scanned_coverages_rate,
|
|
scanned_n_to_world_pose,
|
|
) = ([], [], [])
|
|
#start_time = time.time()
|
|
start_indices = [0]
|
|
total_points = 0
|
|
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"]
|
|
target_point_cloud = (
|
|
DataLoadUtil.load_from_preprocessed_pts(view_path)
|
|
)
|
|
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_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)
|
|
total_points += len(downsampled_target_point_cloud)
|
|
start_indices.append(total_points)
|
|
|
|
|
|
#end_time = time.time()
|
|
#Log.info(f"load data time: {end_time - start_time}")
|
|
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
|
|
)
|
|
|
|
combined_scanned_views_pts = np.concatenate(scanned_views_pts, axis=0)
|
|
voxel_downsampled_combined_scanned_pts_np, inverse = self.voxel_downsample_with_mapping(combined_scanned_views_pts, 0.003)
|
|
random_downsampled_combined_scanned_pts_np, random_downsample_idx = PtsUtil.random_downsample_point_cloud(voxel_downsampled_combined_scanned_pts_np, self.pts_num, require_idx=True)
|
|
|
|
# all_idx_unique = np.arange(len(voxel_downsampled_combined_scanned_pts_np))
|
|
# all_random_downsample_idx = all_idx_unique[random_downsample_idx]
|
|
# scanned_pts_mask = []
|
|
# for idx, start_idx in enumerate(start_indices):
|
|
# if idx == len(start_indices) - 1:
|
|
# break
|
|
# end_idx = start_indices[idx+1]
|
|
# view_inverse = inverse[start_idx:end_idx]
|
|
# view_unique_downsampled_idx = np.unique(view_inverse)
|
|
# view_unique_downsampled_idx_set = set(view_unique_downsampled_idx)
|
|
# mask = np.array([idx in view_unique_downsampled_idx_set for idx in all_random_downsample_idx])
|
|
# #scanned_pts_mask.append(mask)
|
|
data_item = {
|
|
"scanned_pts": np.asarray(scanned_views_pts, dtype=np.float32), # Ndarray(S x Nv x 3)
|
|
"combined_scanned_pts": np.asarray(random_downsampled_combined_scanned_pts_np, dtype=np.float32), # Ndarray(N x 3)
|
|
#"scanned_pts_mask": np.asarray(scanned_pts_mask, dtype=np.bool), # Ndarray(N)
|
|
"scanned_coverage_rate": scanned_coverages_rate, # List(S): Float, range(0, 1)
|
|
"scanned_n_to_world_pose_9d": np.asarray(scanned_n_to_world_pose, dtype=np.float32), # Ndarray(S x 9)
|
|
"best_coverage_rate": nbv_coverage_rate, # Float, range(0, 1)
|
|
"best_to_world_pose_9d": np.asarray(best_to_world_9d, dtype=np.float32), # Ndarray(9)
|
|
"seq_max_coverage_rate": max_coverage_rate, # Float, range(0, 1)
|
|
"scene_name": scene_name, # String
|
|
}
|
|
|
|
return data_item
|
|
|
|
def __len__(self):
|
|
return len(self.datalist)
|
|
|
|
def get_collate_fn(self):
|
|
def collate_fn(batch):
|
|
collate_data = {}
|
|
|
|
''' ------ Varialbe Length ------ '''
|
|
|
|
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["scanned_pts_mask"] = [
|
|
# torch.tensor(item["scanned_pts_mask"]) for item in batch
|
|
# ]
|
|
''' ------ Fixed Length ------ '''
|
|
|
|
collate_data["best_to_world_pose_9d"] = torch.stack(
|
|
[torch.tensor(item["best_to_world_pose_9d"]) for item in batch]
|
|
)
|
|
collate_data["combined_scanned_pts"] = torch.stack(
|
|
[torch.tensor(item["combined_scanned_pts"]) 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",
|
|
"combined_scanned_pts",
|
|
"scanned_pts_mask",
|
|
]:
|
|
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": "/data/hofee/nbv_rec_part2_preprocessed",
|
|
"source": "nbv_reconstruction_dataset",
|
|
"split_file": "/data/hofee/data/sample.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 ------
|