feat: enhance surface category handling in AnnotationFilesGenerator and DeepenToT4Converter

perception_dataset/deepen/deepen_to_t4_converter.py

@@ -6,6 +6,7 @@
 from typing import Any, Dict, List, Optional, Union
 
 from nuscenes.nuscenes import NuScenes
+import yaml
 
 from perception_dataset.abstract_converter import AbstractConverter
 from perception_dataset.deepen.deepen_annotation import (
@@ -57,6 +58,11 @@ def __init__(
         self._label_info: Optional[LabelInfo] = label_info
 
         self._topic_list_yaml: Union[List, Dict] = topic_list
+        if description.get("surface_categories"):
+            with open(description["surface_categories"], "r") as f:
+                self._surface_categories: List[str] = yaml.safe_load(f)
+        else:
+            self._surface_categories = []
 
     def convert(self):
         camera_index: Dict[str, int] = self._description["camera_index"]
@@ -111,7 +117,10 @@ def convert(self):
         for t4data_name, dataset_id in self._t4data_name_to_deepen_dataset_id.items():
             output_dir = osp.join(self._output_base, t4data_name, self._t4_dataset_dir_name)
             input_dir = osp.join(self._input_base, t4data_name)
-            annotation_files_generator = AnnotationFilesGenerator(description=self._description)
+
+            annotation_files_generator = AnnotationFilesGenerator(
+                description=self._description, surface_categories=self._surface_categories
+            )
             annotation_files_generator.convert_one_scene(
                 input_dir=input_dir,
                 output_dir=output_dir,

perception_dataset/fastlabel_to_t4/fastlabel_2d_to_t4_converter.py

@@ -226,7 +226,10 @@ def _format_fastlabel_annotation(self, annotations: Dict[str, List[Dict[str, Any
                                 "sensor_id": self._camera2idx[camera],
                             }
                         )
-                        if self._label_converter.is_object_label(category_label):
+                        if (
+                            self._label_converter.is_object_label(category_label)
+                            and category_label not in self._surface_categories
+                        ):
                             label_t4_dict["two_d_box"] = _convert_polygon_to_bbox(
                                 a["points"][0][0]
                             )

perception_dataset/t4_dataset/annotation_files_generator.py

@@ -8,7 +8,6 @@
 import numpy as np
 from nuscenes.nuscenes import NuScenes
 from pycocotools import mask as cocomask
-import yaml
 
 from perception_dataset.constants import SENSOR_ENUM
 from perception_dataset.t4_dataset.classes.abstract_class import AbstractTable
@@ -26,7 +25,12 @@
 
 
 class AnnotationFilesGenerator:
-    def __init__(self, with_camera: bool = True, description: Dict[str, Dict[str, str]] = {}):
+    def __init__(
+        self,
+        with_camera: bool = True,
+        description: Dict[str, Dict[str, str]] = {},
+        surface_categories: List[str] = [],
+    ):
         # TODO(yukke42): remove the hard coded attribute description
         self._attribute_table = AttributeTable(
             name_to_description={},
@@ -62,12 +66,7 @@ def __init__(self, with_camera: bool = True, description: Dict[str, Dict[str, st
         else:
             self._camera2idx = None
         self._with_lidar = description.get("with_lidar", True)
-
-        if description.get("surface_categories"):
-            with open(description["surface_categories"], "r") as f:
-                self._surface_categories: List[str] = yaml.safe_load(f)
-        else:
-            self._surface_categories = []
+        self._surface_categories: List[str] = surface_categories
 
     def save_tables(self, anno_dir: str):
         for cls_attr in self.__dict__.values():