ablation study

configs/local/inference_config.yaml

@@ -70,7 +70,7 @@ module:
     global_feat: True
     feature_transform: False
   transformer_seq_encoder:
-    embed_dim: 256
+    embed_dim: 320
     num_heads: 4
     ffn_dim: 256
     num_layers: 3

configs/server/server_train_config.yaml

@@ -3,11 +3,11 @@ runner:
   general:
     seed: 0
     device: cuda
-    cuda_visible_devices: "0"
+    cuda_visible_devices: "2"
     parallel: False
     
   experiment:
-    name: train_ab_global_only_with_wp_p++_strong
+    name: newtrain_real_global_only
     root_dir: "experiments"
     use_checkpoint: False
     epoch: -1 # -1 stands for last epoch
@@ -28,18 +28,18 @@ runner:
       - OmniObject3d_test
       - OmniObject3d_val
 
-  pipeline: nbv_reconstruction_pipeline
+  pipeline: nbv_reconstruction_pipeline_global_only
   
 dataset:
   OmniObject3d_train:
     root_dir: "/data/hofee/data/new_full_data"
     model_dir: "../data/scaled_object_meshes"
     source: nbv_reconstruction_dataset
-    split_file: "/data/hofee/data/new_full_data_list/OmniObject3d_train.txt"
+    split_file: "/data/hofee/data/new_full_data_list/new_OmniObject3d_train.txt"
     type: train
     cache: True
     ratio: 1
-    batch_size: 64
+    batch_size: 24
     num_workers: 128
     pts_num: 8192
     load_from_preprocess: True
@@ -48,14 +48,14 @@ dataset:
     root_dir: "/data/hofee/data/new_full_data"
     model_dir: "../data/scaled_object_meshes"
     source: nbv_reconstruction_dataset
-    split_file: "/data/hofee/data/new_full_data_list/OmniObject3d_test.txt"
+    split_file: "/data/hofee/data/new_full_data_list/new_OmniObject3d_test.txt"
     type: test
     cache: True
     filter_degree: 75
     eval_list:
       - pose_diff
     ratio: 1
-    batch_size: 80
+    batch_size: 32
     num_workers: 12
     pts_num: 8192
     load_from_preprocess: True
@@ -64,21 +64,37 @@ dataset:
     root_dir: "/data/hofee/data/new_full_data"
     model_dir: "../data/scaled_object_meshes"
     source: nbv_reconstruction_dataset
-    split_file: "/data/hofee/data/new_full_data_list/OmniObject3d_train.txt"
+    split_file: "/data/hofee/data/new_full_data_list/new_OmniObject3d_train.txt"
     type: test
     cache: True
     filter_degree: 75
     eval_list:
       - pose_diff
     ratio: 0.1
-    batch_size: 80
+    batch_size: 32
     num_workers: 12
     pts_num: 8192
     load_from_preprocess: True
 
 
 pipeline:
-  nbv_reconstruction_pipeline:
+  nbv_reconstruction_pipeline_local:
+    modules:
+      pts_encoder: pointnet++_encoder
+      seq_encoder: transformer_seq_encoder
+      pose_encoder: pose_encoder
+      view_finder: gf_view_finder
+    eps: 1e-5
+    global_scanned_feat: True
+  nbv_reconstruction_pipeline_global:
+    modules:
+      pts_encoder: pointnet++_encoder
+      seq_encoder: transformer_seq_encoder
+      pose_encoder: pose_encoder
+      view_finder: gf_view_finder
+    eps: 1e-5
+    global_scanned_feat: True
+  nbv_reconstruction_pipeline_local_only:
     modules:
       pts_encoder: pointnet++_encoder
       seq_encoder: transformer_seq_encoder
@@ -98,10 +114,9 @@ module:
 
   pointnet++_encoder:
     in_dim: 3
-    params_name: strong
 
   transformer_seq_encoder:
-    embed_dim: 256
+    embed_dim: 1280
     num_heads: 4
     ffn_dim: 256
     num_layers: 3
@@ -110,7 +125,7 @@ module:
   gf_view_finder:
     t_feat_dim: 128
     pose_feat_dim: 256
-    main_feat_dim: 5120
+    main_feat_dim: 1024
     regression_head: Rx_Ry_and_T
     pose_mode: rot_matrix
     per_point_feature: False

core/ab_global_only_pts_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_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)
+        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,91 @@
+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)
+        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/global_pts_pipeline.py

@@ -6,7 +6,7 @@ from PytorchBoot.factory.component_factory import ComponentFactory
 from PytorchBoot.utils import Log
 
 
-@stereotype.pipeline("nbv_reconstruction_global_pts_pipeline")
+@stereotype.pipeline("nbv_reconstruction_pipeline_global")
 class NBVReconstructionGlobalPointsPipeline(nn.Module):
     def __init__(self, config):
         super(NBVReconstructionGlobalPointsPipeline, self).__init__()
@@ -14,7 +14,7 @@ class NBVReconstructionGlobalPointsPipeline(nn.Module):
         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.pose_seq_encoder = ComponentFactory.create(namespace.Stereotype.MODULE, self.module_config["pose_seq_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"]
@@ -73,13 +73,13 @@ class NBVReconstructionGlobalPointsPipeline(nn.Module):
 
         device = next(self.parameters()).device
 
-        pose_feat_seq_list = []
+        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)
-            pose_feat_seq_list.append(self.pose_encoder.encode_pose(scanned_n_to_world_pose_9d))
+            feat_seq_list.append(self.pose_encoder.encode_pose(scanned_n_to_world_pose_9d))
 
