diff --git a/pyproject.toml b/pyproject.toml index 5bc1f88..3bc9496 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "labelformat" -version = "0.1.15" +version = "0.1.16" authors = [{ name = "Lightly.ai" }] description = "A tool for converting computer vision label formats." readme = "README.md" diff --git a/src/labelformat/formats/coco.py b/src/labelformat/formats/coco.py index fdb9df6..fbde286 100644 --- a/src/labelformat/formats/coco.py +++ b/src/labelformat/formats/coco.py @@ -90,6 +90,7 @@ def get_labels(self) -> Iterable[ImageObjectDetection]: bbox=[float(x) for x in ann["bbox"]], format=BoundingBoxFormat.XYWH, ), + confidence=(float(ann["score"]) if "score" in ann else None), ) ) yield ImageObjectDetection( @@ -173,6 +174,8 @@ def save(self, label_input: ObjectDetectionInput) -> None: float(v) for v in obj.box.to_format(BoundingBoxFormat.XYWH) ], } + if obj.confidence is not None: + annotation["score"] = obj.confidence data["annotations"].append(annotation) self.output_file.parent.mkdir(parents=True, exist_ok=True) diff --git a/tests/unit/formats/test_coco.py b/tests/unit/formats/test_coco.py new file mode 100644 index 0000000..c0805c6 --- /dev/null +++ b/tests/unit/formats/test_coco.py @@ -0,0 +1,101 @@ +import json +from pathlib import Path + +import pytest + +from labelformat.formats.coco import COCOObjectDetectionInput, COCOObjectDetectionOutput +from labelformat.model.bounding_box import BoundingBox +from labelformat.model.category import Category +from labelformat.model.image import Image +from labelformat.model.object_detection import ( + ImageObjectDetection, + SingleObjectDetection, +) + +from ... import simple_object_detection_label_input + + +def _create_coco_file(tmp_path: Path, with_score: bool) -> Path: + annotations = [ + { + "image_id": 0, + "category_id": 1, + "bbox": [10.0, 20.0, 20.0, 20.0], + }, + { + "image_id": 0, + "category_id": 0, + "bbox": [50.0, 60.0, 20.0, 20.0], + }, + ] + if with_score: + annotations[0]["score"] = 0.4 + annotations[1]["score"] = 0.8 + data = { + "images": [ + {"id": 0, "file_name": "image.jpg", "width": 100, "height": 200}, + ], + "categories": [ + {"id": 0, "name": "cat"}, + {"id": 1, "name": "dog"}, + {"id": 2, "name": "cow"}, + ], + "annotations": annotations, + } + coco_file = tmp_path / "train.json" + coco_file.write_text(json.dumps(data)) + return coco_file + + +class TestCOCOObjectDetectionInput: + @pytest.mark.parametrize("with_score", [True, False]) + def test_get_labels(self, tmp_path: Path, with_score: bool) -> None: + coco_file = _create_coco_file(tmp_path=tmp_path, with_score=with_score) + label_input = COCOObjectDetectionInput(input_file=coco_file) + labels = list(label_input.get_labels()) + assert labels == [ + ImageObjectDetection( + image=Image(id=0, filename="image.jpg", width=100, height=200), + objects=[ + SingleObjectDetection( + category=Category(id=1, name="dog"), + box=BoundingBox(xmin=10.0, ymin=20.0, xmax=30.0, ymax=40.0), + confidence=0.4 if with_score else None, + ), + SingleObjectDetection( + category=Category(id=0, name="cat"), + box=BoundingBox(xmin=50.0, ymin=60.0, xmax=70.0, ymax=80.0), + confidence=0.8 if with_score else None, + ), + ], + ), + ] + + +class TestCOCOObjectDetectionOutput: + @pytest.mark.parametrize("with_confidence", [True, False]) + def test_save(self, tmp_path: Path, with_confidence: bool) -> None: + output_file = tmp_path / "train.json" + COCOObjectDetectionOutput(output_file=output_file).save( + label_input=simple_object_detection_label_input.get_input( + with_confidence=with_confidence + ) + ) + + output_json = json.loads(output_file.read_text()) + expected_annotations = [ + { + "image_id": 0, + "category_id": 1, + "bbox": [10.0, 20.0, 20.0, 20.0], + }, + { + "image_id": 0, + "category_id": 0, + "bbox": [50.0, 60.0, 20.0, 20.0], + }, + ] + if with_confidence: + expected_annotations[0]["score"] = 0.4 + expected_annotations[1]["score"] = 0.8 + assert output_json["annotations"] == expected_annotations diff --git a/uv.lock b/uv.lock index 1d0a508..1b84305 100644 --- a/uv.lock +++ b/uv.lock @@ -2706,7 +2706,7 @@ wheels = [ [[package]] name = "labelformat" -version = "0.1.15" +version = "0.1.16" source = { editable = "." } dependencies = [ { name = "fsspec", version = "2024.10.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.9'" },