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

Compare commits

base: hofee:985a08d89cc65246c36da8b3256c16de7700bc76
Branches Tags
hofee:master
hofee:ab_new_partial_global
hofee:new_partial_global
hofee:ab_global_only
hofee:hemisphere_random
hofee:ab_global_and_partial_global
hofee:ab_global_and_local
hofee:ab_local_only
..
compare: hofee:8a05b7883d6f1568fa397e27e35d6ced87677dd6
Branches Tags
hofee:ab_new_partial_global
hofee:new_partial_global
hofee:ab_global_only
hofee:hemisphere_random
hofee:master
hofee:ab_global_and_partial_global
hofee:ab_global_and_local
hofee:ab_local_only

No commits in common. "985a08d89cc65246c36da8b3256c16de7700bc76" and "8a05b7883d6f1568fa397e27e35d6ced87677dd6" have entirely different histories.
985a08d89c ... 8a05b7883d

7 changed files with 219 additions and 404 deletions
Show Stats Expand all files Collapse all files

67
configs/local/inference_config.yaml
View File

@ -6,67 +6,71 @@ runner:
cuda_visible_devices: "0,1,2,3,4,5,6,7" cuda_visible_devices: "0,1,2,3,4,5,6,7"
experiment: experiment:
name: overfit_ab_global_only name: w_gf_wo_lf_full
root_dir: "experiments" root_dir: "experiments"
epoch: -1 # -1 stands for last epoch epoch: 1 # -1 stands for last epoch
test: test:
dataset_list: dataset_list:
- OmniObject3d_train - OmniObject3d_train
blender_script_path: "/data/hofee/project/nbv_rec/blender/data_renderer.py" blender_script_path: "/media/hofee/data/project/python/nbv_reconstruction/blender/data_renderer.py"
output_dir: "/data/hofee/data/inference_global_full_on_testset" output_dir: "/media/hofee/data/project/python/nbv_reconstruction/nbv_reconstruction/test/inference_global_full_on_testset"
pipeline: nbv_reconstruction_pipeline pipeline: nbv_reconstruction_global_pts_pipeline
voxel_size: 0.003
dataset: dataset:
OmniObject3d_train: OmniObject3d_train:
root_dir: "/data/hofee/data/new_full_data" root_dir: "/media/hofee/repository/nbv_reconstruction_data_512"
model_dir: "/data/hofee/data/scaled_object_meshes" model_dir: "/media/hofee/data/data/scaled_object_meshes"
source: seq_reconstruction_dataset source: seq_nbv_reconstruction_dataset
split_file: "/data/hofee/data/sample.txt" split_file: "/media/hofee/data/project/python/nbv_reconstruction/nbv_reconstruction/test/test_set_list.txt"
type: test type: test
filter_degree: 75 filter_degree: 75
ratio: 1 ratio: 1
batch_size: 1 batch_size: 1
num_workers: 12 num_workers: 12
pts_num: 8192 pts_num: 4096
load_from_preprocess: True load_from_preprocess: False
OmniObject3d_test:
root_dir: "/data/hofee/data/new_full_data"
model_dir: "/data/hofee/data/scaled_object_meshes"
source: seq_reconstruction_dataset
split_file: "/data/hofee/data/sample.txt"
type: test
filter_degree: 75
eval_list:
- pose_diff
- coverage_rate_increase
ratio: 0.1
batch_size: 1
num_workers: 12
pts_num: 8192
load_from_preprocess: True
pipeline: pipeline:
nbv_reconstruction_pipeline: nbv_reconstruction_local_pts_pipeline:
modules: modules:
pts_encoder: pointnet_encoder pts_encoder: pointnet_encoder
seq_encoder: transformer_seq_encoder seq_encoder: transformer_seq_encoder
pose_encoder: pose_encoder pose_encoder: pose_encoder
view_finder: gf_view_finder view_finder: gf_view_finder
eps: 1e-5 eps: 1e-5
global_scanned_feat: False
nbv_reconstruction_global_pts_pipeline:
modules:
pts_encoder: pointnet_encoder
pose_seq_encoder: transformer_pose_seq_encoder
pose_encoder: pose_encoder
view_finder: gf_view_finder
eps: 1e-5
global_scanned_feat: True global_scanned_feat: True
module: module:
pointnet_encoder: pointnet_encoder:
in_dim: 3 in_dim: 3
out_dim: 1024 out_dim: 1024
global_feat: True global_feat: True
feature_transform: False feature_transform: False
transformer_seq_encoder: transformer_seq_encoder:
embed_dim: 256 pts_embed_dim: 1024
pose_embed_dim: 256
num_heads: 4
ffn_dim: 256
num_layers: 3
output_dim: 2048
transformer_pose_seq_encoder:
pose_embed_dim: 256
num_heads: 4 num_heads: 4
ffn_dim: 256 ffn_dim: 256
num_layers: 3 num_layers: 3
@ -82,8 +86,7 @@ module:
sample_mode: ode sample_mode: ode
sampling_steps: 500 sampling_steps: 500
sde_mode: ve sde_mode: ve
pose_encoder: pose_encoder:
pose_dim: 9 pose_dim: 9
out_dim: 256 out_dim: 256
pts_num_encoder:
out_dim: 64

