From e6db77f4b756169d843eac582dd8e43e733e3771 Mon Sep 17 00:00:00 2001
From: Miguel Mendez <miguelmndez@gmail.com>
Date: Fri, 2 Jul 2021 12:08:32 +0200
Subject: [PATCH 1/3] Add first_n argument

---
 pyodi/apps/paint_annotations.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/pyodi/apps/paint_annotations.py b/pyodi/apps/paint_annotations.py
index 42220a5..45f562c 100644
--- a/pyodi/apps/paint_annotations.py
+++ b/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: 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"]

From 611c3440633c9b1c8727a790b7158d0a3f165442 Mon Sep 17 00:00:00 2001
From: Miguel Mendez <miguelmndez@gmail.com>
Date: Fri, 2 Jul 2021 13:10:00 +0200
Subject: [PATCH 2/3] Add first_n test

---
 tests/apps/test_paint_annotations.py | 38 ++++++++++++++++++++++++++++
 1 file changed, 38 insertions(+)

diff --git a/tests/apps/test_paint_annotations.py b/tests/apps/test_paint_annotations.py
index c9e335f..41a4d0a 100644
--- a/tests/apps/test_paint_annotations.py
+++ b/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

From deb9286b3d542ef47afadaaa2753e32460f7151d Mon Sep 17 00:00:00 2001
From: David de la Iglesia Castro <daviddelaiglesiacastro@gmail.com>
Date: Fri, 2 Jul 2021 11:16:13 +0000
Subject: [PATCH 3/3] Update pyodi/apps/paint_annotations.py

---
 pyodi/apps/paint_annotations.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyodi/apps/paint_annotations.py b/pyodi/apps/paint_annotations.py
index 45f562c..3a55b76 100644
--- a/pyodi/apps/paint_annotations.py
+++ b/pyodi/apps/paint_annotations.py
@@ -44,7 +44,7 @@ def paint_annotations(
     color_key: str = "category_id",
     show_label: bool = True,
     filter_crowd: bool = True,
-    first_n: int = None,
+    first_n: Optional[int] = None,
 ) -> None:
     """Paint `ground_truth_file` or `predictions_file` annotations on `image_folder` images.