first commit

core/ab_global_only_pts_pipeline.py

@@ -0,0 +1,85 @@
+import torch
+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_pipeline_global_only")
+class NBVReconstructionGlobalPointsOnlyPipeline(nn.Module):
+    def __init__(self, config):
+        super(NBVReconstructionGlobalPointsOnlyPipeline, 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.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. - 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):
+        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,
+            "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)
+        repeat_num = data.get("repeat_num", 50)
+        main_feat = main_feat.repeat(repeat_num, 1)
+        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):
+        
+        combined_scanned_pts_batch = data['combined_scanned_pts']   
+        global_scanned_feat = self.pts_encoder.encode_points(combined_scanned_pts_batch)
+        main_feat = global_scanned_feat
+        
+        
+        if torch.isnan(main_feat).any():
+            Log.error("nan in main_feat", True)
+        
+        return main_feat
+    

core/ab_local_only_pts_pipeline.py

@@ -0,0 +1,95 @@
+import torch
+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_pipeline_local_only")
+class NBVReconstructionLocalPointsOnlyPipeline(nn.Module):
+    def __init__(self, config):
+        super(NBVReconstructionLocalPointsOnlyPipeline, 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.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"]
+        
+        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. - 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):
+        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,
+            "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)
+        repeat_num = data.get("repeat_num", 50)
+        main_feat = main_feat.repeat(repeat_num, 1)
+        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_pts_batch = data['scanned_pts']
+        scanned_n_to_world_pose_9d_batch = data['scanned_n_to_world_pose_9d']
+        device = next(self.parameters()).device
+        feat_seq_list = []
+        
+        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 = self.pts_encoder.encode_points(scanned_pts)
+            pose_feat = self.pose_encoder.encode_pose(scanned_n_to_world_pose_9d)
+            seq_feat = torch.cat([pts_feat, pose_feat], dim=-1)
+            feat_seq_list.append(seq_feat)
+        main_feat = self.seq_encoder.encode_sequence(feat_seq_list)
+        
+        
+        if torch.isnan(main_feat).any():
+            Log.error("nan in main_feat", True)
+        
+        return main_feat
+    

core/ab_mlp_pipeline.py

