assignment-a3: adds annotation visualization and document

mmp/a3/annotation.py

@@ -1,7 +1,9 @@
 import glob
 import os
+import re
 from typing import List
 import numpy as np
+from PIL import Image, ImageDraw
 
 
 class AnnotationRect:
@@ -36,3 +38,44 @@ def read_groundtruth_file(path: str) -> List[AnnotationRect]:
     return annotationRects
 
 
+def get_image_with_max_annotations(dir_path: str) -> str:
+    img_pattern = re.compile(r'^(\d+)\.jpg$')
+    files = set(os.listdir(dir_path))
+    max_file = None
+    max_annotations = 0
+
+    for fname in files:
+        match = img_pattern.match(fname)
+        if match:
+            img_file = os.path.join(dir_path, fname)
+            annotations_number = len(read_groundtruth_file(os.path.join(
+                dir_path, f"{match.group(1)}.gt_data.txt")))
+            if (annotations_number > max_annotations):
+                max_file = img_file
+                max_annotations = annotations_number
+    return max_file
+
+
+def visualize_image(image_path: str, output_path='output.jpg', rect_color=(255, 0, 0), width=2):
+    img_pattern = re.compile(r'(.*)(\.jpg)')
+    match = img_pattern.match(image_path)
+    annotations = read_groundtruth_file(f"{match.group(1)}.gt_data.txt")
+
+    img = Image.open(image_path).convert('RGB')
+    draw = ImageDraw.Draw(img)
+
+    for annotation in annotations:
+        draw.rectangle([annotation.x1, annotation.y1, annotation.x2, annotation.y2],
+                       outline=rect_color, width=width)
+
+    img.save(output_path)
+
+
+def main():
+    image_file = get_image_with_max_annotations(
+        "/home/ubuntu/mmp_wise2526_franksim/.data/mmp-public-3.2/train")
+    visualize_image(image_file)
+
+
+if __name__ == "__main__":
+    main()