6
configs/server/server_train_config.yaml
View File

@ -13,7 +13,7 @@ runner:
epoch: -1 # -1 stands for last epoch epoch: -1 # -1 stands for last epoch
max_epochs: 5000 max_epochs: 5000
save_checkpoint_interval: 1 save_checkpoint_interval: 1
test_first: True test_first: False
train: train:
optimizer: optimizer:
@ -54,7 +54,7 @@ dataset:
filter_degree: 75 filter_degree: 75
eval_list: eval_list:
- pose_diff - pose_diff
ratio: 1 ratio: 0.1
batch_size: 80 batch_size: 80
num_workers: 12 num_workers: 12
pts_num: 8192 pts_num: 8192
@ -70,7 +70,7 @@ dataset:
filter_degree: 75 filter_degree: 75
eval_list: eval_list:
- pose_diff - pose_diff
ratio: 0.1 ratio: 0.01
batch_size: 80 batch_size: 80
num_workers: 12 num_workers: 12
pts_num: 8192 pts_num: 8192

2
core/nbv_dataset.py
View File

@ -34,7 +34,7 @@ class NBVReconstructionDataset(BaseDataset):
#self.model_dir = config["model_dir"] #self.model_dir = config["model_dir"]
self.filter_degree = config["filter_degree"] self.filter_degree = config["filter_degree"]
if self.type == namespace.Mode.TRAIN: if self.type == namespace.Mode.TRAIN:
scale_ratio = 1 scale_ratio = 50
self.datalist = self.datalist*scale_ratio self.datalist = self.datalist*scale_ratio
if self.cache: if self.cache:
expr_root = ConfigManager.get("runner", "experiment", "root_dir") expr_root = ConfigManager.get("runner", "experiment", "root_dir")

154
core/old_seq_dataset.py
View File

@ -1,154 +0,0 @@
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 ------+

349
core/seq_dataset.py
View File

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

39
runners/inferencer.py
View File

@ -27,7 +27,6 @@ class Inferencer(Runner):
self.script_path = ConfigManager.get(namespace.Stereotype.RUNNER, "blender_script_path") self.script_path = ConfigManager.get(namespace.Stereotype.RUNNER, "blender_script_path")
self.output_dir = ConfigManager.get(namespace.Stereotype.RUNNER, "output_dir") self.output_dir = ConfigManager.get(namespace.Stereotype.RUNNER, "output_dir")
self.voxel_size = ConfigManager.get(namespace.Stereotype.RUNNER, "voxel_size")
''' Pipeline ''' ''' Pipeline '''
self.pipeline_name = self.config[namespace.Stereotype.PIPELINE] self.pipeline_name = self.config[namespace.Stereotype.PIPELINE]
self.pipeline:torch.nn.Module = ComponentFactory.create(namespace.Stereotype.PIPELINE, self.pipeline_name) self.pipeline:torch.nn.Module = ComponentFactory.create(namespace.Stereotype.PIPELINE, self.pipeline_name)
@ -66,11 +65,16 @@ class Inferencer(Runner):
for dataset_idx, test_set in enumerate(self.test_set_list): for dataset_idx, test_set in enumerate(self.test_set_list):
status_manager.set_progress("inference", "inferencer", f"dataset", dataset_idx, len(self.test_set_list)) status_manager.set_progress("inference", "inferencer", f"dataset", dataset_idx, len(self.test_set_list))
test_set_name = test_set.get_name() test_set_name = test_set.get_name()
test_loader = test_set.get_loader()
total=int(len(test_set)) if test_loader.batch_size > 1:
for i in range(total): Log.error("Batch size should be 1 for inference, found {} in {}".format(test_loader.batch_size, test_set_name), terminate=True)
data = test_set.__getitem__(i)
total=int(len(test_loader))
loop = tqdm(enumerate(test_loader), total=total)
for i, data in loop:
status_manager.set_progress("inference", "inferencer", f"Batch[{test_set_name}]", i+1, total) status_manager.set_progress("inference", "inferencer", f"Batch[{test_set_name}]", i+1, total)
test_set.process_batch(data, self.device)
output = self.predict_sequence(data) output = self.predict_sequence(data)
self.save_inference_result(test_set_name, data["scene_name"][0], output) self.save_inference_result(test_set_name, data["scene_name"][0], output)
@ -84,23 +88,26 @@ class Inferencer(Runner):
''' data for rendering ''' ''' data for rendering '''
scene_path = data["scene_path"][0] scene_path = data["scene_path"][0]
O_to_L_pose = data["O_to_L_pose"][0] O_to_L_pose = data["O_to_L_pose"][0]
voxel_threshold = self.voxel_size voxel_threshold = data["voxel_threshold"][0]
filter_degree = 75 filter_degree = data["filter_degree"][0]
down_sampled_model_pts = data["gt_pts"] model_points_normals = data["model_points_normals"][0]
import ipdb; ipdb.set_trace() model_pts = model_points_normals[:,:3]
first_frame_to_world_9d = data["first_scanned_n_to_world_pose_9d"][0] down_sampled_model_pts = PtsUtil.voxel_downsample_point_cloud(model_pts, voxel_threshold)
first_frame_to_world = np.eye(4) first_frame_to_world_9d = data["first_to_world_9d"][0]
first_frame_to_world[:3,:3] = PoseUtil.rotation_6d_to_matrix_numpy(first_frame_to_world_9d[:6]) first_frame_to_world = torch.eye(4, device=first_frame_to_world_9d.device)
first_frame_to_world[:3,3] = first_frame_to_world_9d[6:] first_frame_to_world[:3,:3] = PoseUtil.rotation_6d_to_matrix_tensor_batch(first_frame_to_world_9d[:,:6])[0]
first_frame_to_world[:3,3] = first_frame_to_world_9d[0,6:]
first_frame_to_world = first_frame_to_world.to(self.device)
''' data for inference ''' ''' data for inference '''
input_data = {} input_data = {}
input_data["combined_scanned_pts"] = torch.tensor(data["first_scanned_pts"][0], dtype=torch.float32).to(self.device) input_data["scanned_pts"] = [data["first_pts"][0].to(self.device)]
input_data["scanned_n_to_world_pose_9d"] = [torch.tensor(data["first_scanned_n_to_world_pose_9d"], dtype=torch.float32).to(self.device)] input_data["scanned_n_to_world_pose_9d"] = [data["first_to_world_9d"][0].to(self.device)]
input_data["mode"] = namespace.Mode.TEST input_data["mode"] = namespace.Mode.TEST
input_pts_N = input_data["combined_scanned_pts"].shape[1] input_data["combined_scanned_pts"] = data["combined_scanned_pts"]
input_pts_N = input_data["scanned_pts"][0].shape[1]
first_frame_target_pts, _ = RenderUtil.render_pts(first_frame_to_world, scene_path, self.script_path, down_sampled_model_pts, voxel_threshold=voxel_threshold, filter_degree=filter_degree, nO_to_nL_pose=O_to_L_pose) first_frame_target_pts, _ = RenderUtil.render_pts(first_frame_to_world, scene_path, self.script_path, model_points_normals, voxel_threshold=voxel_threshold, filter_degree=filter_degree, nO_to_nL_pose=O_to_L_pose)
scanned_view_pts = [first_frame_target_pts] scanned_view_pts = [first_frame_target_pts]
last_pred_cr = self.compute_coverage_rate(scanned_view_pts, None, down_sampled_model_pts, threshold=voxel_threshold) last_pred_cr = self.compute_coverage_rate(scanned_view_pts, None, down_sampled_model_pts, threshold=voxel_threshold)