@@ -0,0 +1,81 @@
+import torch
+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_pipeline_mlp")
+class NBVReconstructionMLPPipeline(nn.Module):
+    def __init__(self, config):
+        super(NBVReconstructionMLPPipeline, 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.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"]
+        
+        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 forward_train(self, data):
+        main_feat = self.get_main_feat(data)
+        ''' get std '''
+        best_to_world_pose_9d_batch = data["best_to_world_pose_9d"]
+        input_data = {
+            "main_feat": main_feat,
+        }
+        pred = self.view_finder(input_data)
+        output = {
+            "pred": pred,
+            "gt": best_to_world_pose_9d_batch,
+        }
+        return output
+    
+    def forward_test(self,data):
+        main_feat = self.get_main_feat(data)
+        estimated_delta_rot_9d, _ = self.view_finder.next_best_view(main_feat)
+        result = {
+            "pred_pose_9d": estimated_delta_rot_9d,
+        }
+        return result
+    
+    
+    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
+        feat_seq_list = []
+        
+        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 = self.pts_encoder.encode_points(scanned_pts)
+            pose_feat = self.pose_encoder.encode_pose(scanned_n_to_world_pose_9d)
+            seq_feat = torch.cat([pts_feat, pose_feat], dim=-1)
+            feat_seq_list.append(seq_feat)
+        main_feat = self.seq_encoder.encode_sequence(feat_seq_list)
+        
+        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
+    

core/evaluation.py

@@ -0,0 +1,109 @@
+import torch
+import numpy as np
+from utils.reconstruction import ReconstructionUtil
+from utils.pose import PoseUtil
+from utils.pts import PtsUtil
+from utils.render import RenderUtil
+import PytorchBoot.stereotype as stereotype
+import PytorchBoot.namespace as namespace
+from PytorchBoot.utils.log_util import Log
+
+
+
+@stereotype.evaluation_method("pose_diff")
+class PoseDiff:
+    def __init__(self, _):
+        pass
+    
+    def evaluate(self, output_list, data_list):
+        results = {namespace.TensorBoard.SCALAR: {}}
+        rot_angle_list = []
+        trans_dist_list = []
+        for output, data in zip(output_list, data_list):
+            gt_pose_9d = data['best_to_world_pose_9d']
+            pred_pose_9d = output['pred_pose_9d']
+            gt_rot_6d = gt_pose_9d[:, :6]
+            gt_trans = gt_pose_9d[:, 6:]
+            pred_rot_6d = pred_pose_9d[:, :6]
+            pred_trans = pred_pose_9d[:, 6:]
+            gt_rot_mat = PoseUtil.rotation_6d_to_matrix_tensor_batch(gt_rot_6d)
+            pred_rot_mat = PoseUtil.rotation_6d_to_matrix_tensor_batch(pred_rot_6d)
+            rotation_angles = PoseUtil.rotation_angle_distance(gt_rot_mat, pred_rot_mat)
+            
+            rot_angle_list.extend(list(rotation_angles))
+            trans_dist = torch.norm(gt_trans-pred_trans, dim=1).mean().item()
+            trans_dist_list.append(trans_dist)
+            
+
+        results[namespace.TensorBoard.SCALAR]["rot_diff"] = float(sum(rot_angle_list) / len(rot_angle_list))
+        results[namespace.TensorBoard.SCALAR]["trans_diff"] = float(sum(trans_dist_list) / len(trans_dist_list))
+        return results
+    
+
+
+@stereotype.evaluation_method("coverage_rate_increase")
+class ConverageRateIncrease:
+    def __init__(self, config):
+        self.config = config
+        self.renderer_path = config["renderer_path"]
+
+    def evaluate(self, output_list, data_list):
+        results = {namespace.TensorBoard.SCALAR: {}}
+        gt_coverate_increase_list = []
+        pred_coverate_increase_list = []
+        cr_diff_list = []
+        for output, data in zip(output_list, data_list):
+            scanned_cr = data['scanned_coverage_rate']
+            gt_cr = data["best_coverage_rate"]
+            scene_path_list = data['scene_path']
+            model_points_normals_list = data['model_points_normals']
+            scanned_view_pts_list = data['scanned_target_pts_list']
+            pred_pose_9ds = output['pred_pose_9d']
+            nO_to_nL_pose_batch = data["nO_to_nL_pose"]
+            voxel_threshold_list = data["voxel_threshold"]
+            filter_degree_list = data["filter_degree"]
+            first_frame_to_world = data["first_frame_to_world"]
+            pred_n_to_world_pose_mats = torch.eye(4, device=pred_pose_9ds.device).unsqueeze(0).repeat(pred_pose_9ds.shape[0], 1, 1)
+            pred_n_to_world_pose_mats[:,:3,:3] = PoseUtil.rotation_6d_to_matrix_tensor_batch(pred_pose_9ds[:, :6])
+            pred_n_to_world_pose_mats[:,:3,3] = pred_pose_9ds[:, 6:]
+            pred_n_to_world_pose_mats = torch.matmul(first_frame_to_world, pred_n_to_world_pose_mats)
+            for idx in range(len(scanned_cr)):
+                model_points_normals = model_points_normals_list[idx]
+                scanned_view_pts = scanned_view_pts_list[idx]
+                voxel_threshold = voxel_threshold_list[idx]
+                model_pts = model_points_normals[:,:3]
+                down_sampled_model_pts = PtsUtil.voxel_downsample_point_cloud(model_pts, voxel_threshold)
+                old_scanned_cr = self.compute_coverage_rate(scanned_view_pts, None, down_sampled_model_pts, threshold=voxel_threshold)
+                gt_coverate_increase_list.append(gt_cr[idx]-old_scanned_cr)
+
+                scene_path = scene_path_list[idx]
+                pred_pose = pred_n_to_world_pose_mats[idx]
+            
+                filter_degree = filter_degree_list[idx]
+                nO_to_nL_pose = nO_to_nL_pose_batch[idx]
+                try:
+                    new_pts, _ = RenderUtil.render_pts(pred_pose, scene_path, self.renderer_path, model_points_normals, voxel_threshold=voxel_threshold, filter_degree=filter_degree, nO_to_nL_pose=nO_to_nL_pose)
+                    pred_cr = self.compute_coverage_rate(scanned_view_pts, new_pts, down_sampled_model_pts, threshold=voxel_threshold)
+                except Exception as e:
+                    Log.warning(f"Error in scene {scene_path}, {e}")
+                    pred_cr = old_scanned_cr
+                pred_coverate_increase_list.append(pred_cr-old_scanned_cr)
+                cr_diff_list.append(gt_cr[idx]-pred_cr)
+                
+        results[namespace.TensorBoard.SCALAR]["gt_cr_increase"] = float(sum(gt_coverate_increase_list) / len(gt_coverate_increase_list))
+        results[namespace.TensorBoard.SCALAR]["pred_cr_increase"] = float(sum(pred_coverate_increase_list) / len(pred_coverate_increase_list))
+        results[namespace.TensorBoard.SCALAR]["cr_diff"] = float(sum(cr_diff_list) / len(cr_diff_list))
+        return results
+    
+    def compute_coverage_rate(self, scanned_view_pts, new_pts, model_pts, threshold=0.005):
+        if new_pts is not None:
+            new_scanned_view_pts = scanned_view_pts + [new_pts]
+        else:
+            new_scanned_view_pts = scanned_view_pts
+        combined_point_cloud = np.vstack(new_scanned_view_pts)
+        down_sampled_combined_point_cloud = PtsUtil.voxel_downsample_point_cloud(combined_point_cloud,threshold)
+        return ReconstructionUtil.compute_coverage_rate(model_pts, down_sampled_combined_point_cloud, threshold)
+
+    
+    
+    

core/global_pts_pipeline.py

@@ -0,0 +1,98 @@
+import torch
+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_pipeline_global")
+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.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"]
+        
+        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. - 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):
+        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,
+            "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)
+        repeat_num = data.get("repeat_num", 50)
+        main_feat = main_feat.repeat(repeat_num, 1)
+        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']
+
+        device = next(self.parameters()).device
+
+        feat_seq_list = []
+        
+        for scanned_n_to_world_pose_9d in scanned_n_to_world_pose_9d_batch:
+            scanned_n_to_world_pose_9d = scanned_n_to_world_pose_9d.to(device)
+            feat_seq_list.append(self.pose_encoder.encode_pose(scanned_n_to_world_pose_9d))
+
+        main_feat = self.seq_encoder.encode_sequence(feat_seq_list)
+        
+        
+        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
+    

