Update core codes, setup codes and add docstrings

README.md

@@ -145,6 +145,7 @@ This work is under the <a rel="license" href="http://creativecommons.org/license
 ## 👏 Acknowledgements
 
 - [OpenMMLab](https://github.com/open-mmlab): Our dataset code uses [MMEngine](https://github.com/open-mmlab/mmengine) and our model is built upon of [MMDetection3D](https://github.com/open-mmlab/mmdetection3d).
+- [PyTorch3D](https://github.com/facebookresearch/pytorch3d): We use some functions supported in PyTorch3D for efficient computations on fundamental 3D data structures.
 - [ScanNet](https://github.com/ScanNet/ScanNet), [3RScan](https://github.com/WaldJohannaU/3RScan), [Matterport3D](https://github.com/niessner/Matterport): Our dataset uses the raw data from these datasets.
 - [ReferIt3D](https://github.com/referit3d/referit3d): We refer to the SR3D's approach to obtaining the language prompt annotations.
 - [SUSTechPOINTS](https://github.com/naurril/SUSTechPOINTS): Our annotation tool is developed based on the open-source framework used by SUSTechPOINTS.

embodiedscan/converter/generate_image_scannet.py

@@ -23,6 +23,7 @@ class RGBDFrame:
     """Class for single ScanNet RGB-D image processing."""
 
     def load(self, file_handle):
+        """Load basic information of a given RGBD frame."""
         self.camera_to_world = np.asarray(struct.unpack(
             'f' * 16, file_handle.read(16 * 4)),
                                           dtype=np.float32).reshape(4, 4)
@@ -38,10 +39,12 @@ def load(self, file_handle):
                           file_handle.read(self.depth_size_bytes)))
 
     def decompress_depth(self, compression_type):
+        """Decompress the depth data."""
         assert compression_type == 'zlib_ushort'
         return zlib.decompress(self.depth_data)
 
     def decompress_color(self, compression_type):
+        """Decompress the RGB image data."""
         assert compression_type == 'jpeg'
         return imageio.imread(self.color_data)
 
@@ -57,6 +60,7 @@ def __init__(self, filename, fast=False):
         self.load(filename, fast)
 
     def load(self, filename, fast):
+        """Load a single scene data with multiple RGBD frames."""
         with open(filename, 'rb') as f:
             version = struct.unpack('I', f.read(4))[0]
             assert self.version == version
@@ -99,6 +103,7 @@ def load(self, filename, fast):
                     self.frames.append(frame)
 
     def export_depth_images(self, output_path):
+        """Export depth images to the output path."""
         if not os.path.exists(output_path):
             os.makedirs(output_path)
         for f in range(len(self.frames)):
@@ -111,6 +116,7 @@ def export_depth_images(self, output_path):
                              self.index_to_str(self.index[f]) + '.png'), depth)
 
     def export_color_images(self, output_path):
+        """Export RGB images to the output path."""
         if not os.path.exists(output_path):
             os.makedirs(output_path)
         for f in range(len(self.frames)):
@@ -122,15 +128,18 @@ def export_color_images(self, output_path):
 
     @staticmethod
     def index_to_str(index):
+        """Convert the sample index to string."""
         return str(index).zfill(5)
 
     @staticmethod
     def save_mat_to_file(matrix, filename):
+        """Save a matrix to file."""
         with open(filename, 'w') as f:
             for line in matrix:
                 np.savetxt(f, line[np.newaxis], fmt='%f')
 
     def export_poses(self, output_path):
+        """Export camera poses to the output path."""
         if not os.path.exists(output_path):
             os.makedirs(output_path)
         for f in range(len(self.frames)):
@@ -140,12 +149,14 @@ def export_poses(self, output_path):
                              self.index_to_str(self.index[f]) + '.txt'))
 
     def export_intrinsics(self, output_path):
