From c4b33fe1d6ecc201bfd9942dbc0910ccd573a300 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Tue, 4 Jul 2023 17:31:33 +0200 Subject: [PATCH 1/2] =?UTF-8?q?=E2=9C=85=20ready=20for=20test=20and=20revi?= =?UTF-8?q?ew?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/dataset/formats/coco.py | 24 ++++- test/dataset/formats/test_coco.py | 136 +++++++++++++++++++++++++++- 2 files changed, 157 insertions(+), 3 deletions(-) diff --git a/supervision/dataset/formats/coco.py b/supervision/dataset/formats/coco.py index 51adff01..c1d0a77a 100644 --- a/supervision/dataset/formats/coco.py +++ b/supervision/dataset/formats/coco.py @@ -6,7 +6,10 @@ from typing import Dict, List, Tuple import cv2 import numpy as np -from supervision.dataset.utils import approximate_mask_with_polygons +from supervision.dataset.utils import ( + approximate_mask_with_polygons, + map_detections_class_id, +) from supervision.detection.core import Detections from supervision.detection.utils import polygon_to_mask from supervision.utils.file import read_json_file, save_json_file @@ -19,6 +22,18 @@ def coco_categories_to_classes(coco_categories: List[dict]) -> List[str]: ] +def build_coco_class_index_mapping( + coco_categories: List[dict], target_classes: List[str] +) -> Dict[int, int]: + source_class_to_index = { + category["name"]: category["id"] for category in coco_categories + } + return { + source_class_to_index[target_class_name]: target_class_index + for target_class_index, target_class_name in enumerate(target_classes) + } + + def classes_to_coco_categories(classes: List[str]) -> List[dict]: return [ { @@ -124,6 +139,9 @@ def load_coco_annotations( ) -> Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]: coco_data = read_json_file(file_path=annotations_path) classes = coco_categories_to_classes(coco_categories=coco_data["categories"]) + class_index_mapping = build_coco_class_index_mapping( + coco_categories=coco_data["categories"], target_classes=classes + ) coco_images = coco_data["images"] coco_annotations_groups = group_coco_annotations_by_image_id( coco_annotations=coco_data["annotations"] @@ -147,6 +165,10 @@ def load_coco_annotations( resolution_wh=(image_width, image_height), with_masks=force_masks, ) + annotation = map_detections_class_id( + source_to_target_mapping=class_index_mapping, + detections=annotation, + ) images[image_name] = image annotations[image_name] = annotation diff --git a/test/dataset/formats/test_coco.py b/test/dataset/formats/test_coco.py index ca1ceb6c..cbbb2e21 100644 --- a/test/dataset/formats/test_coco.py +++ b/test/dataset/formats/test_coco.py @@ -1,11 +1,11 @@ from contextlib import ExitStack as DoesNotRaise -from typing import List, Tuple +from typing import List, Tuple, Dict import pytest from supervision import Detections from supervision.dataset.formats.coco import classes_to_coco_categories, coco_categories_to_classes, \ - group_coco_annotations_by_image_id, coco_annotations_to_detections + group_coco_annotations_by_image_id, coco_annotations_to_detections, build_coco_class_index_mapping import numpy as np @@ -290,3 +290,135 @@ def test_coco_annotations_to_detections( with_masks=with_masks ) assert result == expected_result + +@pytest.mark.parametrize( + "coco_categories, target_classes, expected_result, exception", + [ + ( + [], + [], + {}, + DoesNotRaise() + ), # empty coco categories + ( + [ + { + "id": 0, + "name": "fashion-assistant", + "supercategory": "none" + } + ], + [ + "fashion-assistant" + ], + { + 0: 0 + }, + DoesNotRaise() + ), # single coco category starting from 0 + ( + [ + { + "id": 1, + "name": "fashion-assistant", + "supercategory": "none" + } + ], + [ + "fashion-assistant" + ], + { + 1: 0 + }, + DoesNotRaise() + ), # single coco category starting from 1 + ( + [ + { + "id": 0, + "name": "fashion-assistant", + "supercategory": "none" + }, + { + "id": 2, + "name": "hoodie", + "supercategory": "fashion-assistant" + }, + { + "id": 1, + "name": "baseball cap", + "supercategory": "fashion-assistant" + } + ], + [ + "fashion-assistant", + "baseball cap", + "hoodie" + ], + { + 0: 0, + 1: 1, + 2: 2 + }, + DoesNotRaise() + ), # three coco categories + ( + [ + { + "id": 2, + "name": "hoodie", + "supercategory": "fashion-assistant" + }, + { + "id": 1, + "name": "baseball cap", + "supercategory": "fashion-assistant" + } + ], + [ + "baseball cap", + "hoodie" + ], + { + 2: 1, + 1: 0 + }, + DoesNotRaise() + ), # two coco categories + ( + [ + { + "id": 3, + "name": "hoodie", + "supercategory": "fashion-assistant" + }, + { + "id": 1, + "name": "baseball cap", + "supercategory": "fashion-assistant" + } + ], + [ + "baseball cap", + "hoodie" + ], + { + 3: 1, + 1: 0 + }, + DoesNotRaise() + ), # two coco categories with missing category + ] +) +def test_build_coco_class_index_mapping( + coco_categories: List[dict], + target_classes: List[str], + expected_result: Dict[int, int], + exception: Exception +) -> None: + with exception: + result = build_coco_class_index_mapping( + coco_categories=coco_categories, + target_classes=target_classes + ) + assert result == expected_result From 4f5f87acb858c7b7d250badd1edbbd2752b59815 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Wed, 5 Jul 2023 11:37:44 +0200 Subject: [PATCH 2/2] =?UTF-8?q?1=EF=B8=8F=E2=83=A3=20changing=20logic=20of?= =?UTF-8?q?=20`to=5Fcoco`=20to=20start=20`image=5Fid`=20and=20`annotation?= =?UTF-8?q?=5Fid`=20from=201?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/dataset/formats/coco.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/supervision/dataset/formats/coco.py b/supervision/dataset/formats/coco.py index c1d0a77a..a6066ead 100644 --- a/supervision/dataset/formats/coco.py +++ b/supervision/dataset/formats/coco.py @@ -199,8 +199,7 @@ def save_coco_annotations( coco_images = [] coco_categories = classes_to_coco_categories(classes=classes) - image_id = 0 - annotation_id = 0 + image_id, annotation_id = 1, 1 for image_name, image in images.items(): image_height, image_width, _ = image.shape