core/local_pts_pipeline.py

@@ -0,0 +1,99 @@
+import torch
+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_pipeline_local")
+class NBVReconstructionLocalPointsPipeline(nn.Module):
+    def __init__(self, config):
+        super(NBVReconstructionLocalPointsPipeline, 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.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"]
+        
+        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. - 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):
+        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,
+            "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)
+        repeat_num = data.get("repeat_num", 50)
+        main_feat = main_feat.repeat(repeat_num, 1)
+        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_pts_batch = data['scanned_pts']
+        scanned_n_to_world_pose_9d_batch = data['scanned_n_to_world_pose_9d']
+        device = next(self.parameters()).device
+        feat_seq_list = []
+    
+        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 = self.pts_encoder.encode_points(scanned_pts)
+            pose_feat = self.pose_encoder.encode_pose(scanned_n_to_world_pose_9d)
+            seq_feat = torch.cat([pts_feat, pose_feat], dim=-1)
+            feat_seq_list.append(seq_feat)
+        main_feat = self.seq_encoder.encode_sequence(feat_seq_list)
+        
+        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
+    

core/loss.py

@@ -0,0 +1,27 @@
+import torch
+import PytorchBoot.stereotype as stereotype
+
+@stereotype.loss_function("gf_loss")
+class GFLoss:
+    def __init__(self, _):
+        pass
+
+    def compute(self, output, _):
+        estimated_score = output['estimated_score']
+        target_score = output['target_score']
+        std = output['std']
+        bs = estimated_score.shape[0]
+        loss_weighting = std ** 2
+        loss = torch.mean(torch.sum((loss_weighting * (estimated_score - target_score) ** 2).view(bs, -1), dim=-1))
+        return loss
+    
+@stereotype.loss_function("mse_loss")
+class MSELoss:
+    def __init__(self,_):
+        pass
+    
+    def compute(self, output, _):
+        pred_pose = output["pred"]
+        gt_pose = output["gt"]
+        loss = torch.mean(torch.sum((pred_pose - gt_pose) ** 2, dim=-1))
+        return loss

core/nbv_dataset.py

@@ -0,0 +1,282 @@
+import numpy as np
+from PytorchBoot.dataset import BaseDataset
+import PytorchBoot.namespace as namespace
+import PytorchBoot.stereotype as stereotype
+from PytorchBoot.config import ConfigManager
+from PytorchBoot.utils.log_util import Log
+import torch
+import os
+import sys
+import time
+
+sys.path.append(r"/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("nbv_reconstruction_dataset")
+class NBVReconstructionDataset(BaseDataset):
+    def __init__(self, config):
+        super(NBVReconstructionDataset, self).__init__(config)
+        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.type = config["type"]
+        self.cache = config.get("cache")
+        self.load_from_preprocess = config.get("load_from_preprocess", False)
+
+        if self.type == namespace.Mode.TEST:
+            #self.model_dir = config["model_dir"]
+            self.filter_degree = config["filter_degree"]
+        if self.type == namespace.Mode.TRAIN:
+            scale_ratio = 1
+            self.datalist = self.datalist*scale_ratio
+        if self.cache:
+            expr_root = ConfigManager.get("runner", "experiment", "root_dir")
+            expr_name = ConfigManager.get("runner", "experiment", "name")
+            self.cache_dir = os.path.join(expr_root, expr_name, "cache")
+            # self.preprocess_cache()
+
+    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
+            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"]
+                if max_coverage_rate > scene_max_coverage_rate:
+                    scene_max_coverage_rate = max_coverage_rate
+                max_coverage_rate_list.append(max_coverage_rate)
+                
+            if max_coverage_rate_list:
+                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):
+        Log.info("preprocessing cache...")
+        for item_idx in range(len(self.datalist)):
+            self.__getitem__(item_idx)
+        Log.success("finish preprocessing cache.")
+
+    def load_from_cache(self, scene_name, curr_frame_idx):
+        cache_name = f"{scene_name}_{curr_frame_idx}.txt"
+        cache_path = os.path.join(self.cache_dir, cache_name)
+        if os.path.exists(cache_path):
+            data = np.loadtxt(cache_path)
+            return data
+        else:
+            return None
+
+    def save_to_cache(self, scene_name, curr_frame_idx, data):
+        cache_name = f"{scene_name}_{curr_frame_idx}.txt"
+        cache_path = os.path.join(self.cache_dir, cache_name)
+        try:
+            np.savetxt(cache_path, data)
+        except Exception as e:
+            Log.error(f"Save cache failed: {e}")
+
+    def voxel_downsample_with_mapping(self, point_cloud, voxel_size=0.003):
+        voxel_indices = np.floor(point_cloud / voxel_size).astype(np.int32)
+        unique_voxels, inverse, counts = np.unique(voxel_indices, axis=0, return_inverse=True, return_counts=True)
+        idx_sort = np.argsort(inverse)
+        idx_unique = idx_sort[np.cumsum(counts)-counts]
+        downsampled_points = point_cloud[idx_unique]
+        return downsampled_points, inverse
+
+
+    def __getitem__(self, index):
+        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["seq_max_coverage_rate"]
+        scene_name = data_item_info["scene_name"]
+        (
+            scanned_views_pts,
+            scanned_coverages_rate,
+            scanned_n_to_world_pose,
+        ) = ([], [], [])
+        #start_time = time.time()
+        start_indices = [0]
+        total_points = 0
+        for view in scanned_views:
+            frame_idx = view[0]
+            coverage_rate = view[1]
+            view_path = DataLoadUtil.get_path(self.root_dir, scene_name, frame_idx)
+            cam_info = DataLoadUtil.load_cam_info(view_path, binocular=True)
+            
+            n_to_world_pose = cam_info["cam_to_world"]  
+            target_point_cloud = (
+                DataLoadUtil.load_from_preprocessed_pts(view_path)
+            )
+            downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(
+                target_point_cloud, self.pts_num
+            )
+            scanned_views_pts.append(downsampled_target_point_cloud)
+            scanned_coverages_rate.append(coverage_rate)
+            n_to_world_6d = PoseUtil.matrix_to_rotation_6d_numpy(
+                np.asarray(n_to_world_pose[:3, :3])
+            )
+            n_to_world_trans = n_to_world_pose[:3, 3]
+            n_to_world_9d = np.concatenate([n_to_world_6d, n_to_world_trans], axis=0)
+            scanned_n_to_world_pose.append(n_to_world_9d)
+            total_points += len(downsampled_target_point_cloud)
+            start_indices.append(total_points)
+
+
+        #end_time = time.time()
+        #Log.info(f"load data time: {end_time - start_time}")
+        nbv_idx, nbv_coverage_rate = nbv[0], nbv[1]
+        nbv_path = DataLoadUtil.get_path(self.root_dir, scene_name, nbv_idx)
+        cam_info = DataLoadUtil.load_cam_info(nbv_path)
+        best_frame_to_world = cam_info["cam_to_world"]
+
+        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, inverse = self.voxel_downsample_with_mapping(combined_scanned_views_pts, 0.003)
+        random_downsampled_combined_scanned_pts_np, random_downsample_idx = PtsUtil.random_downsample_point_cloud(voxel_downsampled_combined_scanned_pts_np, self.pts_num, require_idx=True)
+
+        # all_idx_unique = np.arange(len(voxel_downsampled_combined_scanned_pts_np))
+        # all_random_downsample_idx = all_idx_unique[random_downsample_idx]
+        # scanned_pts_mask = []
+        # for idx, start_idx in enumerate(start_indices):
+        #     if idx == len(start_indices) - 1:
+        #         break
+        #     end_idx = start_indices[idx+1]
+        #     view_inverse = inverse[start_idx:end_idx]
+        #     view_unique_downsampled_idx = np.unique(view_inverse)
+        #     view_unique_downsampled_idx_set = set(view_unique_downsampled_idx)
+        #     mask = np.array([idx in view_unique_downsampled_idx_set for idx in all_random_downsample_idx])
+        #     #scanned_pts_mask.append(mask)
+        data_item = {
+            "scanned_pts": np.asarray(scanned_views_pts, dtype=np.float32), # Ndarray(S x Nv x 3)
+            "combined_scanned_pts": np.asarray(random_downsampled_combined_scanned_pts_np, dtype=np.float32), # Ndarray(N x 3)
+            #"scanned_pts_mask": np.asarray(scanned_pts_mask, dtype=np.bool), # Ndarray(N)
+            "scanned_coverage_rate": scanned_coverages_rate, # List(S): Float, range(0, 1)
+            "scanned_n_to_world_pose_9d": np.asarray(scanned_n_to_world_pose, dtype=np.float32), # Ndarray(S x 9)
+            "best_coverage_rate": nbv_coverage_rate, # Float, range(0, 1)
+            "best_to_world_pose_9d": np.asarray(best_to_world_9d, dtype=np.float32), # Ndarray(9)
+            "seq_max_coverage_rate": max_coverage_rate, # Float, range(0, 1)
+            "scene_name": scene_name, # String
+        }
+
+        return data_item
+
+    def __len__(self):
+        return len(self.datalist)
+
+    def get_collate_fn(self):
+        def collate_fn(batch):
+            collate_data = {}
+            
+            ''' ------ Varialbe Length ------ '''
+            
+            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["scanned_pts_mask"] = [
+            #     torch.tensor(item["scanned_pts_mask"]) for item in batch
+            # ]
+            ''' ------ Fixed Length ------ '''
+            
+            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]
+            )
+            
+            for key in batch[0].keys():
+                if key not in [
+                    "scanned_pts",
+                    "scanned_n_to_world_pose_9d",
+                    "best_to_world_pose_9d",
+                    "combined_scanned_pts",
+                    "scanned_pts_mask",
+                ]:
+                    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": "/data/hofee/nbv_rec_part2_preprocessed",
+        "source": "nbv_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 = NBVReconstructionDataset(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 ------

