Add first_n argument

pyodi/apps/paint_annotations.py

@@ -44,6 +44,7 @@ def paint_annotations(
     color_key: str = "category_id",
     show_label: bool = True,
     filter_crowd: bool = True,
+    first_n: Optional[int] = None,
 ) -> None:
     """Paint `ground_truth_file` or `predictions_file` annotations on `image_folder` images.
 
@@ -59,6 +60,8 @@ def paint_annotations(
         color_key: Choose the key in annotations on which the color will depend. Defaults to 'category_id'.
         show_label: Choose whether to show label and score threshold on image. Default True.
         filter_crowd: Filter out crowd annotations or not. Default True.
+        first_n: Paint only first n annotations and stop after that.
+            If None, all images will be painted.
     """
     Path(output_folder).mkdir(exist_ok=True, parents=True)
 
@@ -85,8 +88,9 @@ def paint_annotations(
             image_id_to_annotations[annotation["image_id"]].append(annotation)
 
     image_data = ground_truth["images"]
+    first_n = first_n or len(image_data)
 
-    for image in image_data:
+    for image in image_data[:first_n]:
 
         image_filename = image["file_name"]
         image_id = image["id"]

tests/apps/test_paint_annotations.py

@@ -2,6 +2,7 @@
 from pathlib import Path
 
 import numpy as np
+import pytest
 from matplotlib import cm
 from PIL import Image
 
@@ -141,3 +142,40 @@ def test_crowd_annotations_skipped_when_filter_crowd(tmp_path):
 
     assert np.array_equal(result_image_1[3, 4], color_2)
     assert np.array_equal(result_image_1[5, 6], color_2)
+
+
+@pytest.mark.parametrize("first_n", [1, 3, None])
+def test_first_n(tmp_path, first_n):
+
+    images_folder = tmp_path / "images"
+    Path(images_folder).mkdir(exist_ok=True, parents=True)
+
+    output_folder = tmp_path / "test_result"
+    Path(output_folder).mkdir(exist_ok=True, parents=True)
+
+    images, annotations = [], []
+    for i in range(5):
+        images.append(
+            {"id": 0, "width": 10, "height": 10, "file_name": f"test_{i}.png"}
+        )
+        image = np.zeros((10, 10, 3), dtype=np.uint8)
+        Image.fromarray(image).save(images_folder / f"test_{i}.png")
+
+        annotations.append(
+            {"image_id": i, "category_id": 0, "bbox": [0, 0, 2, 2], "score": 1}
+        )
+
+    categories = [{"id": 0, "name": "", "supercategory": "object"}]
+
+    coco_data = dict(images=images, annotations=annotations, categories=categories)
+
+    with open(tmp_path / "test_annotation.json", "w") as json_file:
+        json.dump(coco_data, json_file)
+
+    first_n = first_n or len(images)
+
+    paint_annotations(
+        tmp_path / "test_annotation.json", images_folder, output_folder, first_n=first_n
+    )
+
+    assert len(list(Path(output_folder).iterdir())) == first_n