visualize label grid

mmp/a4/label_grid.py

@@ -1,7 +1,10 @@
 from typing import Sequence
 import numpy as np
+from PIL import Image, ImageDraw
+import math
 
-from ..a3.annotation import AnnotationRect
+from ..a3.annotation import AnnotationRect, read_groundtruth_file
+from .anchor_grid import get_anchor_grid
 
 
 def iou(rect1: AnnotationRect, rect2: AnnotationRect) -> float:
@@ -24,11 +27,71 @@ def get_label_grid(
     anchor_grid: np.ndarray, gts: Sequence[AnnotationRect], min_iou: float
 ) -> tuple[np.ndarray, ...]:
     label_grid = np.empty(anchor_grid.shape[:-1], dtype=bool)
-    for (width, ratio, row, col), item in np.ndenumerate(anchor_grid):
+    for idx in np.ndindex(anchor_grid.shape[:-1]):
         for gt in gts:
-            calculated_iou = iou(item, gt)
-            label_grid[width, ratio, row, col] = False
+            calculated_iou = iou(AnnotationRect.fromarray(anchor_grid[idx]), gt)
+            label_grid[idx] = False
             if calculated_iou >= min_iou:
-                label_grid[width, ratio, row, col] = True
+                label_grid[idx] = True
                 break
     return label_grid
+
+
+def calculate_label_grid_for_image(
+    image: str,
+    scale_factor: float,
+    anchor_widths: Sequence[float],
+    aspect_ratios: Sequence[float],
+    min_iou: float,
+):
+    im_width, im_height = Image.open(image).size
+
+    anchor_grid = get_anchor_grid(
+        num_rows=math.floor(im_height / scale_factor),
+        num_cols=math.floor(im_width / scale_factor),
+        scale_factor=scale_factor,
+        anchor_widths=anchor_widths,
+        aspect_ratios=aspect_ratios,
+    )
+    gts = read_groundtruth_file(".data/mmp-public-3.2/train/02242500.gt_data.txt")
+    label_grid = get_label_grid(anchor_grid=anchor_grid, gts=gts, min_iou=min_iou)
+    annotations = [
+        AnnotationRect.fromarray(anchor_grid[idx])
+        for idx in np.ndindex(anchor_grid.shape[:-1])
+        if label_grid[idx]
+    ]
+    draw_annotation_rects(image, annotations, output_path="output2.jpg")
+
+
+def draw_annotation_rects(
+    image: str,
+    annotations: Sequence[AnnotationRect],
+    rect_color=(255, 0, 0),
+    rect_width=2,
+    output_path="output.jpg",
+):
+    img = Image.open(image).convert("RGB")
+    draw = ImageDraw.Draw(img)
+
+    for annotation in annotations:
+        draw.rectangle(
+            [annotation.x1, annotation.y1, annotation.x2, annotation.y2],
+            outline=rect_color,
+            width=rect_width,
+        )
+
+    img.save(output_path)
+
+
+def main():
+    calculate_label_grid_for_image(
+        ".data/mmp-public-3.2/train/02242500.jpg",
+        8,
+        anchor_widths=[16, 32, 64, 96, 128, 144, 150],
+        aspect_ratios=[1, 4 / 3, 5 / 3, 2, 2.5, 3],
+        min_iou=0.70,
+    )
+
+
+if __name__ == "__main__":
+    main()