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

global: upd

Browse Source
This commit is contained in:
hofee 2024-10-31 16:02:26 +00:00
parent 097712c0ea
commit b221036e8b
3 changed files with 332 additions and 155 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
core/nbv_dataset.py
View File

@ -48,7 +48,7 @@ class NBVReconstructionDataset(BaseDataset):
for line in f: for line in f:
scene_name = line.strip() scene_name = line.strip()
scene_name_list.append(scene_name) scene_name_list.append(scene_name)
return scene_name_list[:10] return scene_name_list
def get_datalist(self): def get_datalist(self):
datalist = [] datalist = []

154
core/old_seq_dataset.py Normal file
View File

@ -0,0 +1,154 @@
import numpy as np
from PytorchBoot.dataset import BaseDataset
import PytorchBoot.namespace as namespace
import PytorchBoot.stereotype as stereotype
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")
from utils.data_load import DataLoadUtil
from utils.pose import PoseUtil
from utils.pts import PtsUtil
@stereotype.dataset("old_seq_nbv_reconstruction_dataset")
class SeqNBVReconstructionDataset(BaseDataset):
def __init__(self, config):
super(SeqNBVReconstructionDataset, self).__init__(config)
self.type = config["type"]
if self.type != namespace.Mode.TEST:
Log.error("Dataset <seq_nbv_reconstruction_dataset> Only support test mode", terminate=True)
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.model_dir = config["model_dir"]
self.filter_degree = config["filter_degree"]
self.load_from_preprocess = config.get("load_from_preprocess", False)
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
scene_max_cr_idx = 0
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
scene_max_cr_idx = seq_idx
label_path = DataLoadUtil.get_label_path(self.root_dir, scene_name, scene_max_cr_idx)
label_data = DataLoadUtil.load_label(label_path)
first_frame = label_data["best_sequence"][0]
best_seq_len = len(label_data["best_sequence"])
datalist.append({
"scene_name": scene_name,
"first_frame": first_frame,
"max_coverage_rate": scene_max_coverage_rate,
"best_seq_len": best_seq_len,
"label_idx": scene_max_cr_idx,
})
return datalist
def __getitem__(self, index):
data_item_info = self.datalist[index]
first_frame_idx = data_item_info["first_frame"][0]
first_frame_coverage = data_item_info["first_frame"][1]
max_coverage_rate = data_item_info["max_coverage_rate"]
scene_name = data_item_info["scene_name"]
first_cam_info = DataLoadUtil.load_cam_info(DataLoadUtil.get_path(self.root_dir, scene_name, first_frame_idx), binocular=True)
first_view_path = DataLoadUtil.get_path(self.root_dir, scene_name, first_frame_idx)
first_left_cam_pose = first_cam_info["cam_to_world"]
first_center_cam_pose = first_cam_info["cam_to_world_O"]
first_target_point_cloud = DataLoadUtil.load_from_preprocessed_pts(first_view_path)
first_pts_num = first_target_point_cloud.shape[0]
first_downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(first_target_point_cloud, self.pts_num)
first_to_world_rot_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(first_left_cam_pose[:3,:3]))
first_to_world_trans = first_left_cam_pose[:3,3]
first_to_world_9d = np.concatenate([first_to_world_rot_6d, first_to_world_trans], axis=0)
diag = DataLoadUtil.get_bbox_diag(self.model_dir, scene_name)
voxel_threshold = diag*0.02
first_O_to_first_L_pose = np.dot(np.linalg.inv(first_left_cam_pose), first_center_cam_pose)
scene_path = os.path.join(self.root_dir, scene_name)
model_points_normals = DataLoadUtil.load_points_normals(self.root_dir, scene_name)
data_item = {
"first_pts_num": np.asarray(
first_pts_num, dtype=np.int32
),
"first_pts": np.asarray([first_downsampled_target_point_cloud],dtype=np.float32),
"combined_scanned_pts": np.asarray(first_downsampled_target_point_cloud,dtype=np.float32),
"first_to_world_9d": np.asarray([first_to_world_9d],dtype=np.float32),
"scene_name": scene_name,
"max_coverage_rate": max_coverage_rate,
"voxel_threshold": voxel_threshold,
"filter_degree": self.filter_degree,
"O_to_L_pose": first_O_to_first_L_pose,
"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"],
"first_frame_id": first_frame_idx,
}
return data_item
def __len__(self):
return len(self.datalist)
def get_collate_fn(self):
def collate_fn(batch):
collate_data = {}
collate_data["first_pts"] = [torch.tensor(item['first_pts']) for item in batch]
collate_data["first_to_world_9d"] = [torch.tensor(item['first_to_world_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 ["first_pts", "first_to_world_9d", "combined_scanned_pts"]:
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": "/home/data/hofee/project/nbv_rec/data/nbv_rec_data_512_preproc_npy",
"split_file": "/home/data/hofee/project/nbv_rec/data/OmniObject3d_train.txt",
"model_dir": "/home/data/hofee/project/nbv_rec/data/scaled_object_meshes",
"ratio": 0.005,
"batch_size": 2,
"filter_degree": 75,
"num_workers": 0,
"pts_num": 32684,
"type": namespace.Mode.TEST,
"load_from_preprocess": True
}
ds = SeqNBVReconstructionDataset(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 ------+

331
core/seq_dataset.py
View File

@ -1,154 +1,177 @@
import numpy as np import numpy as np
from PytorchBoot.dataset import BaseDataset from PytorchBoot.dataset import BaseDataset
import PytorchBoot.namespace as namespace import PytorchBoot.namespace as namespace
import PytorchBoot.stereotype as stereotype import PytorchBoot.stereotype as stereotype
from PytorchBoot.utils.log_util import Log from PytorchBoot.config import ConfigManager
import torch from PytorchBoot.utils.log_util import Log
import os import torch
import sys import os
sys.path.append(r"/home/data/hofee/project/nbv_rec/nbv_reconstruction") import sys
from utils.data_load import DataLoadUtil sys.path.append(r"/data/hofee/project/nbv_rec/nbv_reconstruction")
from utils.pose import PoseUtil
from utils.pts import PtsUtil from utils.data_load import DataLoadUtil
from utils.pose import PoseUtil
@stereotype.dataset("seq_nbv_reconstruction_dataset") from utils.pts import PtsUtil
class SeqNBVReconstructionDataset(BaseDataset):
def __init__(self, config):
super(SeqNBVReconstructionDataset, self).__init__(config) @stereotype.dataset("seq_reconstruction_dataset")
self.type = config["type"] class SeqReconstructionDataset(BaseDataset):
if self.type != namespace.Mode.TEST: def __init__(self, config):
Log.error("Dataset <seq_nbv_reconstruction_dataset> Only support test mode", terminate=True) super(SeqReconstructionDataset, self).__init__(config)
self.config = config self.config = config
self.root_dir = config["root_dir"] self.root_dir = config["root_dir"]
self.split_file_path = config["split_file"] self.split_file_path = config["split_file"]
self.scene_name_list = self.load_scene_name_list() self.scene_name_list = self.load_scene_name_list()
self.datalist = self.get_datalist() self.datalist = self.get_datalist()
self.pts_num = config["pts_num"]
self.pts_num = config["pts_num"]
self.model_dir = config["model_dir"] self.type = config["type"]
self.filter_degree = config["filter_degree"] self.cache = config.get("cache")
self.load_from_preprocess = config.get("load_from_preprocess", False) self.load_from_preprocess = config.get("load_from_preprocess", False)
if self.type == namespace.Mode.TEST:
def load_scene_name_list(self): #self.model_dir = config["model_dir"]
scene_name_list = [] self.filter_degree = config["filter_degree"]
with open(self.split_file_path, "r") as f: if self.type == namespace.Mode.TRAIN:
for line in f: scale_ratio = 1
scene_name = line.strip() self.datalist = self.datalist*scale_ratio
scene_name_list.append(scene_name) if self.cache:
return scene_name_list expr_root = ConfigManager.get("runner", "experiment", "root_dir")
expr_name = ConfigManager.get("runner", "experiment", "name")
def get_datalist(self): self.cache_dir = os.path.join(expr_root, expr_name, "cache")
datalist = [] # self.preprocess_cache()
for scene_name in self.scene_name_list:
seq_num = DataLoadUtil.get_label_num(self.root_dir, scene_name) def load_scene_name_list(self):
scene_max_coverage_rate = 0 scene_name_list = []
scene_max_cr_idx = 0 with open(self.split_file_path, "r") as f:
for line in f:
for seq_idx in range(seq_num): scene_name = line.strip()
label_path = DataLoadUtil.get_label_path(self.root_dir, scene_name, seq_idx) scene_name_list.append(scene_name)
label_data = DataLoadUtil.load_label(label_path) return scene_name_list
max_coverage_rate = label_data["max_coverage_rate"]
if max_coverage_rate > scene_max_coverage_rate: def get_datalist(self):
scene_max_coverage_rate = max_coverage_rate datalist = []
scene_max_cr_idx = seq_idx for scene_name in self.scene_name_list:
seq_num = DataLoadUtil.get_label_num(self.root_dir, scene_name)
label_path = DataLoadUtil.get_label_path(self.root_dir, scene_name, scene_max_cr_idx) scene_max_coverage_rate = 0
label_data = DataLoadUtil.load_label(label_path) max_coverage_rate_list = []
first_frame = label_data["best_sequence"][0] scene_max_cr_idx = 0
best_seq_len = len(label_data["best_sequence"]) for seq_idx in range(seq_num):
datalist.append({ label_path = DataLoadUtil.get_label_path(
"scene_name": scene_name, self.root_dir, scene_name, seq_idx
"first_frame": first_frame, )
"max_coverage_rate": scene_max_coverage_rate, label_data = DataLoadUtil.load_label(label_path)
"best_seq_len": best_seq_len, max_coverage_rate = label_data["max_coverage_rate"]
"label_idx": scene_max_cr_idx, if max_coverage_rate > scene_max_coverage_rate:
}) scene_max_coverage_rate = max_coverage_rate
return datalist scene_max_cr_idx = seq_idx
max_coverage_rate_list.append(max_coverage_rate)
def __getitem__(self, index): best_label_path = DataLoadUtil.get_label_path(self.root_dir, scene_name, scene_max_cr_idx)
data_item_info = self.datalist[index] best_label_data = DataLoadUtil.load_label(best_label_path)
first_frame_idx = data_item_info["first_frame"][0] first_frame = best_label_data["best_sequence"][0]
first_frame_coverage = data_item_info["first_frame"][1] best_seq_len = len(best_label_data["best_sequence"])
max_coverage_rate = data_item_info["max_coverage_rate"] datalist.append({
scene_name = data_item_info["scene_name"] "scene_name": scene_name,
first_cam_info = DataLoadUtil.load_cam_info(DataLoadUtil.get_path(self.root_dir, scene_name, first_frame_idx), binocular=True) "first_frame": first_frame,
first_view_path = DataLoadUtil.get_path(self.root_dir, scene_name, first_frame_idx) "best_seq_len": best_seq_len,
first_left_cam_pose = first_cam_info["cam_to_world"] "max_coverage_rate": scene_max_coverage_rate,
first_center_cam_pose = first_cam_info["cam_to_world_O"] "label_idx": scene_max_cr_idx,
first_target_point_cloud = DataLoadUtil.load_from_preprocessed_pts(first_view_path) })
first_pts_num = first_target_point_cloud.shape[0] return datalist
first_downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(first_target_point_cloud, self.pts_num)
first_to_world_rot_6d = PoseUtil.matrix_to_rotation_6d_numpy(np.asarray(first_left_cam_pose[:3,:3])) def preprocess_cache(self):
first_to_world_trans = first_left_cam_pose[:3,3] Log.info("preprocessing cache...")
first_to_world_9d = np.concatenate([first_to_world_rot_6d, first_to_world_trans], axis=0) for item_idx in range(len(self.datalist)):
diag = DataLoadUtil.get_bbox_diag(self.model_dir, scene_name) self.__getitem__(item_idx)
voxel_threshold = diag*0.02 Log.success("finish preprocessing cache.")
first_O_to_first_L_pose = np.dot(np.linalg.inv(first_left_cam_pose), first_center_cam_pose)
scene_path = os.path.join(self.root_dir, scene_name) def load_from_cache(self, scene_name, curr_frame_idx):
model_points_normals = DataLoadUtil.load_points_normals(self.root_dir, scene_name) cache_name = f"{scene_name}_{curr_frame_idx}.txt"
cache_path = os.path.join(self.cache_dir, cache_name)
data_item = { if os.path.exists(cache_path):
"first_pts_num": np.asarray( data = np.loadtxt(cache_path)
first_pts_num, dtype=np.int32 return data
), else:
"first_pts": np.asarray([first_downsampled_target_point_cloud],dtype=np.float32), return None
"combined_scanned_pts": np.asarray(first_downsampled_target_point_cloud,dtype=np.float32),
"first_to_world_9d": np.asarray([first_to_world_9d],dtype=np.float32), def save_to_cache(self, scene_name, curr_frame_idx, data):
"scene_name": scene_name, cache_name = f"{scene_name}_{curr_frame_idx}.txt"
"max_coverage_rate": max_coverage_rate, cache_path = os.path.join(self.cache_dir, cache_name)
"voxel_threshold": voxel_threshold, try:
"filter_degree": self.filter_degree, np.savetxt(cache_path, data)
"O_to_L_pose": first_O_to_first_L_pose, except Exception as e:
"first_frame_coverage": first_frame_coverage, Log.error(f"Save cache failed: {e}")
"scene_path": scene_path,
"model_points_normals": model_points_normals, def __getitem__(self, index):
"best_seq_len": data_item_info["best_seq_len"], data_item_info = self.datalist[index]
"first_frame_id": first_frame_idx, max_coverage_rate = data_item_info["max_coverage_rate"]
} scene_name = data_item_info["scene_name"]
return data_item (
scanned_views_pts,
def __len__(self): scanned_coverages_rate,
return len(self.datalist) scanned_n_to_world_pose,
) = ([], [], [])
def get_collate_fn(self): view = data_item_info["first_frame"]
def collate_fn(batch): frame_idx = view[0]
collate_data = {} coverage_rate = view[1]
collate_data["first_pts"] = [torch.tensor(item['first_pts']) for item in batch] view_path = DataLoadUtil.get_path(self.root_dir, scene_name, frame_idx)
collate_data["first_to_world_9d"] = [torch.tensor(item['first_to_world_9d']) for item in batch] cam_info = DataLoadUtil.load_cam_info(view_path, binocular=True)
collate_data["combined_scanned_pts"] = torch.stack([torch.tensor(item['combined_scanned_pts']) for item in batch])
for key in batch[0].keys(): n_to_world_pose = cam_info["cam_to_world"]
if key not in ["first_pts", "first_to_world_9d", "combined_scanned_pts"]: target_point_cloud = (
collate_data[key] = [item[key] for item in batch] DataLoadUtil.load_from_preprocessed_pts(view_path)
return collate_data )
return collate_fn downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(
target_point_cloud, self.pts_num
# -------------- Debug ---------------- # )
if __name__ == "__main__": scanned_views_pts.append(downsampled_target_point_cloud)
import torch scanned_coverages_rate.append(coverage_rate)
seed = 0 n_to_world_6d = PoseUtil.matrix_to_rotation_6d_numpy(
torch.manual_seed(seed) np.asarray(n_to_world_pose[:3, :3])
np.random.seed(seed) )
config = { n_to_world_trans = n_to_world_pose[:3, 3]
"root_dir": "/home/data/hofee/project/nbv_rec/data/nbv_rec_data_512_preproc_npy", n_to_world_9d = np.concatenate([n_to_world_6d, n_to_world_trans], axis=0)
"split_file": "/home/data/hofee/project/nbv_rec/data/OmniObject3d_train.txt", scanned_n_to_world_pose.append(n_to_world_9d)
"model_dir": "/home/data/hofee/project/nbv_rec/data/scaled_object_meshes",
"ratio": 0.005, # combined_scanned_views_pts = np.concatenate(scanned_views_pts, axis=0)
"batch_size": 2, # voxel_downsampled_combined_scanned_pts_np = PtsUtil.voxel_downsample_point_cloud(combined_scanned_views_pts, 0.002)
"filter_degree": 75, # random_downsampled_combined_scanned_pts_np = PtsUtil.random_downsample_point_cloud(voxel_downsampled_combined_scanned_pts_np, self.pts_num)
"num_workers": 0,
"pts_num": 32684, data_item = {
"type": namespace.Mode.TEST, "first_scanned_pts": np.asarray(scanned_views_pts, dtype=np.float32), # Ndarray(S x Nv x 3)
"load_from_preprocess": True "first_scanned_coverage_rate": scanned_coverages_rate, # List(S): Float, range(0, 1)
} "first_scanned_n_to_world_pose_9d": np.asarray(scanned_n_to_world_pose, dtype=np.float32), # Ndarray(S x 9)
ds = SeqNBVReconstructionDataset(config) "seq_max_coverage_rate": max_coverage_rate, # Float, range(0, 1)
print(len(ds)) "scene_name": scene_name, # String
#ds.__getitem__(10) }
dl = ds.get_loader(shuffle=True)
for idx, data in enumerate(dl): return data_item
data = ds.process_batch(data, "cuda:0")
print(data) def __len__(self):
# ------ Debug Start ------ return len(self.datalist)
import ipdb;ipdb.set_trace()
# ------ Debug End ------+
# -------------- Debug ---------------- #
if __name__ == "__main__":
import torch
seed = 0
torch.manual_seed(seed)
np.random.seed(seed)
config = {
"root_dir": "/data/hofee/data/new_full_data",
"source": "seq_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 = SeqReconstructionDataset(config)
print(len(ds))
print(ds.__getitem__(10))