+        """Export the intrinsic matrix to the output path."""
         if not os.path.exists(output_path):
             os.makedirs(output_path)
         self.save_mat_to_file(self.intrinsic_color,
                               os.path.join(output_path, 'intrinsic.txt'))
 
     def export_depth_intrinsics(self, output_path):
+        """Export the depth intrinsic matrix to the output path."""
         if not os.path.exists(output_path):
             os.makedirs(output_path)
         self.save_mat_to_file(self.intrinsic_depth,
@@ -167,6 +178,7 @@ def process_scene(path, fast, idx):
 
 
 def process_directory(path, fast, nproc):
+    """Process the files in a directory with parallel support."""
     mmengine.track_parallel_progress(func=partial(process_scene, path, fast),
                                      tasks=os.listdir(path),
                                      nproc=nproc)

embodiedscan/embodied_dataset.py

@@ -5,13 +5,40 @@
 import mmengine
 import numpy as np
 from mmdet3d.registry import DATASETS
-from mmdet3d.structures import get_box_type
 from mmengine.dataset import BaseDataset
 from mmengine.fileio import load
 
+from embodiedscan.structures import get_box_type
+
 
 @DATASETS.register_module()
 class EmbodiedScanDataset(BaseDataset):
+    r"""EmbodiedScan Dataset.
+
+    This class serves as the API for experiments on the EmbodiedScan Dataset.
+
+    Please refer to `EmbodiedScan Dataset
+    <https://github.com/OpenRobotLab/EmbodiedScan>`_  for data downloading.
+
+    Args:
+        data_root (str): Path of dataset root.
+        ann_file (str): Path of annotation file.
+        metainfo (dict, optional): Meta information for dataset, such as class
+            information. Defaults to None.
+        pipeline (List[dict]): Pipeline used for data processing.
+            Defaults to [].
+        test_mode (bool): Whether the dataset is in test mode.
+            Defaults to False.
+        load_eval_anns (bool): Whether to load evaluation annotations.
+            Defaults to True. Only take effect when test_mode is True.
+        filter_empty_gt (bool): Whether to filter the data with empty GT.
+            If it's set to be True, the example with empty annotations after
+            data pipeline will be dropped and a random example will be chosen
+            in `__getitem__`. Defaults to True.
+        remove_dontcare (bool): Whether to remove objects that we do not care.
+            Defaults to False.
+        box_type_3d (str): To be deprecated?
+    """
 
     def __init__(self,
                  data_root: str,

embodiedscan/eval/det_metric.py

@@ -1,19 +1,20 @@
-# Copyright (c) OpenMMLab. All rights reserved.
+# Copyright (c) OpenRobotLab. All rights reserved.
 import logging
 from collections import OrderedDict
 from typing import Dict, List, Optional, Sequence, Union
 
 import numpy as np
 from mmdet3d.evaluation import indoor_eval
 from mmdet3d.registry import METRICS
-from mmdet3d.structures import get_box_type
 from mmdet.evaluation import eval_map
 from mmengine.dist import (broadcast_object_list, collect_results,
                            is_main_process)
 from mmengine.evaluator import BaseMetric
 from mmengine.evaluator.metric import _to_cpu
 from mmengine.logging import MMLogger, print_log
 
+from embodiedscan.structures import get_box_type
+
 
 @METRICS.register_module()
 class IndoorDetMetric(BaseMetric):

embodiedscan/eval/indoor_eval.py

@@ -1,4 +1,4 @@
-# Copyright (c) OpenMMLab. All rights reserved.
+# Copyright (c) OpenRobotLab. All rights reserved.
 import numpy as np
 import torch
 from mmengine.logging import print_log

embodiedscan/eval/occupancy_metric.py

@@ -1,4 +1,4 @@
-# Copyright (c) OpenMMLab. All rights reserved.
+# Copyright (c) OpenRobotLab. All rights reserved.
 import logging
 from typing import Dict, Optional, Sequence