xiaoxue_admin
/
supervision
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

✅ ready for test and review

This commit is contained in:
SkalskiP 2023-07-04 17:31:33 +02:00
parent 3e4cb6df42
commit c4b33fe1d6
2 changed files with 157 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
supervision/dataset/formats/coco.py
View File

@ -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

136
test/dataset/formats/test_coco.py
View File

@ -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