-        main_feat = self.pose_seq_encoder.encode_sequence(pose_feat_seq_list)
+        main_feat = self.seq_encoder.encode_sequence(feat_seq_list)
         
         
         combined_scanned_pts_batch = data['combined_scanned_pts']   

core/local_pts_pipeline.py

@@ -5,7 +5,7 @@ import PytorchBoot.stereotype as stereotype
 from PytorchBoot.factory.component_factory import ComponentFactory
 from PytorchBoot.utils import Log
 
-@stereotype.pipeline("nbv_reconstruction_local_pts_pipeline")
+@stereotype.pipeline("nbv_reconstruction_pipeline_local")
 class NBVReconstructionLocalPointsPipeline(nn.Module):
     def __init__(self, config):
         super(NBVReconstructionLocalPointsPipeline, self).__init__()
@@ -70,23 +70,18 @@ class NBVReconstructionLocalPointsPipeline(nn.Module):
     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 = []
+        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_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)
+            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']   

core/nbv_dataset.py

@@ -4,10 +4,10 @@ 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")
 
@@ -51,7 +51,7 @@ class NBVReconstructionDataset(BaseDataset):
                 scene_name_list.append(scene_name)
         return scene_name_list
 
-    def get_datalist(self, bias=False):
+    def get_datalist(self):
         datalist = []
         for scene_name in self.scene_name_list:
             seq_num = DataLoadUtil.get_label_num(self.root_dir, scene_name)
@@ -80,18 +80,16 @@ class NBVReconstructionDataset(BaseDataset):
                     for data_pair in label_data["data_pairs"]:
                         scanned_views = data_pair[0]
                         next_best_view = data_pair[1]
-                        accept_probability = scanned_views[-1][1]
-                        if accept_probability > np.random.rand():
-                            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,
-                                }
-                            )
+                        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):
@@ -117,8 +115,13 @@ class NBVReconstructionDataset(BaseDataset):
         except Exception as e:
             Log.error(f"Save cache failed: {e}")
 
-    def voxel_downsample_with_mask(self, pts, voxel_size):
-        pass
+    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):
@@ -132,6 +135,9 @@ class NBVReconstructionDataset(BaseDataset):
             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]
@@ -153,8 +159,12 @@ class NBVReconstructionDataset(BaseDataset):
             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)
@@ -167,14 +177,27 @@ class NBVReconstructionDataset(BaseDataset):
         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)
         
+        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)
@@ -200,7 +223,9 @@ class NBVReconstructionDataset(BaseDataset):
             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(
@@ -209,12 +234,14 @@ class NBVReconstructionDataset(BaseDataset):
             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
@@ -230,10 +257,9 @@ if __name__ == "__main__":
     torch.manual_seed(seed)
     np.random.seed(seed)
     config = {
-        "root_dir": "/data/hofee/data/new_full_data",
-        "model_dir": "../data/scaled_object_meshes",
+        "root_dir": "/data/hofee/nbv_rec_part2_preprocessed",
         "source": "nbv_reconstruction_dataset",
-        "split_file": "/data/hofee/data/new_full_data_list/OmniObject3d_train.txt",
+        "split_file": "/data/hofee/data/sample.txt",
         "load_from_preprocess": True,
         "ratio": 0.5,
         "batch_size": 2,

core/pipeline.py

@@ -90,26 +90,51 @@ class NBVReconstructionPipeline(nn.Module):
         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 = self.pts_encoder.encode_points(
-            combined_scanned_pts_batch, require_per_point_feat=False
+        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)
-
-        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)  # Tensor(S x 9)
+        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 = pose_feat_seq
+
+            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
+        return main_feat

runners/inferencer.py

@@ -92,7 +92,8 @@ class Inferencer(Runner):
                         output = self.predict_sequence(data)
                         self.save_inference_result(test_set_name, data["scene_name"], output)
                     except Exception as e:
-                        Log.error(f"Error in scene {scene_name}, {e}")
+                        print(e)
+                        Log.error(f"Error, {e}")
                         continue
                     
             status_manager.set_progress("inference", "inferencer", f"dataset", len(self.test_set_list), len(self.test_set_list))
@@ -116,7 +117,9 @@ class Inferencer(Runner):
         
         ''' data for inference '''
         input_data = {}
+        
         input_data["combined_scanned_pts"] = torch.tensor(data["first_scanned_pts"][0], dtype=torch.float32).to(self.device).unsqueeze(0)
+        input_data["scanned_pts_mask"] = [torch.zeros(input_data["combined_scanned_pts"].shape[1], dtype=torch.bool).to(self.device).unsqueeze(0)]
         input_data["scanned_n_to_world_pose_9d"] = [torch.tensor(data["first_scanned_n_to_world_pose_9d"], dtype=torch.float32).to(self.device)]
         input_data["mode"] = namespace.Mode.TEST
         input_pts_N = input_data["combined_scanned_pts"].shape[1]
@@ -254,6 +257,14 @@ class Inferencer(Runner):
 
         return result
 
+    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 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]

runners/simulator.py

@@ -1,5 +1,5 @@
-import pybullet as p
-import pybullet_data
+# import pybullet as p
+# import pybullet_data
 import numpy as np
 import os
 import time