145 lines
5.8 KiB
Python
145 lines
5.8 KiB
Python
import torch
|
|
import time
|
|
from torch import nn
|
|
import PytorchBoot.namespace as namespace
|
|
import PytorchBoot.stereotype as stereotype
|
|
from PytorchBoot.factory.component_factory import ComponentFactory
|
|
from PytorchBoot.utils import Log
|
|
|
|
|
|
@stereotype.pipeline("nbv_reconstruction_global_pts_n_num_pipeline")
|
|
class NBVReconstructionGlobalPointsPipeline(nn.Module):
|
|
def __init__(self, config):
|
|
super(NBVReconstructionGlobalPointsPipeline, self).__init__()
|
|
self.config = config
|
|
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.pts_num_encoder = ComponentFactory.create(
|
|
namespace.Stereotype.MODULE, self.module_config["pts_num_encoder"]
|
|
)
|
|
|
|
self.transformer_seq_encoder = ComponentFactory.create(
|
|
namespace.Stereotype.MODULE, self.module_config["transformer_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"]
|
|
|
|
if mode == namespace.Mode.TRAIN:
|
|
return self.forward_train(data)
|
|
elif mode == namespace.Mode.TEST:
|
|
return self.forward_test(data)
|
|
else:
|
|
Log.error("Unknown mode: {}".format(mode), True)
|
|
|
|
def pertube_data(self, gt_delta_9d):
|
|
bs = gt_delta_9d.shape[0]
|
|
random_t = (
|
|
torch.rand(bs, device=gt_delta_9d.device) * (1.0 - self.eps) + self.eps
|
|
)
|
|
random_t = random_t.unsqueeze(-1)
|
|
mu, std = self.view_finder.marginal_prob(gt_delta_9d, random_t)
|
|
std = std.view(-1, 1)
|
|
z = torch.randn_like(gt_delta_9d)
|
|
perturbed_x = mu + z * std
|
|
target_score = -z * std / (std**2)
|
|
return perturbed_x, random_t, target_score, std
|
|
|
|
def forward_train(self, data):
|
|
start_time = time.time()
|
|
main_feat = self.get_main_feat(data)
|
|
end_time = time.time()
|
|
print("get_main_feat time: ", end_time - start_time)
|
|
""" 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,
|
|
"main_feat": main_feat,
|
|
}
|
|
estimated_score = self.view_finder(input_data)
|
|
output = {
|
|
"estimated_score": estimated_score,
|
|
"target_score": target_score,
|
|
"std": std,
|
|
}
|
|
return output
|
|
|
|
def forward_test(self, data):
|
|
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_main_feat(self, data):
|
|
scanned_n_to_world_pose_9d_batch = data[
|
|
"scanned_n_to_world_pose_9d"
|
|
] # List(B): Tensor(S x 9)
|
|
scanned_pts_mask_batch = data[
|
|
"scanned_pts_mask"
|
|
] # Tensor(B x N)
|
|
|
|
device = next(self.parameters()).device
|
|
|
|
embedding_list_batch = []
|
|
|
|
combined_scanned_pts_batch = data["combined_scanned_pts"] # Tensor(B x N x 3)
|
|
global_scanned_feat, perpoint_scanned_feat_batch = self.pts_encoder.encode_points(
|
|
combined_scanned_pts_batch, require_per_point_feat=True
|
|
) # global_scanned_feat: Tensor(B x Dg), perpoint_scanned_feat: Tensor(B x N x Dl)
|
|
|
|
for scanned_n_to_world_pose_9d, scanned_mask, perpoint_scanned_feat in zip(
|
|
scanned_n_to_world_pose_9d_batch,
|
|
scanned_pts_mask_batch,
|
|
perpoint_scanned_feat_batch,
|
|
):
|
|
scanned_target_pts_num = [] # List(S): Int
|
|
partial_feat_seq = []
|
|
|
|
seq_len = len(scanned_n_to_world_pose_9d)
|
|
for seq_idx in range(seq_len):
|
|
partial_idx_in_combined_pts = scanned_mask == seq_idx # Ndarray(V), N->V idx mask
|
|
partial_perpoint_feat = perpoint_scanned_feat[partial_idx_in_combined_pts] # Ndarray(V x Dl)
|
|
partial_feat = torch.mean(partial_perpoint_feat, dim=0) # Tensor(Dl)
|
|
partial_feat_seq.append(partial_feat)
|
|
scanned_target_pts_num.append(partial_perpoint_feat.shape[0])
|
|
|
|
|
|
scanned_target_pts_num = torch.tensor(scanned_target_pts_num, dtype=torch.float32).unsqueeze(-1).to(device) # Tensor(S x 1)
|
|
scanned_n_to_world_pose_9d = scanned_n_to_world_pose_9d.to(device) # Tensor(S x 9)
|
|
pose_feat_seq = self.pose_encoder.encode_pose(scanned_n_to_world_pose_9d) # Tensor(S x Dp)
|
|
pts_num_feat_seq = self.pts_num_encoder.encode_pts_num(scanned_target_pts_num) # Tensor(S x Dn)
|
|
partial_feat_seq = torch.stack(partial_feat_seq) # Tensor(S x Dl)
|
|
seq_embedding = torch.cat([pose_feat_seq, pts_num_feat_seq, partial_feat_seq], dim=-1) # Tensor(S x (Dp+Dn+Dl))
|
|
embedding_list_batch.append(seq_embedding) # List(B): Tensor(S x (Dp+Dn+Dl))
|
|
|
|
|
|
seq_feat = self.transformer_seq_encoder.encode_sequence(embedding_list_batch) # Tensor(B x Ds)
|
|
main_feat = torch.cat([seq_feat, global_scanned_feat], dim=-1) # Tensor(B x (Ds+Dg))
|
|
|
|
if torch.isnan(main_feat).any():
|
|
Log.error("nan in main_feat", True)
|
|
|
|
return main_feat
|