feat: enhance surface category handling in AnnotationFilesGenerator a…

…nd DeepenToT4Converter Signed-off-by: Shunsuke Miura <[email protected]>
tier4 · Dec 17, 2024 · 87fd385 · 87fd385
1 parent c77c682
commit 87fd385
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 9 deletions.
diff --git a/perception_dataset/deepen/deepen_to_t4_converter.py b/perception_dataset/deepen/deepen_to_t4_converter.py
@@ -4,6 +4,7 @@
 import re
 import shutil
 from typing import Any, Dict, List, Optional, Union
+import yaml
 
 from nuscenes.nuscenes import NuScenes
 
@@ -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,8 @@ 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,

diff --git a/perception_dataset/fastlabel_to_t4/fastlabel_2d_to_t4_converter.py b/perception_dataset/fastlabel_to_t4/fastlabel_2d_to_t4_converter.py
@@ -226,7 +226,7 @@ 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]
                             )

diff --git a/perception_dataset/t4_dataset/annotation_files_generator.py b/perception_dataset/t4_dataset/annotation_files_generator.py
@@ -26,7 +26,7 @@
 
 
 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 +62,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():