separate resize functionality from CLI

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")

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)

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: