From 4e4fcb2ce50fbd2257fb9cb59a2e86a307e58a68 Mon Sep 17 00:00:00 2001 From: hofee Date: Thu, 19 Sep 2024 12:12:48 +0000 Subject: [PATCH 1/2] add points_normals under display_table_world_space --- configs/server/train_config.yaml | 2 +- runners/strategy_generator.py | 1 - utils/data_load.py | 20 +++----------------- 3 files changed, 4 insertions(+), 19 deletions(-) diff --git a/configs/server/train_config.yaml b/configs/server/train_config.yaml index 6c47be4..3881eb9 100644 --- a/configs/server/train_config.yaml +++ b/configs/server/train_config.yaml @@ -9,7 +9,7 @@ runner: experiment: name: new_test_overfit_to_world root_dir: "experiments" - use_checkpoint: False + use_checkpoint: True epoch: -1 # -1 stands for last epoch max_epochs: 5000 save_checkpoint_interval: 3 diff --git a/runners/strategy_generator.py b/runners/strategy_generator.py index 4b27e88..5280178 100644 --- a/runners/strategy_generator.py +++ b/runners/strategy_generator.py @@ -117,7 +117,6 @@ class StrategyGenerator(Runner): if self.save_mesh: DataLoadUtil.save_target_mesh_at_world_space(root, model_dir, scene_name) - DataLoadUtil.save_downsampled_world_model_points(root, scene_name, down_sampled_model_pts) def generate_data_pairs(self, useful_view): data_pairs = [] diff --git a/utils/data_load.py b/utils/data_load.py index afa7b04..589337a 100644 --- a/utils/data_load.py +++ b/utils/data_load.py @@ -17,27 +17,11 @@ class DataLoadUtil: path = os.path.join(root,scene_name, f"label.json") return path - @staticmethod - def get_sampled_model_points_path(root, scene_name): - path = os.path.join(root,scene_name, f"sampled_model_points.txt") - return path - @staticmethod def get_scene_seq_length(root, scene_name): camera_params_path = os.path.join(root, scene_name, "camera_params") return len(os.listdir(camera_params_path)) - @staticmethod - def load_downsampled_world_model_points(root, scene_name): - model_path = DataLoadUtil.get_sampled_model_points_path(root, scene_name) - model_points = np.loadtxt(model_path) - return model_points - - @staticmethod - def save_downsampled_world_model_points(root, scene_name, model_points): - model_path = DataLoadUtil.get_sampled_model_points_path(root, scene_name) - np.savetxt(model_path, model_points) - @staticmethod def load_mesh_at(model_dir, object_name, world_object_pose): model_path = os.path.join(model_dir, object_name, "mesh.obj") @@ -269,7 +253,9 @@ class DataLoadUtil: return overlapping_points @staticmethod - def load_points_normals(root, scene_name): + def load_points_normals(root, scene_name, display_table_as_world_space_origin=True): points_path = os.path.join(root, scene_name, "points_and_normals.txt") points_normals = np.loadtxt(points_path) + if display_table_as_world_space_origin: + points_normals[:,:3] = points_normals[:,:3] - DataLoadUtil.DISPLAY_TABLE_POSITION return points_normals \ No newline at end of file From 18333e68313907a0bd3da0a7cf9098a703ec0803 Mon Sep 17 00:00:00 2001 From: hofee Date: Fri, 20 Sep 2024 06:43:19 +0000 Subject: [PATCH 2/2] change dataset.py to nbv_dataset.py --- core/{dataset.py => nbv_dataset.py} | 0 modules/transformer_seq_encoder.py | 8 +------- 2 files changed, 1 insertion(+), 7 deletions(-) rename core/{dataset.py => nbv_dataset.py} (100%) diff --git a/core/dataset.py b/core/nbv_dataset.py similarity index 100% rename from core/dataset.py rename to core/nbv_dataset.py diff --git a/modules/transformer_seq_encoder.py b/modules/transformer_seq_encoder.py index 8b22b4f..1eae505 100644 --- a/modules/transformer_seq_encoder.py +++ b/modules/transformer_seq_encoder.py @@ -22,7 +22,6 @@ class TransformerSequenceEncoder(nn.Module): self.fc = nn.Linear(embed_dim, config["output_dim"]) def encode_sequence(self, pts_embedding_list_batch, pose_embedding_list_batch): - # Combine features and pad sequences combined_features_batch = [] lengths = [] @@ -36,16 +35,11 @@ class TransformerSequenceEncoder(nn.Module): combined_tensor = pad_sequence(combined_features_batch, batch_first=True) # Shape: [batch_size, max_seq_len, embed_dim] - # Prepare mask for padding 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 encoding + transformer_output = self.transformer_encoder(combined_tensor, src_key_padding_mask=padding_mask) - - # Mean pooling final_feature = transformer_output.mean(dim=1) - - # Fully connected layer final_output = self.fc(final_feature) return final_output