add global_feat
This commit is contained in:
@@ -7,12 +7,11 @@ 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")
|
||||
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
|
||||
from utils.reconstruction import ReconstructionUtil
|
||||
|
||||
|
||||
@stereotype.dataset("nbv_reconstruction_dataset")
|
||||
@@ -35,7 +34,7 @@ class NBVReconstructionDataset(BaseDataset):
|
||||
self.model_dir = config["model_dir"]
|
||||
self.filter_degree = config["filter_degree"]
|
||||
if self.type == namespace.Mode.TRAIN:
|
||||
scale_ratio = 1
|
||||
scale_ratio = 10
|
||||
self.datalist = self.datalist*scale_ratio
|
||||
if self.cache:
|
||||
expr_root = ConfigManager.get("runner", "experiment", "root_dir")
|
||||
@@ -56,20 +55,34 @@ class NBVReconstructionDataset(BaseDataset):
|
||||
def get_datalist(self):
|
||||
datalist = []
|
||||
for scene_name in self.scene_name_list:
|
||||
label_path = DataLoadUtil.get_label_path_old(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]
|
||||
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"]
|
||||
datalist.append(
|
||||
{
|
||||
"scanned_views": scanned_views,
|
||||
"next_best_view": next_best_view,
|
||||
"max_coverage_rate": max_coverage_rate,
|
||||
"scene_name": scene_name,
|
||||
}
|
||||
)
|
||||
if max_coverage_rate > scene_max_coverage_rate:
|
||||
scene_max_coverage_rate = max_coverage_rate
|
||||
max_coverage_rate_list.append(max_coverage_rate)
|
||||
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):
|
||||
@@ -102,7 +115,7 @@ class NBVReconstructionDataset(BaseDataset):
|
||||
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"]
|
||||
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 = [], [], []
|
||||
|
||||
@@ -151,13 +164,18 @@ class NBVReconstructionDataset(BaseDataset):
|
||||
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 = PtsUtil.voxel_downsample_point_cloud(combined_scanned_views_pts, 0.002)
|
||||
random_downsampled_combined_scanned_pts_np = PtsUtil.random_downsample_point_cloud(voxel_downsampled_combined_scanned_pts_np, self.pts_num)
|
||||
data_item = {
|
||||
"scanned_pts": np.asarray(scanned_views_pts,dtype=np.float32),
|
||||
"combined_scanned_pts": np.asarray(random_downsampled_combined_scanned_pts_np,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,
|
||||
"seq_max_coverage_rate": max_coverage_rate,
|
||||
"scene_name": scene_name
|
||||
}
|
||||
|
||||
@@ -195,10 +213,11 @@ class NBVReconstructionDataset(BaseDataset):
|
||||
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])
|
||||
collate_data["combined_scanned_pts"] = torch.stack([torch.tensor(item['combined_scanned_pts']) 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"]:
|
||||
if key not in ["scanned_pts", "scanned_n_to_world_pose_9d", "best_to_world_pose_9d", "first_frame_to_world", "combined_scanned_pts"]:
|
||||
collate_data[key] = [item[key] for item in batch]
|
||||
return collate_data
|
||||
return collate_fn
|
||||
@@ -211,11 +230,11 @@ if __name__ == "__main__":
|
||||
torch.manual_seed(seed)
|
||||
np.random.seed(seed)
|
||||
config = {
|
||||
"root_dir": "/media/hofee/repository/nbv_reconstruction_data_512",
|
||||
"model_dir": "/media/hofee/data/data/scaled_object_meshes",
|
||||
"root_dir": "/home/data/hofee/project/nbv_rec/data/nbv_rec_data_512_preproc_npy",
|
||||
"model_dir": "/home/data/hofee/project/nbv_rec/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": False,
|
||||
"split_file": "/home/data/hofee/project/nbv_rec/data/OmniObject3d_test.txt",
|
||||
"load_from_preprocess": True,
|
||||
"ratio": 0.5,
|
||||
"batch_size": 2,
|
||||
"filter_degree": 75,
|
||||
|
||||
@@ -5,16 +5,20 @@ import PytorchBoot.stereotype as stereotype
|
||||
from PytorchBoot.factory.component_factory import ComponentFactory
|
||||
from PytorchBoot.utils import Log
|
||||
|
||||
from utils.pts import PtsUtil
|
||||
|
||||
@stereotype.pipeline("nbv_reconstruction_pipeline")
|
||||
class NBVReconstructionPipeline(nn.Module):
|
||||
def __init__(self, config):
|
||||
super(NBVReconstructionPipeline, self).__init__()
|
||||
self.config = config
|
||||
self.pts_encoder = ComponentFactory.create(namespace.Stereotype.MODULE, config["pts_encoder"])
|
||||
self.pose_encoder = ComponentFactory.create(namespace.Stereotype.MODULE, config["pose_encoder"])
|
||||
self.seq_encoder = ComponentFactory.create(namespace.Stereotype.MODULE, config["seq_encoder"])
|
||||
self.view_finder = ComponentFactory.create(namespace.Stereotype.MODULE, config["view_finder"])
|
||||
self.eps = 1e-5
|
||||
self.module_config = config["modules"]
|
||||
self.pts_encoder = ComponentFactory.create(namespace.Stereotype.MODULE, self.module_config["pts_encoder"])
|
||||
self.pose_encoder = ComponentFactory.create(namespace.Stereotype.MODULE, self.module_config["pose_encoder"])
|
||||
self.seq_encoder = ComponentFactory.create(namespace.Stereotype.MODULE, self.module_config["seq_encoder"])
|
||||
self.view_finder = ComponentFactory.create(namespace.Stereotype.MODULE, self.module_config["view_finder"])
|
||||
self.eps = float(self.config["eps"])
|
||||
self.enable_global_scanned_feat = self.config["global_scanned_feat"]
|
||||
|
||||
def forward(self, data):
|
||||
mode = data["mode"]
|
||||
@@ -38,14 +42,14 @@ class NBVReconstructionPipeline(nn.Module):
|
||||
return perturbed_x, random_t, target_score, std
|
||||
|
||||
def forward_train(self, data):
|
||||
seq_feat = self.get_seq_feat(data)
|
||||
main_feat = self.get_main_feat(data)
|
||||
''' get std '''
|
||||
best_to_world_pose_9d_batch = data["best_to_world_pose_9d"]
|
||||
perturbed_x, random_t, target_score, std = self.pertube_data(best_to_world_pose_9d_batch)
|
||||
input_data = {
|
||||
"sampled_pose": perturbed_x,
|
||||
"t": random_t,
|
||||
"seq_feat": seq_feat,
|
||||
"main_feat": main_feat,
|
||||
}
|
||||
estimated_score = self.view_finder(input_data)
|
||||
output = {
|
||||
@@ -56,29 +60,44 @@ class NBVReconstructionPipeline(nn.Module):
|
||||
return output
|
||||
|
||||
def forward_test(self,data):
|
||||
seq_feat = self.get_seq_feat(data)
|
||||
estimated_delta_rot_9d, in_process_sample = self.view_finder.next_best_view(seq_feat)
|
||||
main_feat = self.get_main_feat(data)
|
||||
estimated_delta_rot_9d, in_process_sample = self.view_finder.next_best_view(main_feat)
|
||||
result = {
|
||||
"pred_pose_9d": estimated_delta_rot_9d,
|
||||
"in_process_sample": in_process_sample
|
||||
}
|
||||
return result
|
||||
|
||||
def get_seq_feat(self, data):
|
||||
|
||||
def get_main_feat(self, data):
|
||||
scanned_pts_batch = data['scanned_pts']
|
||||
scanned_n_to_world_pose_9d_batch = data['scanned_n_to_world_pose_9d']
|
||||
|
||||
|
||||
device = next(self.parameters()).device
|
||||
|
||||
|
||||
|
||||
pts_feat_seq_list = []
|
||||
pose_feat_seq_list = []
|
||||
device = next(self.parameters()).device
|
||||
|
||||
for scanned_pts,scanned_n_to_world_pose_9d in zip(scanned_pts_batch,scanned_n_to_world_pose_9d_batch):
|
||||
|
||||
scanned_pts = scanned_pts.to(device)
|
||||
scanned_n_to_world_pose_9d = scanned_n_to_world_pose_9d.to(device)
|
||||
pts_feat_seq_list.append(self.pts_encoder.encode_points(scanned_pts))
|
||||
pose_feat_seq_list.append(self.pose_encoder.encode_pose(scanned_n_to_world_pose_9d))
|
||||
|
||||
main_feat = self.seq_encoder.encode_sequence(pts_feat_seq_list, pose_feat_seq_list)
|
||||
|
||||
seq_feat = self.seq_encoder.encode_sequence(pts_feat_seq_list, pose_feat_seq_list)
|
||||
if torch.isnan(seq_feat).any():
|
||||
Log.error("nan in seq_feat", True)
|
||||
return seq_feat
|
||||
if self.enable_global_scanned_feat:
|
||||
combined_scanned_pts_batch = data['combined_scanned_pts']
|
||||
global_scanned_feat = self.pts_encoder.encode_points(combined_scanned_pts_batch)
|
||||
main_feat = torch.cat([main_feat, global_scanned_feat], dim=-1)
|
||||
|
||||
|
||||
if torch.isnan(main_feat).any():
|
||||
Log.error("nan in main_feat", True)
|
||||
|
||||
return main_feat
|
||||
|
||||
|
||||
Reference in New Issue
Block a user