add global_pts_pipeline and pose_seq_encooder

app_train.py

@@ -5,4 +5,4 @@ from PytorchBoot.runners.trainer import DefaultTrainer
 class TrainApp:
     @staticmethod
     def start():
-        DefaultTrainer("configs/server/train_config.yaml").run()
+        DefaultTrainer("configs/server/server_train_config.yaml").run()

configs/server/server_split_dataset_config.yaml

@@ -9,8 +9,8 @@ runner:
     name: debug
     root_dir: "experiments"
   
-  split:
+  split: # 
-    root_dir: "../data/sample_for_training_preprocessed/sample_preprocessed_scenes"
+    root_dir: "/home/data/hofee/project/nbv_rec/data/nbv_rec_data_512_preproc_npy"
     type: "unseen_instance" # "unseen_category"
     datasets:
       OmniObject3d_train: 

configs/server/server_train_config.yaml

@@ -0,0 +1,145 @@
+
+runner:
+  general:
+    seed: 0
+    device: cuda
+    cuda_visible_devices: "0"
+    parallel: False
+    
+  experiment:
+    name: overfit_w_global_feat_wo_local_pts_feat_small
+    root_dir: "experiments"
+    use_checkpoint: False
+    epoch: -1 # -1 stands for last epoch
+    max_epochs: 5000
+    save_checkpoint_interval: 1
+    test_first: True  
+
+  train:
+    optimizer:
+      type: Adam
+      lr: 0.0001
+    losses: 
+      - gf_loss
+    dataset: OmniObject3d_train
+  test:
+    frequency: 3 # test frequency
+    dataset_list:
+      #- OmniObject3d_test
+      - OmniObject3d_val
+
+  pipeline: nbv_reconstruction_global_pts_pipeline
+  
+dataset:
+  OmniObject3d_train:
+    root_dir: "/home/data/hofee/project/nbv_rec/data/nbv_rec_data_512_preproc_npy"
+    model_dir: "../data/scaled_object_meshes"
+    source: nbv_reconstruction_dataset
+    split_file: "/home/data/hofee/project/nbv_rec/data/OmniObject3d_sample.txt"
+    type: train
+    cache: True
+    ratio: 1
+    batch_size: 160
+    num_workers: 16
+    pts_num: 4096
+    load_from_preprocess: True
+
+  OmniObject3d_test:
+    root_dir: "/home/data/hofee/project/nbv_rec/data/nbv_rec_data_512_preproc_npy"
+    model_dir: "../data/scaled_object_meshes"
+    source: nbv_reconstruction_dataset
+    split_file: "/home/data/hofee/project/nbv_rec/data/OmniObject3d_test.txt"
+    type: test
+    cache: True
+    filter_degree: 75
+    eval_list:
+      - pose_diff
+    ratio: 0.05
+    batch_size: 1
+    num_workers: 12
+    pts_num: 4096
+    load_from_preprocess: True
+
+  OmniObject3d_val:
+    root_dir: "/home/data/hofee/project/nbv_rec/data/nbv_rec_data_512_preproc_npy"
+    model_dir: "../data/scaled_object_meshes"
+    source: nbv_reconstruction_dataset
+    split_file: "/home/data/hofee/project/nbv_rec/data/OmniObject3d_sample.txt"
+    type: test
+    cache: True
+    filter_degree: 75
+    eval_list:
+      - pose_diff
+    ratio: 1
+    batch_size: 1
+    num_workers: 12
+    pts_num: 4096
+    load_from_preprocess: True
+
+
+pipeline:
+  nbv_reconstruction_local_pts_pipeline:
+    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_global_pts_pipeline:
+    modules:
+      pts_encoder: pointnet_encoder
+      pose_seq_encoder: transformer_pose_seq_encoder
+      pose_encoder: pose_encoder
+      view_finder: gf_view_finder
+    eps: 1e-5
+    global_scanned_feat: True
+
+
+    
+module:
+
+  pointnet_encoder:
+    in_dim: 3
+    out_dim: 1024
+    global_feat: True
+    feature_transform: False
+
+  transformer_seq_encoder:
+    pts_embed_dim: 1024
+    pose_embed_dim: 256
+    num_heads: 4
+    ffn_dim: 256
+    num_layers: 3
+    output_dim: 2048
+
+  transformer_pose_seq_encoder:
+    pose_embed_dim: 256
+    num_heads: 4
+    ffn_dim: 256
+    num_layers: 3
+    output_dim: 1024
+
+  gf_view_finder:
+    t_feat_dim: 128
+    pose_feat_dim: 256
+    main_feat_dim: 2048
+    regression_head: Rx_Ry_and_T
+    pose_mode: rot_matrix
+    per_point_feature: False
+    sample_mode: ode
+    sampling_steps: 500
+    sde_mode: ve
+
+  pose_encoder:
+    pose_dim: 9
+    out_dim: 256
+
+loss_function:
+  gf_loss:
+
+evaluation_method:
+  pose_diff:
+  coverage_rate_increase: 
+    renderer_path: "../blender/data_renderer.py" 

