From 87fd38530c4ab6b746fb529afe1d4a74d0bc3170 Mon Sep 17 00:00:00 2001 From: Shunsuke Miura Date: Tue, 17 Dec 2024 14:52:14 +0900 Subject: [PATCH] feat: enhance surface category handling in AnnotationFilesGenerator and DeepenToT4Converter Signed-off-by: Shunsuke Miura --- perception_dataset/deepen/deepen_to_t4_converter.py | 9 ++++++++- .../fastlabel_to_t4/fastlabel_2d_to_t4_converter.py | 2 +- .../t4_dataset/annotation_files_generator.py | 9 ++------- 3 files changed, 11 insertions(+), 9 deletions(-) diff --git a/perception_dataset/deepen/deepen_to_t4_converter.py b/perception_dataset/deepen/deepen_to_t4_converter.py index 6e9adf28..d5ba40d5 100644 --- 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 index 3ebbcf8d..f2787fbe 100644 --- 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 index ffbaa381..69175596 100644 --- 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():