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

change to-1 to to_world

Browse Source
This commit is contained in:
hofee
2024-09-19 00:20:26 +08:00
parent 935069d68c
commit 8d5d6d5df4
3 changed files with 29 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
core/pipeline.py
View File

@@ -40,8 +40,8 @@ class NBVReconstructionPipeline(nn.Module):
def forward_train(self, data):
seq_feat = self.get_seq_feat(data)
''' get std '''
best_to_1_pose_9d_batch = data["best_to_1_pose_9d"]
perturbed_x, random_t, target_score, std = self.pertube_data(best_to_1_pose_9d_batch)
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,
@@ -66,16 +66,16 @@ class NBVReconstructionPipeline(nn.Module):
def get_seq_feat(self, data):
scanned_pts_batch = data['scanned_pts']
scanned_n_to_1_pose_9d_batch = data['scanned_n_to_1_pose_9d']
scanned_n_to_world_pose_9d_batch = data['scanned_n_to_world_pose_9d']
pts_feat_seq_list = []
pose_feat_seq_list = []
device = next(self.parameters()).device
for scanned_pts,scanned_n_to_1_pose_9d in zip(scanned_pts_batch,scanned_n_to_1_pose_9d_batch):
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_1_pose_9d = scanned_n_to_1_pose_9d.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_1_pose_9d))
pose_feat_seq_list.append(self.pose_encoder.encode_pose(scanned_n_to_world_pose_9d))
seq_feat = self.seq_encoder.encode_sequence(pts_feat_seq_list, pose_feat_seq_list)
if torch.isnan(seq_feat).any():