configs/server/train_config.yaml

@@ -1,106 +0,0 @@
-
-runner:
-  general:
-    seed: 0
-    device: cuda
-    cuda_visible_devices: "0,1,2,3,4,5,6,7"
-    parallel: False
-    
-  experiment:
-    name: new_test_overfit_to_world_preprocessed
-    root_dir: "experiments"
-    use_checkpoint: False
-    epoch: -1 # -1 stands for last epoch
-    max_epochs: 5000
-    save_checkpoint_interval: 3
-    test_first: True  
-
-  train:
-    optimizer:
-      type: Adam
-      lr: 0.0001
-    losses: 
-      - gf_loss
-    dataset: OmniObject3d_train
-  test:
-    frequency: 3 # test frequency
-    dataset_list:
-      - OmniObject3d_test
-
-  pipeline: nbv_reconstruction_pipeline
-  
-dataset:
-  OmniObject3d_train:
-    root_dir: "../data/sample_for_training_preprocessed/sample_preprocessed_scenes"
-    model_dir: "../data/scaled_object_meshes"
-    source: nbv_reconstruction_dataset
-    split_file: "../data/sample_for_training_preprocessed/OmniObject3d_train.txt"
-    type: train
-    cache: True
-    ratio: 1
-    batch_size: 160
-    num_workers: 16
-    pts_num: 4096
-    load_from_preprocess: True
-
-  OmniObject3d_test:
-    root_dir: "../data/sample_for_training_preprocessed/sample_preprocessed_scenes"
-    model_dir: "../data/scaled_object_meshes"
-    source: nbv_reconstruction_dataset
-    split_file: "../data/sample_for_training_preprocessed/OmniObject3d_train.txt"
-    type: test
-    cache: True
-    filter_degree: 75
-    eval_list:
-      - pose_diff
-    ratio: 0.1
-    batch_size: 1
-    num_workers: 12
-    pts_num: 4096
-    load_from_preprocess: True
-
-pipeline:
-  nbv_reconstruction_pipeline:
-    pts_encoder: pointnet_encoder
-    seq_encoder: transformer_seq_encoder
-    pose_encoder: pose_encoder
-    view_finder: gf_view_finder
-    
-module:
-
-  pointnet_encoder:
-    in_dim: 3
-    out_dim: 1024
-    global_feat: True
-    feature_transform: False
-
-  transformer_seq_encoder:
-    pts_embed_dim: 1024
-    pose_embed_dim: 256
-    num_heads: 4
-    ffn_dim: 256
-    num_layers: 3
-    output_dim: 2048
-
-  gf_view_finder:
-    t_feat_dim: 128
-    pose_feat_dim: 256
-    main_feat_dim: 2048
-    regression_head: Rx_Ry_and_T
-    pose_mode: rot_matrix
-    per_point_feature: False
-    sample_mode: ode
-    sampling_steps: 500
-    sde_mode: ve
-
-  pose_encoder:
-    pose_dim: 9
-    out_dim: 256
-
-loss_function:
-  gf_loss:
-
-evaluation_method:
-  pose_diff:
-  coverage_rate_increase: 
-    renderer_path: "../blender/data_renderer.py" 

