From 7bc6753ff3b5d064b591258af3399b53364ad027 Mon Sep 17 00:00:00 2001 From: tlpss Date: Thu, 19 Oct 2023 11:54:48 +0200 Subject: [PATCH] separate resize functionality from CLI --- airo-dataset-tools/airo_dataset_tools/cli.py | 25 ++--------- .../coco_tools/transform_dataset.py | 45 ++++++++++++------- .../cvat_labeling/convert_cvat_to_coco.py | 2 +- 3 files changed, 34 insertions(+), 38 deletions(-) diff --git a/airo-dataset-tools/airo_dataset_tools/cli.py b/airo-dataset-tools/airo_dataset_tools/cli.py index 7c37a6e1..7043c83a 100644 --- a/airo-dataset-tools/airo_dataset_tools/cli.py +++ b/airo-dataset-tools/airo_dataset_tools/cli.py @@ -5,13 +5,11 @@ from typing import List, Optional import click -from airo_dataset_tools.coco_tools.albumentations import PillowResize from airo_dataset_tools.coco_tools.coco_instances_to_yolo import create_yolo_dataset_from_coco_instances_dataset from airo_dataset_tools.coco_tools.fiftyone_viewer import view_coco_dataset from airo_dataset_tools.coco_tools.split_dataset import split_and_save_coco_dataset -from airo_dataset_tools.coco_tools.transform_dataset import apply_transform_to_coco_dataset +from airo_dataset_tools.coco_tools.transform_dataset import resize_coco_keypoints_dataset from airo_dataset_tools.cvat_labeling.convert_cvat_to_coco import cvat_image_to_coco -from airo_dataset_tools.data_parsers.coco import CocoKeypointsDataset @click.group() @@ -64,31 +62,14 @@ def convert_cvat_to_coco_cli( @click.argument("annotations-json-path", type=click.Path(exists=True)) @click.option("--width", type=int, required=True) @click.option("--height", type=int, required=True) -def resize_coco_keypoints_dataset(annotations_json_path: str, width: int, height: int) -> None: +def resize_coco_keypoints_dataset_cli(annotations_json_path: str, width: int, height: int) -> None: """Resize a COCO dataset. Will create a new directory with the resized dataset on the same level as the original dataset. Dataset is assumed to be /dir annotations.json # contains relative paths w.r.t. /dir ... """ - coco_dataset_dir = os.path.dirname(annotations_json_path) - annotations_file_name = os.path.basename(annotations_json_path) - dataset_parent_dir = os.path.dirname(coco_dataset_dir) - transformed_dataset_dir = os.path.join( - dataset_parent_dir, f"{annotations_file_name.split('.')[0]}_resized_{width}x{height}" - ) - os.makedirs(transformed_dataset_dir, exist_ok=True) - - transforms = [PillowResize(height, width)] - coco_json = json.load(open(annotations_json_path, "r")) - coco_dataset = CocoKeypointsDataset(**coco_json) - transformed_dataset = apply_transform_to_coco_dataset( - transforms, coco_dataset, coco_dataset_dir, transformed_dataset_dir - ) - - transformed_dataset_dict = transformed_dataset.dict(exclude_none=True) - with open(os.path.join(transformed_dataset_dir, annotations_file_name), "w") as f: - json.dump(transformed_dataset_dict, f) + resize_coco_keypoints_dataset(annotations_json_path, width, height) @cli.command(name="coco-instances-to-yolo") diff --git a/airo-dataset-tools/airo_dataset_tools/coco_tools/transform_dataset.py b/airo-dataset-tools/airo_dataset_tools/coco_tools/transform_dataset.py index 8b121d1d..ca9a363c 100644 --- a/airo-dataset-tools/airo_dataset_tools/coco_tools/transform_dataset.py +++ b/airo-dataset-tools/airo_dataset_tools/coco_tools/transform_dataset.py @@ -4,6 +4,7 @@ import albumentations as A import numpy as np import tqdm +from airo_dataset_tools.coco_tools.albumentations import PillowResize from airo_dataset_tools.data_parsers.coco import ( CocoImage, CocoInstanceAnnotation, @@ -13,6 +14,7 @@ ) from airo_dataset_tools.segmentation_mask_converter import BinarySegmentationMask from PIL import Image +import json def apply_transform_to_coco_dataset( # type: ignore # noqa: C901 @@ -153,26 +155,39 @@ def apply_transform_to_coco_dataset( # type: ignore # noqa: C901 return coco_dataset +def resize_coco_keypoints_dataset(annotations_json_path: str, width: int, height: int) -> None: + """Resize a COCO dataset. Will create a new directory with the resized dataset on the same level as the original dataset. + Dataset is assumed to be + /dir + annotations.json # contains relative paths w.r.t. /dir + ... + """ + coco_dataset_dir = os.path.dirname(annotations_json_path) + annotations_file_name = os.path.basename(annotations_json_path) + dataset_parent_dir = os.path.dirname(coco_dataset_dir) + transformed_dataset_dir = os.path.join( + dataset_parent_dir, f"{annotations_file_name.split('.')[0]}_resized_{width}x{height}" + ) + os.makedirs(transformed_dataset_dir, exist_ok=True) + + transforms = [PillowResize(height, width)] + coco_json = json.load(open(annotations_json_path, "r")) + coco_dataset = CocoKeypointsDataset(**coco_json) + transformed_dataset = apply_transform_to_coco_dataset( + transforms, coco_dataset, coco_dataset_dir, transformed_dataset_dir + ) + + transformed_dataset_dict = transformed_dataset.dict(exclude_none=True) + with open(os.path.join(transformed_dataset_dir, annotations_file_name), "w") as f: + json.dump(transformed_dataset_dict, f) + + if __name__ == "__main__": """example usage of the above function to resize all images in a coco dataset. Copy the following lines into your own codebase and modify as needed.""" - import json import pathlib path = pathlib.Path(__file__).parents[1] / "cvat_labeling" / "example" / "coco.json" coco_json_path = str(path) - coco_dir = os.path.dirname(coco_json_path) - coco_file_name = os.path.basename(coco_json_path) - coco_target_dir = os.path.join(os.path.dirname(coco_dir), "transformed") - os.makedirs(coco_target_dir, exist_ok=True) - - transforms = [A.Resize(128, 128)] - - coco_json = json.load(open(coco_json_path, "r")) - coco_dataset = CocoKeypointsDataset(**coco_json) - transformed_dataset = apply_transform_to_coco_dataset(transforms, coco_dataset, coco_dir, coco_target_dir) - - transformed_dataset_dict = transformed_dataset.dict(exclude_none=True) - with open(os.path.join(coco_target_dir, coco_file_name), "w") as f: - json.dump(transformed_dataset_dict, f) + resize_coco_keypoints_dataset(coco_json_path, 640, 480) diff --git a/airo-dataset-tools/airo_dataset_tools/cvat_labeling/convert_cvat_to_coco.py b/airo-dataset-tools/airo_dataset_tools/cvat_labeling/convert_cvat_to_coco.py index 6f15a795..fae31dd1 100644 --- a/airo-dataset-tools/airo_dataset_tools/cvat_labeling/convert_cvat_to_coco.py +++ b/airo-dataset-tools/airo_dataset_tools/cvat_labeling/convert_cvat_to_coco.py @@ -137,7 +137,7 @@ def _validate_coco_categories_are_in_cvat( if coco_category.name == category_str: break for category_keypoint in coco_category.keypoints: - assert category_keypoint in semantic_types, f"semantic type {category_keypoint.name} not found" + assert category_keypoint in semantic_types, f"semantic type {category_keypoint} not found" def _get_n_category_instances_in_image(cvat_image: ImageItem, category_name: str) -> int: