save_coco_annotations logic simplification
This commit is contained in:
SkalskiP committed
1 parent
4d82460358
commit
f63336a824
2 files changed
+117
-39
No files matched your search
@@ -31,24 +31,16 @@ def classes_to_coco_categories(classes: List[str]) -> List[dict]:
|
||||
]
|
||||
|
||||
|
||||
def _extract_image_info(annotation_data: dict) -> List[dict]:
|
||||
image_infos = annotation_data.get("images", None)
|
||||
return image_infos
|
||||
|
||||
|
||||
def _annotations_dict(annotation_data: dict) -> Tuple[np.ndarray, dict]:
|
||||
annotations_infos = annotation_data.get("annotations", None)
|
||||
image_id_label_id_pair = np.zeros((0, 2))
|
||||
def group_coco_annotations_by_image_id(
|
||||
coco_annotations: List[dict],
|
||||
) -> Dict[int, List[dict]]:
|
||||
annotations = {}
|
||||
for anno in annotations_infos:
|
||||
_annotations = {}
|
||||
for key, item in anno.items():
|
||||
_annotations[key] = anno[key]
|
||||
id_pair = np.array([[anno["image_id"], anno["id"]]])
|
||||
image_id_label_id_pair = np.append(image_id_label_id_pair, id_pair, axis=0)
|
||||
annotations[anno["id"]] = _annotations
|
||||
|
||||
return image_id_label_id_pair, annotations
|
||||
for annotation in coco_annotations:
|
||||
image_id = annotation["image_id"]
|
||||
if image_id not in annotations:
|
||||
annotations[image_id] = []
|
||||
annotations[image_id].append(annotation)
|
||||
return annotations
|
||||
|
||||
|
||||
def _polygons_to_masks(
|
||||
@@ -147,36 +139,29 @@ def load_coco_annotations(
|
||||
Returns:
|
||||
Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]: A tuple containing a list of class names, a dictionary with image names as keys and images as values, and a dictionary with image names as keys and corresponding Detections instances as values.
|
||||
"""
|
||||
annotation_data = read_json_file(file_path=annotations_path)
|
||||
|
||||
classes = coco_categories_to_classes(coco_categories=annotation_data["categories"])
|
||||
coco_images = annotation_data["images"]
|
||||
image_id_label_id_pair, annotation_dict = _annotations_dict(
|
||||
annotation_data=annotation_data
|
||||
coco_data = read_json_file(file_path=annotations_path)
|
||||
classes = coco_categories_to_classes(coco_categories=coco_data["categories"])
|
||||
coco_images = coco_data["images"]
|
||||
coco_annotations_groups = group_coco_annotations_by_image_id(
|
||||
coco_annotations=coco_data["annotations"]
|
||||
)
|
||||
|
||||
images = {}
|
||||
annotations = {}
|
||||
|
||||
for coco_image in coco_images:
|
||||
image_name = coco_image["file_name"]
|
||||
image_name, image_width, image_height = (
|
||||
coco_image["file_name"],
|
||||
coco_image["width"],
|
||||
coco_image["height"],
|
||||
)
|
||||
image_annotations = coco_annotations_groups.get(coco_image["id"])
|
||||
image_path = os.path.join(images_directory_path, image_name)
|
||||
|
||||
image = cv2.imread(str(image_path))
|
||||
|
||||
# Filter annotations based on image id
|
||||
per_image_label_ids = image_id_label_id_pair[
|
||||
image_id_label_id_pair[:, 0] == coco_image["id"]
|
||||
][:, 1]
|
||||
|
||||
image_annotations = []
|
||||
for per_image_label_id in per_image_label_ids:
|
||||
image_annotations.append(annotation_dict[int(per_image_label_id)])
|
||||
|
||||
w, h = coco_image["width"], coco_image["height"]
|
||||
|
||||
annotation = coco_annotations_to_detections(
|
||||
image_annotations=image_annotations,
|
||||
resolution_wh=(w, h),
|
||||
resolution_wh=(image_width, image_height),
|
||||
with_masks=force_masks,
|
||||
)
|
||||
|
||||
|
||||
@@ -1,9 +1,27 @@
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
from typing import List
|
||||
from typing import List, Tuple
|
||||
|
||||
import pytest
|
||||
|
||||
from supervision.dataset.formats.coco import classes_to_coco_categories, coco_categories_to_classes
|
||||
from supervision.dataset.formats.coco import classes_to_coco_categories, coco_categories_to_classes, \
|
||||
group_coco_annotations_by_image_id
|
||||
|
||||
|
||||
def generate_cock_coco_annotation(
|
||||
annotation_id: int = 0,
|
||||
image_id: int = 0,
|
||||
category_id: int = 0,
|
||||
bbox: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0),
|
||||
area: float = 0.0
|
||||
) -> dict:
|
||||
return {
|
||||
"id": annotation_id,
|
||||
"image_id": image_id,
|
||||
"category_id": category_id,
|
||||
"bbox": list(bbox),
|
||||
"area": area,
|
||||
"iscrowd": 0
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -130,3 +148,78 @@ def test_classes_to_coco_categories_and_back_to_classes(classes: List[str], exce
|
||||
coco_categories = classes_to_coco_categories(classes=classes)
|
||||
result = coco_categories_to_classes(coco_categories=coco_categories)
|
||||
assert result == classes
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"coco_annotations, expected_result, exception",
|
||||
[
|
||||
(
|
||||
[],
|
||||
{},
|
||||
DoesNotRaise()
|
||||
), # empty coco annotations
|
||||
(
|
||||
[
|
||||
generate_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0)
|
||||
],
|
||||
{
|
||||
0: [
|
||||
generate_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0)
|
||||
]
|
||||
},
|
||||
DoesNotRaise()
|
||||
), # single coco annotation
|
||||
(
|
||||
[
|
||||
generate_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0),
|
||||
generate_cock_coco_annotation(annotation_id=1, image_id=1, category_id=0)
|
||||
],
|
||||
{
|
||||
0: [
|
||||
generate_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0)
|
||||
],
|
||||
1: [
|
||||
generate_cock_coco_annotation(annotation_id=1, image_id=1, category_id=0)
|
||||
]
|
||||
},
|
||||
DoesNotRaise()
|
||||
), # two coco annotations
|
||||
(
|
||||
[
|
||||
generate_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0),
|
||||
generate_cock_coco_annotation(annotation_id=1, image_id=1, category_id=1),
|
||||
generate_cock_coco_annotation(annotation_id=2, image_id=1, category_id=2),
|
||||
generate_cock_coco_annotation(annotation_id=3, image_id=2, category_id=3),
|
||||
generate_cock_coco_annotation(annotation_id=4, image_id=3, category_id=1),
|
||||
generate_cock_coco_annotation(annotation_id=5, image_id=3, category_id=2),
|
||||
generate_cock_coco_annotation(annotation_id=5, image_id=3, category_id=3),
|
||||
],
|
||||
{
|
||||
0: [
|
||||
generate_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0),
|
||||
],
|
||||
1: [
|
||||
generate_cock_coco_annotation(annotation_id=1, image_id=1, category_id=1),
|
||||
generate_cock_coco_annotation(annotation_id=2, image_id=1, category_id=2),
|
||||
],
|
||||
2: [
|
||||
generate_cock_coco_annotation(annotation_id=3, image_id=2, category_id=3),
|
||||
],
|
||||
3: [
|
||||
generate_cock_coco_annotation(annotation_id=4, image_id=3, category_id=1),
|
||||
generate_cock_coco_annotation(annotation_id=5, image_id=3, category_id=2),
|
||||
generate_cock_coco_annotation(annotation_id=5, image_id=3, category_id=3),
|
||||
]
|
||||
},
|
||||
DoesNotRaise()
|
||||
), # two coco annotations
|
||||
]
|
||||
)
|
||||
def test_group_coco_annotations_by_image_id(
|
||||
coco_annotations: List[dict],
|
||||
expected_result: dict,
|
||||
exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = group_coco_annotations_by_image_id(coco_annotations=coco_annotations)
|
||||
assert result == expected_result
|
||||
Reference in new issue
Block a user