core/global_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_global_pts_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.pose_seq_encoder = ComponentFactory.create(namespace.Stereotype.MODULE, self.module_config["pose_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_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 = []
+        
+        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))
+
+        main_feat = self.pose_seq_encoder.encode_sequence(pose_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/local_pts_pipeline.py

@@ -5,16 +5,18 @@ import PytorchBoot.stereotype as stereotype
 from PytorchBoot.factory.component_factory import ComponentFactory
 from PytorchBoot.utils import Log
 
-@stereotype.pipeline("nbv_reconstruction_pipeline")
+@stereotype.pipeline("nbv_reconstruction_local_pts_pipeline")
-class NBVReconstructionPipeline(nn.Module):
+class NBVReconstructionLocalPointsPipeline(nn.Module):
     def __init__(self, config):
-        super(NBVReconstructionPipeline, self).__init__()
+        super(NBVReconstructionLocalPointsPipeline, self).__init__()
         self.config = config
-        self.pts_encoder = ComponentFactory.create(namespace.Stereotype.MODULE, config["pts_encoder"])
+        self.module_config = config["modules"]
-        self.pose_encoder = ComponentFactory.create(namespace.Stereotype.MODULE, config["pose_encoder"])
+        self.pts_encoder = ComponentFactory.create(namespace.Stereotype.MODULE, self.module_config["pts_encoder"])
-        self.seq_encoder = ComponentFactory.create(namespace.Stereotype.MODULE, config["seq_encoder"])
+        self.pose_encoder = ComponentFactory.create(namespace.Stereotype.MODULE, self.module_config["pose_encoder"])
-        self.view_finder = ComponentFactory.create(namespace.Stereotype.MODULE, config["view_finder"])
+        self.seq_encoder = ComponentFactory.create(namespace.Stereotype.MODULE, self.module_config["seq_encoder"])
-        self.eps = 1e-5
+        self.view_finder = ComponentFactory.create(namespace.Stereotype.MODULE, self.module_config["view_finder"])
+        self.eps = float(self.config["eps"])
+        self.enable_global_scanned_feat = self.config["global_scanned_feat"]
         
     def forward(self, data):
         mode = data["mode"]
@@ -38,14 +40,14 @@ class NBVReconstructionPipeline(nn.Module):
         return perturbed_x, random_t, target_score, std
     
     def forward_train(self, data):
-        seq_feat = self.get_seq_feat(data)
+        main_feat = self.get_main_feat(data)
         ''' get std '''
         best_to_world_pose_9d_batch = data["best_to_world_pose_9d"]
         perturbed_x, random_t, target_score, std = self.pertube_data(best_to_world_pose_9d_batch)
         input_data = {
             "sampled_pose": perturbed_x,
             "t": random_t,
-            "seq_feat": seq_feat,
+            "main_feat": main_feat,
         }
         estimated_score = self.view_finder(input_data)
         output = {
@@ -56,29 +58,44 @@ class NBVReconstructionPipeline(nn.Module):
         return output
     
     def forward_test(self,data):
-        seq_feat = self.get_seq_feat(data)
+        main_feat = self.get_main_feat(data)
-        estimated_delta_rot_9d, in_process_sample = self.view_finder.next_best_view(seq_feat)
+        estimated_delta_rot_9d, in_process_sample = self.view_finder.next_best_view(main_feat)
         result = {
             "pred_pose_9d": estimated_delta_rot_9d,
             "in_process_sample": in_process_sample
         }
         return result
     
-    def get_seq_feat(self, data):
+    
+    def get_main_feat(self, data):
         scanned_pts_batch = data['scanned_pts']
         scanned_n_to_world_pose_9d_batch = data['scanned_n_to_world_pose_9d']
+        
+
+        device = next(self.parameters()).device
+
+        
+
         pts_feat_seq_list = []
         pose_feat_seq_list = []
-        device = next(self.parameters()).device
+        
         for scanned_pts,scanned_n_to_world_pose_9d in zip(scanned_pts_batch,scanned_n_to_world_pose_9d_batch):
             
             scanned_pts = scanned_pts.to(device)
             scanned_n_to_world_pose_9d = scanned_n_to_world_pose_9d.to(device)
             pts_feat_seq_list.append(self.pts_encoder.encode_points(scanned_pts))
             pose_feat_seq_list.append(self.pose_encoder.encode_pose(scanned_n_to_world_pose_9d))
+
+        main_feat = self.seq_encoder.encode_sequence(pts_feat_seq_list, pose_feat_seq_list)
         
-        seq_feat = self.seq_encoder.encode_sequence(pts_feat_seq_list, pose_feat_seq_list)
+        if self.enable_global_scanned_feat:
-        if torch.isnan(seq_feat).any():
+            combined_scanned_pts_batch = data['combined_scanned_pts']   
-            Log.error("nan in seq_feat", True)
+            global_scanned_feat = self.pts_encoder.encode_points(combined_scanned_pts_batch)
-        return seq_feat
+            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/nbv_dataset.py

@@ -7,12 +7,11 @@ 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")
+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
-from utils.reconstruction import ReconstructionUtil
 
 
 @stereotype.dataset("nbv_reconstruction_dataset")
@@ -35,7 +34,7 @@ class NBVReconstructionDataset(BaseDataset):
             self.model_dir = config["model_dir"]
             self.filter_degree = config["filter_degree"]
         if self.type == namespace.Mode.TRAIN:
-            scale_ratio = 1
+            scale_ratio = 100
             self.datalist = self.datalist*scale_ratio
         if self.cache:
             expr_root = ConfigManager.get("runner", "experiment", "root_dir")
@@ -56,20 +55,35 @@ class NBVReconstructionDataset(BaseDataset):
     def get_datalist(self):
         datalist = []
         for scene_name in self.scene_name_list:
-            label_path = DataLoadUtil.get_label_path_old(self.root_dir, scene_name)
+            seq_num = DataLoadUtil.get_label_num(self.root_dir, scene_name)
-            label_data = DataLoadUtil.load_label(label_path)
+            scene_max_coverage_rate = 0
-            for data_pair in label_data["data_pairs"]:
+            max_coverage_rate_list = []
-                scanned_views = data_pair[0]
+
-                next_best_view = data_pair[1]
+            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"]
-                datalist.append(
+                if max_coverage_rate > scene_max_coverage_rate:
-                    {
+                    scene_max_coverage_rate = max_coverage_rate
-                        "scanned_views": scanned_views,
+                max_coverage_rate_list.append(max_coverage_rate)
-                        "next_best_view": next_best_view,
+            mean_coverage_rate = np.mean(max_coverage_rate_list)
-                        "max_coverage_rate": max_coverage_rate,
+
-                        "scene_name": scene_name,
+            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
+                        })
+                    break # TODO: for small version debug
         return datalist
 
     def preprocess_cache(self):
@@ -102,7 +116,7 @@ class NBVReconstructionDataset(BaseDataset):
         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["max_coverage_rate"]
+        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 = [], [], []
       
@@ -151,13 +165,18 @@ class NBVReconstructionDataset(BaseDataset):
         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 = 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)
         data_item = {
             "scanned_pts": np.asarray(scanned_views_pts,dtype=np.float32),
+            "combined_scanned_pts": np.asarray(random_downsampled_combined_scanned_pts_np,dtype=np.float32),
             "scanned_coverage_rate": scanned_coverages_rate,
             "scanned_n_to_world_pose_9d": np.asarray(scanned_n_to_world_pose,dtype=np.float32),
             "best_coverage_rate": nbv_coverage_rate,
             "best_to_world_pose_9d": np.asarray(best_to_world_9d,dtype=np.float32),
-            "max_coverage_rate": max_coverage_rate,
+            "seq_max_coverage_rate": max_coverage_rate,
             "scene_name": scene_name
         }
         
@@ -195,10 +214,11 @@ class NBVReconstructionDataset(BaseDataset):
             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["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])
             if "first_frame_to_world" in batch[0]:
                 collate_data["first_frame_to_world"] = torch.stack([torch.tensor(item["first_frame_to_world"]) 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", "first_frame_to_world"]:
+                if key not in ["scanned_pts", "scanned_n_to_world_pose_9d", "best_to_world_pose_9d", "first_frame_to_world", "combined_scanned_pts"]:
                     collate_data[key] = [item[key] for item in batch]
             return collate_data
         return collate_fn
@@ -211,11 +231,11 @@ if __name__ == "__main__":
     torch.manual_seed(seed)
     np.random.seed(seed)
     config = {
-        "root_dir": "/media/hofee/repository/nbv_reconstruction_data_512",
+        "root_dir": "/home/data/hofee/project/nbv_rec/data/nbv_rec_data_512_preproc_npy",
-        "model_dir": "/media/hofee/data/data/scaled_object_meshes",
+        "model_dir": "/home/data/hofee/project/nbv_rec/data/scaled_object_meshes",
         "source": "nbv_reconstruction_dataset",
-        "split_file": "/media/hofee/data/project/python/nbv_reconstruction/sample_for_training/OmniObject3d_train.txt",
+        "split_file": "/home/data/hofee/project/nbv_rec/data/OmniObject3d_test.txt",
-        "load_from_preprocess": False,
+        "load_from_preprocess": True,
         "ratio": 0.5,
         "batch_size": 2,
         "filter_degree": 75,

modules/func_lib/samplers.py

@@ -32,7 +32,7 @@ def cond_ode_sampler(
     init_x=None,
 ):
     pose_dim = PoseUtil.get_pose_dim(pose_mode)
-    batch_size = data["seq_feat"].shape[0]
+    batch_size = data["main_feat"].shape[0]
     init_x = (
         prior((batch_size, pose_dim), T=T).to(device)
         if init_x is None

modules/gf_view_finder.py

@@ -80,13 +80,13 @@ class GradientFieldViewFinder(nn.Module):
         """
         Args:
             data, dict {
-                'seq_feat': [bs, c]
+                'main_feat': [bs, c]
                 'pose_sample': [bs, pose_dim]
                 't': [bs, 1]
             }
         """
 
-        seq_feat = data['seq_feat']
+        main_feat = data['main_feat']
         sampled_pose = data['sampled_pose']
         t = data['t']
         t_feat = self.t_encoder(t.squeeze(1))
@@ -95,7 +95,7 @@ class GradientFieldViewFinder(nn.Module):
         if self.per_point_feature:
             raise NotImplementedError
         else:
-            total_feat = torch.cat([seq_feat, t_feat, pose_feat], dim=-1)
+            total_feat = torch.cat([main_feat, t_feat, pose_feat], dim=-1)
         _, std = self.marginal_prob_fn(total_feat, t)
 
         if self.regression_head == 'Rx_Ry_and_T':
@@ -134,9 +134,9 @@ class GradientFieldViewFinder(nn.Module):
 
         return in_process_sample, res
 
-    def next_best_view(self, seq_feat):
+    def next_best_view(self, main_feat):
         data = {
-            'seq_feat': seq_feat,
+            'main_feat': main_feat,
         }
         in_process_sample, res = self.sample(data)
         return res.to(dtype=torch.float32), in_process_sample

modules/transformer_pose_seq_encoder.py

@@ -0,0 +1,63 @@
+import torch
+from torch import nn
+from torch.nn.utils.rnn import pad_sequence
+import PytorchBoot.stereotype as stereotype
+
+
+@stereotype.module("transformer_pose_seq_encoder")
+class TransformerPoseSequenceEncoder(nn.Module):
+    def __init__(self, config):
+        super(TransformerPoseSequenceEncoder, self).__init__()
+        self.config = config
+        embed_dim = config["pose_embed_dim"]
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=embed_dim,
+            nhead=config["num_heads"],
+            dim_feedforward=config["ffn_dim"],
+            batch_first=True,
+        )
+        self.transformer_encoder = nn.TransformerEncoder(
+            encoder_layer, num_layers=config["num_layers"]
+        )
+        self.fc = nn.Linear(embed_dim, config["output_dim"])
+
+    def encode_sequence(self, pose_embedding_list_batch):
+        
+        lengths = []
+        
+        for pose_embedding_list in  pose_embedding_list_batch:
+            lengths.append(len(pose_embedding_list))
+
+        combined_tensor = pad_sequence(pose_embedding_list_batch, batch_first=True)  # Shape: [batch_size, max_seq_len, embed_dim]
+
+        max_len = max(lengths)
+        padding_mask = torch.tensor([([0] * length + [1] * (max_len - length)) for length in lengths], dtype=torch.bool).to(combined_tensor.device)
+
+        transformer_output = self.transformer_encoder(combined_tensor, src_key_padding_mask=padding_mask)
+        final_feature = transformer_output.mean(dim=1)
+        final_output = self.fc(final_feature)
+
+        return final_output
+
+
+if __name__ == "__main__":
+    config = {
+        "pose_embed_dim": 256,
+        "num_heads": 4,  
+        "ffn_dim": 256,
+        "num_layers": 3, 
+        "output_dim": 1024, 
+    }
+
+    encoder = TransformerPoseSequenceEncoder(config)
+    seq_len = [5, 8, 9, 4]
+    batch_size = 4
+
+    pose_embedding_list_batch = [
+        torch.randn(seq_len[idx], config["pose_embed_dim"]) for idx in range(batch_size)
+    ]
+    output_feature = encoder.encode_sequence(
+        pose_embedding_list_batch
+    )
+    print("Encoded Feature:", output_feature)
+    print("Feature Shape:", output_feature.shape)