core/old_seq_dataset.py

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

core/pipeline.py

@@ -0,0 +1,140 @@
+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_pipeline")
+class NBVReconstructionPipeline(nn.Module):
+    def __init__(self, config):
+        super(NBVReconstructionPipeline, 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.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"])
+
+    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):
+        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,
+            "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)
+        repeat_num = data.get("repeat_num", 1)
+        main_feat = main_feat.repeat(repeat_num, 1)
+        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"] # List(B): Tensor(S x N)
+
+        scanned_pts_mask_batch = data["scanned_pts_mask"] # List(B): Tensor(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, per_point_feat_batch = self.pts_encoder.encode_points(
+            combined_scanned_pts_batch, require_per_point_feat=True
+        )  # global_scanned_feat: Tensor(B x Dg)
+        batch_size = len(scanned_n_to_world_pose_9d_batch)
+        for i in range(batch_size):
+            seq_len = len(scanned_n_to_world_pose_9d_batch[i])
+            scanned_n_to_world_pose_9d = scanned_n_to_world_pose_9d_batch[i].to(device)  # Tensor(S x 9)
+            scanned_pts_mask = scanned_pts_mask_batch[i] # Tensor(S x N)
+            per_point_feat = per_point_feat_batch[i] # Tensor(N x Dp)
+            partial_point_feat_seq = []
+            for j in range(seq_len):
+                partial_per_point_feat = per_point_feat[scanned_pts_mask[j]]
+                if partial_per_point_feat.shape[0] == 0:
+                    partial_point_feat = torch.zeros(per_point_feat.shape[1], device=device)
+                else:
+                    partial_point_feat = torch.mean(partial_per_point_feat, dim=0) # Tensor(Dp)
+                partial_point_feat_seq.append(partial_point_feat)
+            partial_point_feat_seq = torch.stack(partial_point_feat_seq, dim=0) # Tensor(S x Dp)
+            
+            pose_feat_seq = self.pose_encoder.encode_pose(scanned_n_to_world_pose_9d)  # Tensor(S x Dp) 
+
+            seq_embedding = torch.cat([partial_point_feat_seq, pose_feat_seq], dim=-1)
+            
+            embedding_list_batch.append(seq_embedding) # List(B): Tensor(S x (Dp))
+        
+        seq_feat = self.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():
+            for i in range(len(main_feat)):
+                if torch.isnan(main_feat[i]).any():
+                    scanned_pts_mask = scanned_pts_mask_batch[i]
+                    Log.info(f"scanned_pts_mask shape: {scanned_pts_mask.shape}")
+                    Log.info(f"scanned_pts_mask sum: {scanned_pts_mask.sum()}")
+                    import ipdb
+                    ipdb.set_trace()
+            Log.error("nan in main_feat", True)
+
+        return main_feat

core/seq_dataset.py

@@ -0,0 +1,209 @@
+import numpy as np
+from PytorchBoot.dataset import BaseDataset
+import PytorchBoot.namespace as namespace
+import PytorchBoot.stereotype as stereotype
+from PytorchBoot.config import ConfigManager
+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")
+
+from utils.data_load import DataLoadUtil
+from utils.pose import PoseUtil
+from utils.pts import PtsUtil
+
+
+@stereotype.dataset("seq_reconstruction_dataset")
+class SeqReconstructionDataset(BaseDataset):
+    def __init__(self, config):
+        super(SeqReconstructionDataset, self).__init__(config)
+        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.type = config["type"]
+        self.cache = config.get("cache")
+        self.load_from_preprocess = config.get("load_from_preprocess", False)
+
+        if self.type == namespace.Mode.TEST:
+            #self.model_dir = config["model_dir"]
+            self.filter_degree = config["filter_degree"]
+        if self.type == namespace.Mode.TRAIN:
+            scale_ratio = 1
+            self.datalist = self.datalist*scale_ratio
+        if self.cache:
+            expr_root = ConfigManager.get("runner", "experiment", "root_dir")
+            expr_name = ConfigManager.get("runner", "experiment", "name")
+            self.cache_dir = os.path.join(expr_root, expr_name, "cache")
+            # self.preprocess_cache()
+
+    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()
+                if not os.path.exists(os.path.join(self.root_dir, scene_name)):
+                    continue
+                scene_name_list.append(scene_name)
+        return scene_name_list    
+    
+    def get_scene_name_list(self):
+        return self.scene_name_list
+
+
+    def get_datalist(self):
+        datalist = []
+        total = len(self.scene_name_list)
+        for idx, scene_name in enumerate(self.scene_name_list):
+            print(f"processing {scene_name} ({idx}/{total})")
+            scene_max_cr_idx = 0
+            frame_len = DataLoadUtil.get_scene_seq_length(self.root_dir, scene_name)
+
+            for i in range(10,frame_len):
+                path = DataLoadUtil.get_path(self.root_dir, scene_name, i)
+                pts = DataLoadUtil.load_from_preprocessed_pts(path, "npy")
+                print(pts.shape)
+                if pts.shape[0] == 0:
+                    continue
+                else:
+                    break
+            print(i)
+            datalist.append({
+                "scene_name": scene_name,
+                "first_frame": i,
+                "best_seq_len": -1,
+                "max_coverage_rate": 1.0,
+                "label_idx": scene_max_cr_idx,
+            })  
+        return datalist
+
+    def preprocess_cache(self):
+        Log.info("preprocessing cache...")
+        for item_idx in range(len(self.datalist)):
+            self.__getitem__(item_idx)
+        Log.success("finish preprocessing cache.")
+
+    def load_from_cache(self, scene_name, curr_frame_idx):
+        cache_name = f"{scene_name}_{curr_frame_idx}.txt"
+        cache_path = os.path.join(self.cache_dir, cache_name)
+        if os.path.exists(cache_path):
+            data = np.loadtxt(cache_path)
+            return data
+        else:
+            return None
+
+    def save_to_cache(self, scene_name, curr_frame_idx, data):
+        cache_name = f"{scene_name}_{curr_frame_idx}.txt"
+        cache_path = os.path.join(self.cache_dir, cache_name)
+        try:
+            np.savetxt(cache_path, data)
+        except Exception as e:
+            Log.error(f"Save cache failed: {e}")
+
+    def seq_combined_pts(self, scene, frame_idx_list):
+        all_combined_pts = []   
+        for i in frame_idx_list:
+            path = DataLoadUtil.get_path(self.root_dir, scene, i)
+            pts = DataLoadUtil.load_from_preprocessed_pts(path,"npy")
+            if pts.shape[0] == 0:
+                continue
+            all_combined_pts.append(pts)
+        all_combined_pts = np.vstack(all_combined_pts)
+        downsampled_all_pts = PtsUtil.voxel_downsample_point_cloud(all_combined_pts, 0.003)
+        return downsampled_all_pts
+
+    def __getitem__(self, index):
+        data_item_info = self.datalist[index]
+        max_coverage_rate = data_item_info["max_coverage_rate"]
+        best_seq_len = data_item_info["best_seq_len"]
+        scene_name = data_item_info["scene_name"]
+        (
+            scanned_views_pts,
+            scanned_coverages_rate,
+            scanned_n_to_world_pose,
+        ) = ([], [], [])
+        view = data_item_info["first_frame"]
+        frame_idx = view
+        view_path = DataLoadUtil.get_path(self.root_dir, scene_name, frame_idx)
+        cam_info = DataLoadUtil.load_cam_info(view_path, binocular=True)
+        
+        n_to_world_pose = cam_info["cam_to_world"]  
+        target_point_cloud = (
+            DataLoadUtil.load_from_preprocessed_pts(view_path)
+        )
+        downsampled_target_point_cloud = PtsUtil.random_downsample_point_cloud(
+            target_point_cloud, self.pts_num
+        )
+        scanned_views_pts.append(downsampled_target_point_cloud)
+        
+        n_to_world_6d = PoseUtil.matrix_to_rotation_6d_numpy(
+            np.asarray(n_to_world_pose[:3, :3])
+        ) 
+        first_left_cam_pose = cam_info["cam_to_world"]
+        first_center_cam_pose = cam_info["cam_to_world_O"]
+        first_O_to_first_L_pose = np.dot(np.linalg.inv(first_left_cam_pose), first_center_cam_pose)
+        n_to_world_trans = n_to_world_pose[:3, 3]
+        n_to_world_9d = np.concatenate([n_to_world_6d, n_to_world_trans], axis=0)
+        scanned_n_to_world_pose.append(n_to_world_9d)
+
+        frame_list = []
+        for i in range(DataLoadUtil.get_scene_seq_length(self.root_dir, scene_name)):
+            frame_list.append(i)
+        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_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)
+            "best_seq_len": best_seq_len, # Int
+            "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
+    from tqdm import tqdm
+    import pickle
+    import os
+
+    seed = 0
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    
+    config = {
+        "root_dir": "/media/hofee/data/data/test_bottle/view",
+        "source": "seq_reconstruction_dataset",
+        "split_file": "/media/hofee/data/data/test_bottle/test_bottle.txt",
+        "load_from_preprocess": True,
+        "filter_degree": 75,
+        "num_workers": 0,
+        "pts_num": 8192,
+        "type": namespace.Mode.TEST,
+    }
+    
+    output_dir = "/media/hofee/data/data/test_bottle/preprocessed_dataset"
+    os.makedirs(output_dir, exist_ok=True)
+
+    ds = SeqReconstructionDataset(config)
+    for i in tqdm(range(len(ds)), desc="processing dataset"):
+        output_path = os.path.join(output_dir, f"item_{i}.pkl")
+        item = ds.__getitem__(i)
+        for key, value in item.items():
+            if isinstance(value, np.ndarray):
+                item[key] = value.tolist()
+        #import ipdb; ipdb.set_trace()
+        with open(output_path, "wb") as f:
+            pickle.dump(item, f)

core/seq_dataset_preprocessed.py

@@ -0,0 +1,82 @@
+import numpy as np
+from PytorchBoot.dataset import BaseDataset
+import PytorchBoot.namespace as namespace
+import PytorchBoot.stereotype as stereotype
+from PytorchBoot.config import ConfigManager
+from PytorchBoot.utils.log_util import Log
+import pickle
+import torch
+import os
+import sys
+
+sys.path.append(r"C:\Document\Local Project\nbv_rec\nbv_reconstruction")
+
+from utils.data_load import DataLoadUtil
+from utils.pose import PoseUtil
+from utils.pts import PtsUtil
+
+@stereotype.dataset("seq_reconstruction_dataset_preprocessed")
+class SeqReconstructionDatasetPreprocessed(BaseDataset):
+    def __init__(self, config):
+        super(SeqReconstructionDatasetPreprocessed, self).__init__(config)
+        self.config = config
+        self.root_dir = config["root_dir"]
+        self.real_root_dir = r"/media/hofee/repository/final_test_set/view"
+        self.item_list = os.listdir(self.root_dir)
+
+    def __getitem__(self, index):
+        data = pickle.load(open(os.path.join(self.root_dir, self.item_list[index]), "rb"))
+        data_item = {
+            "first_scanned_pts": np.asarray(data["first_scanned_pts"], dtype=np.float32), # Ndarray(S x Nv x 3)
+            "first_scanned_n_to_world_pose_9d": np.asarray(data["first_scanned_n_to_world_pose_9d"], dtype=np.float32), # Ndarray(S x 9)
+            "seq_max_coverage_rate": data["seq_max_coverage_rate"], # Float, range(0, 1)
+            "best_seq_len": data["best_seq_len"], # Int
+            "scene_name": data["scene_name"], # String
+            "gt_pts": np.asarray(data["gt_pts"], dtype=np.float32), # Ndarray(N x 3)
+            "scene_path": os.path.join(self.real_root_dir, data["scene_name"]), # String
+            "O_to_L_pose": np.asarray(data["O_to_L_pose"], dtype=np.float32),
+        }
+        return data_item
+
+    def __len__(self):
+        return len(self.item_list)
+
+# -------------- Debug ---------------- #
+if __name__ == "__main__":
+    import torch
+
+    seed = 0
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+    '''
+    OmniObject3d_test:
+    root_dir: "H:\\AI\\Datasets\\packed_test_data"
+    model_dir: "H:\\AI\\Datasets\\scaled_object_meshes"
+    source: seq_reconstruction_dataset
+    split_file: "H:\\AI\\Datasets\\data_list\\OmniObject3d_test.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
+    '''
+    config = {
+        "root_dir": "/media/hofee/data/data/test_bottle/preprocessed_dataset",
+        "source": "seq_reconstruction_dataset",
+        "split_file": "H:\\AI\\Datasets\\data_list\\OmniObject3d_test.txt",
+        "load_from_preprocess": True,
+        "ratio": 1,
+        "filter_degree": 75,
+        "num_workers": 0,
+        "pts_num": 8192,
+        "type": "test",
+    }
+    ds = SeqReconstructionDataset(config)
+    print(len(ds))
+    print(ds.__getitem__(10))
+