6
utils/render.py
View File

@ -10,7 +10,7 @@ from utils.pts import PtsUtil
class RenderUtil: class RenderUtil:
@staticmethod @staticmethod
def render_pts(cam_pose, scene_path, script_path, voxel_threshold=0.005, filter_degree=75, nO_to_nL_pose=None, require_full_scene=False): def render_pts(cam_pose, scene_path, script_path, model_points_normals, voxel_threshold=0.005, filter_degree=75, nO_to_nL_pose=None, require_full_scene=False):
nO_to_world_pose = DataLoadUtil.get_real_cam_O_from_cam_L(cam_pose, nO_to_nL_pose, scene_path=scene_path) nO_to_world_pose = DataLoadUtil.get_real_cam_O_from_cam_L(cam_pose, nO_to_nL_pose, scene_path=scene_path)
@ -34,10 +34,10 @@ class RenderUtil:
return None return None
path = os.path.join(temp_dir, "tmp") path = os.path.join(temp_dir, "tmp")
point_cloud = DataLoadUtil.get_target_point_cloud_world_from_path(path, binocular=True) point_cloud = DataLoadUtil.get_target_point_cloud_world_from_path(path, binocular=True)
normals = DataLoadUtil.get_target_normals_world_from_path(path, binocular=True)
cam_params = DataLoadUtil.load_cam_info(path, binocular=True) cam_params = DataLoadUtil.load_cam_info(path, binocular=True)
filtered_point_cloud = PtsUtil.filter_points(point_cloud, normals, cam_pose=cam_params["cam_to_world"], voxel_size=voxel_threshold, theta=filter_degree) ''' TODO: old code: filter_points api is changed, need to update the code '''
filtered_point_cloud = PtsUtil.filter_points(point_cloud, model_points_normals, cam_pose=cam_params["cam_to_world"], voxel_size=voxel_threshold, theta=filter_degree)
full_scene_point_cloud = None full_scene_point_cloud = None
if require_full_scene: if require_full_scene:
depth_L, depth_R = DataLoadUtil.load_depth(path, cam_params['near_plane'], cam_params['far_plane'], binocular=True) depth_L, depth_R = DataLoadUtil.load_depth(path, cam_params['near_plane'], cam_params['far_plane'], binocular=True)