diff --git a/pyodi/apps/coco/coco_merge.py b/pyodi/apps/coco/coco_merge.py index 8c2d46e..d62c07c 100644 --- a/pyodi/apps/coco/coco_merge.py +++ b/pyodi/apps/coco/coco_merge.py @@ -16,19 +16,22 @@ # API REFERENCE """ # noqa: E501 import json -from typing import Any, Dict +from typing import Any, Dict, Optional from loguru import logger @logger.catch(reraise=True) -def coco_merge(input_extend: str, input_add: str, output_file: str,) -> str: +def coco_merge( + input_extend: str, input_add: str, output_file: str, indent: Optional[int] = None, +) -> str: """Merge COCO annotation files. Args: input_extend: Path to input file to be extended. input_add: Path to input file to be added. output_file : Path to output file with merged annotations. + indent: Argument passed to `json.dump`. See https://docs.python.org/3/library/json.html#json.dump. """ with open(input_extend, "r") as f: data_extend = json.load(f) @@ -88,6 +91,6 @@ def coco_merge(input_extend: str, input_add: str, output_file: str,) -> str: ) with open(output_file, "w") as f: - json.dump(output, f, indent=2) + json.dump(output, f, indent=indent) return output_file diff --git a/setup.cfg b/setup.cfg index 38619af..c6fe274 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,6 +1,6 @@ [metadata] name = pyodi -version = 0.0.6 +version = 0.0.7 author = Pyodi description = Object Detection Insights long_description = file: README.md @@ -38,6 +38,7 @@ dev = mkdocs==1.1.2 mkdocstrings==0.14.0 mkdocs-material==6.2.8 + mock==4.0.3 mypy==0.800 pre-commit==2.10.1 pydocstyle==5.1.1 diff --git a/tests/apps/test_coco_merge.py b/tests/apps/test_coco_merge.py index 1fb4b15..2e6b352 100644 --- a/tests/apps/test_coco_merge.py +++ b/tests/apps/test_coco_merge.py @@ -1,5 +1,8 @@ import json +import pytest +from mock import ANY, patch + from pyodi.apps.coco.coco_merge import coco_merge @@ -56,3 +59,29 @@ def test_coco_merge(tmp_path): range(len(data["annotations"])) ) assert [i["category_id"] for i in data["annotations"]] == [1, 2, 1, 3, 3, 1] + + +@pytest.mark.parametrize("indent", [None, 2]) +def test_coco_merge_with_json_indent(tmp_path, indent): + images = [{"id": 0, "file_name": "0.jpg"}] + anns1 = [{"image_id": 0, "category_id": 0, "id": 0}] + anns2 = [{"image_id": 0, "category_id": 1, "id": 0}] + categories = [{"id": 0, "name": "excavator"}, {"id": 1, "name": "bus"}] + + coco1 = dict(images=images, annotations=anns1, categories=categories) + coco2 = dict(images=images, annotations=anns2, categories=categories) + + tmp_files = [] + for i, coco_data in enumerate([coco1, coco2]): + tmp_files.append(tmp_path / f"{i}.json") + with open(tmp_files[-1], "w") as f: + json.dump(coco_data, f) + + with patch("json.dump") as mock: + coco_merge( + input_extend=tmp_path / "0.json", + input_add=tmp_path / "1.json", + output_file=tmp_path / "result.json", + indent=indent, + ) + mock.assert_called_once_with(ANY, ANY, indent=indent)