From f1db743eb5a64c07e2ad602f6178786059bb68a6 Mon Sep 17 00:00:00 2001 From: Shunsuke Miura <37187849+miursh@users.noreply.github.com> Date: Mon, 13 May 2024 10:57:28 +0900 Subject: [PATCH] feat: merge T4 Datasets with 2D Camera annotations and 3D LiDAR annotations (#96) * add t4 2d-3d merger Signed-off-by: Shunsuke Miura * add missing surface_ann output Signed-off-by: Shunsuke Miura * Update the document, remove unsused codes Signed-off-by: Shunsuke Miura * add error handling, remove unnecessary prints Signed-off-by: Shunsuke Miura * cosmetic change Signed-off-by: Shunsuke Miura * add an output Signed-off-by: Shunsuke Miura * pre-commit Signed-off-by: Shunsuke Miura * Update config/label/attribute.yaml Co-authored-by: kminoda <44218668+kminoda@users.noreply.github.com> --------- Signed-off-by: Shunsuke Miura Co-authored-by: kminoda <44218668+kminoda@users.noreply.github.com> --- config/label/attribute.yaml | 4 +- config/merge_2d_t4dataset_to_3d.yaml | 7 + docs/tools_overview.md | 15 +- perception_dataset/convert.py | 18 ++ .../t4_dataset/t4_dataset_2d3d_merger.py | 154 ++++++++++++++++++ 5 files changed, 195 insertions(+), 3 deletions(-) create mode 100644 config/merge_2d_t4dataset_to_3d.yaml create mode 100644 perception_dataset/t4_dataset/t4_dataset_2d3d_merger.py diff --git a/config/label/attribute.yaml b/config/label/attribute.yaml index 271e8da0..58f7ddb7 100644 --- a/config/label/attribute.yaml +++ b/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: diff --git a/config/merge_2d_t4dataset_to_3d.yaml b/config/merge_2d_t4dataset_to_3d.yaml new file mode 100644 index 00000000..f3379b14 --- /dev/null +++ b/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 diff --git a/docs/tools_overview.md b/docs/tools_overview.md index adde162e..3625e468 100644 --- a/docs/tools_overview.md +++ b/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 diff --git a/perception_dataset/convert.py b/perception_dataset/convert.py index ac97758f..15215abd 100644 --- a/perception_dataset/convert.py +++ b/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() diff --git a/perception_dataset/t4_dataset/t4_dataset_2d3d_merger.py b/perception_dataset/t4_dataset/t4_dataset_2d3d_merger.py new file mode 100644 index 00000000..18eaa4a6 --- /dev/null +++ b/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