feat: merge T4 Datasets with 2D Camera annotations and 3D LiDAR annot…

…ations (#96)

* add t4 2d-3d merger

Signed-off-by: Shunsuke Miura <[email protected]>

* add missing surface_ann output

Signed-off-by: Shunsuke Miura <[email protected]>

* Update the document, remove unsused codes

Signed-off-by: Shunsuke Miura <[email protected]>

* add error handling, remove unnecessary prints

Signed-off-by: Shunsuke Miura <[email protected]>

* cosmetic change

Signed-off-by: Shunsuke Miura <[email protected]>

* add an output

Signed-off-by: Shunsuke Miura <[email protected]>

* pre-commit

Signed-off-by: Shunsuke Miura <[email protected]>

* Update config/label/attribute.yaml

Co-authored-by: kminoda <[email protected]>

---------

Signed-off-by: Shunsuke Miura <[email protected]>
Co-authored-by: kminoda <[email protected]>

config/label/attribute.yaml

@@ -1,5 +1,6 @@
 pedestrian_state:
-    sitting: [pedestrian_state.sitting_lying_down, pedestrian_state.sitting]
+    # The word `siting` is a typo included intentionally since the existing T4dataset also had the same typo.
+    sitting: [pedestrian_state.sitting_lying_down, pedestrian_state.siting_lying_down, pedestrian_state.sitting]
     lying_down: [pedestrian_state.lying_down]
     standing: [pedestrian_state.standing, pedestrian_state.moving]
 vehicle_state:
@@ -24,6 +25,7 @@ extremities_state:
 emergency_vehicle_lights_state:
     on: [emergency_vehicle_lights_state.on]
     off: [emergency_vehicle_lights_state.off]
+    unknown: [emergency_vehicle_lights_state.unknown]
 object_state:
     still: [object_state.still]
 truncation_state:

config/merge_2d_t4dataset_to_3d.yaml

@@ -0,0 +1,7 @@
+task: merge_2d_t4dataset_to_3d
+conversion:
+  input_base: data/t4_format_2d_only
+  output_base: data/t4_format_3d_only
+  dataset_corresponding:
+    # output 3D dataset name: input 2D dataset name
+    # (e.g.)DBv2.0_1-1_3d: DBv2.0_1-1_2d

docs/tools_overview.md

@@ -130,10 +130,10 @@ python -m perception_dataset.convert --config config/convert_deepen_to_t4_sample
 
 ## FastLabel
 
-This step converts FastLabel 2D annotations to T4 format (2D only).
-
 ### Conversion from FastLabel JSON Format to T4 Format
 
+This step converts FastLabel 2D annotations to T4 format (2D only).
+
 input: T4 format data (3D annotated or non-annotated) + FastLabel annotations (JSON format)  
 output: T4 format data (2D annotated)
 
@@ -142,6 +142,17 @@ python -m perception_dataset.convert --config config/convert_fastlabel_2d_to_t4.
 # To overwrite T4-format data, use the --overwrite option
 ```
 
+### Merge 2D T4 format data into 3D T4 format data
+
+This step merges 2D-id-linked T4 format dataset into originally 3D-labeled T4 format data.
+
+input: T4 format data (3D annotated) + T4 format data (2D annotated/ID-linked)  
+output: T4 format data (2D&3D annotated)
+
+```bash
+python -m perception_dataset.convert --config config/merge_2d_t4dataset_to_3d.yaml
+```
+
 ## Rosbag with objects
 
 ### Synthetic bag to T4 format

perception_dataset/convert.py

@@ -342,6 +342,24 @@ def main():
         logger.info(f"[BEGIN] Converting Fastlabel data ({input_base}) to T4 data ({output_base})")
         converter.convert()
         logger.info(f"[END] Converting Fastlabel data ({input_base}) to T4 data ({output_base})")
+
+    elif task == "merge_2d_t4dataset_to_3d":
+        from perception_dataset.t4_dataset.t4_dataset_2d3d_merger import T4dataset2D3DMerger
+
+        input_base = config_dict["conversion"]["input_base"]
+        output_base = config_dict["conversion"]["output_base"]
+        dataset_corresponding = config_dict["conversion"]["dataset_corresponding"]
+
+        converter = T4dataset2D3DMerger(
+            input_base=input_base,
+            output_base=output_base,
+            dataset_corresponding=dataset_corresponding,
+        )
+
+        logger.info(f"[BEGIN] Merging T4 dataset ({input_base}) into T4 dataset ({output_base})")
+        converter.convert()
+        logger.info(f"[Done] Merging T4 dataset ({input_base}) into T4 dataset ({output_base})")
+
     else:
         raise NotImplementedError()
 

perception_dataset/t4_dataset/t4_dataset_2d3d_merger.py

@@ -0,0 +1,154 @@
+import json
+from pathlib import Path
+from typing import Dict
+
+from perception_dataset.abstract_converter import AbstractConverter
+from perception_dataset.utils.logger import configure_logger
+
+logger = configure_logger(modname=__name__)
+
+
+class T4dataset2D3DMerger(AbstractConverter):
+    def __init__(
+        self,
+        input_base: str,
+        output_base: str,
+        dataset_corresponding: Dict[str, int],
+    ):
+        self._input_base = Path(input_base)
+        self._output_base = Path(output_base)
+        self._t4dataset_name_to_merge: Dict[str, str] = dataset_corresponding
+
+    def convert(self):
+        for output_3d_t4dataset_name in self._t4dataset_name_to_merge.keys():
+            logger.info(f"Merge 2D annotation to {output_3d_t4dataset_name}")
+            input_t4dataset_name = self._t4dataset_name_to_merge[output_3d_t4dataset_name]
+            input_2d_annotation_dir = self._input_base / input_t4dataset_name / "annotation"
+            if not input_2d_annotation_dir.exists():
+                input_2d_annotation_dir = (
+                    self._input_base / input_t4dataset_name / "t4_dataset/annotation"
+                )
+            if not input_2d_annotation_dir.exists():
+                logger.warning(f"input_dir {input_2d_annotation_dir} not exists.")
+                continue
+
+            output_3d_annotation_dir = self._output_base / output_3d_t4dataset_name / "annotation"
+            if not output_3d_annotation_dir.exists():
+                logger.warning(f"output_dir {output_3d_annotation_dir} not exists.")
+                continue
+
+            out_attribute, attribute_in_out_token_map = self._merge_json_files(
+                input_2d_annotation_dir, output_3d_annotation_dir, "attribute.json"
+            )
+            out_category, category_in_out_token_map = self._merge_json_files(
+                input_2d_annotation_dir, output_3d_annotation_dir, "category.json"
+            )
+            out_instance, instance_in_out_token_map = self._merge_json_files(
+                input_2d_annotation_dir, output_3d_annotation_dir, "instance.json"
+            )
+            out_visibility, visibility_in_out_token_map = self._merge_json_files(
+                input_2d_annotation_dir, output_3d_annotation_dir, "visibility.json"
+            )
+
+            out_object_ann = self._update_object_ann(
+                input_2d_annotation_dir,
+                attribute_in_out_token_map,
+                category_in_out_token_map,
+                instance_in_out_token_map,
+            )
+            out_surface_ann = self._update_surface_ann(
+                input_2d_annotation_dir, category_in_out_token_map
+            )
+            with open(output_3d_annotation_dir / "attribute.json", "w") as f:
+                json.dump(out_attribute, f, indent=4)
+            with open(output_3d_annotation_dir / "category.json", "w") as f:
+                json.dump(out_category, f, indent=4)
+            with open(output_3d_annotation_dir / "instance.json", "w") as f:
+                json.dump(out_instance, f, indent=4)
+            with open(output_3d_annotation_dir / "visibility.json", "w") as f:
+                json.dump(out_visibility, f, indent=4)
+            with open(output_3d_annotation_dir / "object_ann.json", "w") as f:
+                json.dump(out_object_ann, f, indent=4)
+            with open(output_3d_annotation_dir / "surface_ann.json", "w") as f:
+                json.dump(out_surface_ann, f, indent=4)
+
+    def _merge_json_files(self, input_dir, output_dir, filename):
+        """
+        Merge the input json file to the output json file
+        Args:
+            input_dir: input directory
+            output_dir: output directory
+            filename: json file name
+        return:
+            out_json_data: list of output json data
+            in_out_token_map: mapping of input token to output token for the same name data
+        """
+        with open(input_dir / filename) as f:
+            in_data: list[dict[str, str]] = json.load(f)
+        with open(output_dir / filename) as f:
+            out_data: list[dict[str, str]] = json.load(f)
+
+        in_out_token_map = {}
+        for in_d in in_data:
+            for out_d in out_data:
+                if "name" in in_d.keys():
+                    if in_d["name"] == out_d["name"]:
+                        in_out_token_map[in_d["token"]] = out_d["token"]
+                        break
+                elif "token" in in_d.keys():
+                    if in_d["token"] == out_d["token"]:
+                        in_out_token_map[in_d["token"]] = out_d["token"]
+                        break
+
+        out_data += [d for d in in_data if d["token"] not in in_out_token_map.keys()]
+
+        return out_data, in_out_token_map
+
+    def _update_object_ann(
+        self,
+        input_2d_annotation_dir,
+        attribute_in_out_token_map,
+        category_in_out_token_map,
+        instance_in_out_token_map,
+    ):
+        """
+        Update the attribute token, category token, and instance token in object annotation json file
+        Args:
+            input_2d_annotation_dir: input directory
+            attribute_in_out_token_map: mapping of input attribute token to output attribute token
+            category_in_out_token_map: mapping of input category token to output category token
+            instance_in_out_token_map: mapping of input instance token to output instance token
+        Return:
+            object_ann: list of updated object annotation data
+        """
+        with open(input_2d_annotation_dir / "object_ann.json") as f:
+            object_ann: list[dict[str, str]] = json.load(f)
+
+        for obj in object_ann:
+            for attribute_token in obj["attribute_tokens"]:
+                if attribute_token in attribute_in_out_token_map.keys():
+                    obj["attribute_tokens"].remove(attribute_token)
+                    obj["attribute_tokens"].append(attribute_in_out_token_map[attribute_token])
+            if obj["category_token"] in category_in_out_token_map.keys():
+                obj["category_token"] = category_in_out_token_map[obj["category_token"]]
+            if obj["instance_token"] in instance_in_out_token_map.keys():
+                obj["instance_token"] = instance_in_out_token_map[obj["instance_token"]]
+
+        return object_ann
+
+    def _update_surface_ann(self, input_2d_annotation_dir, category_in_out_token_map):
+        """
+        Update the category token in surface annotation json file
+        Args:
+            input_2d_annotation_dir: input directory
+            category_in_out_token_map: mapping of input category token to output category token
+        Return:
+            surface_ann: list of updated surface annotation data
+        """
+        with open(input_2d_annotation_dir / "surface_ann.json") as f:
+            surface_ann: list[dict[str, str]] = json.load(f)
+
+        for surface in surface_ann:
+            if surface["category_token"] in category_in_out_token_map.keys():
+                surface["category_token"] = category_in_out_token_map[surface["category_token"]]
+        return surface_ann