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

add global_feat

Browse Source
This commit is contained in:
hofee
2024-09-24 09:10:25 +00:00
parent b209ce050c
commit 43f22ad91b
7 changed files with 123 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

49
core/pipeline.py
View File

@@ -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