mmp_wise2526_franksim/mmp/a3/annotation.py

import os
import re
from typing import List
import numpy as np
from PIL import Image, ImageDraw


class AnnotationRect:
    """Exercise 3.1"""

    def __init__(self, x1, y1, x2, y2):
        self.x1 = x1
        self.x2 = x2
        self.y1 = y1
        self.y2 = y2

    def area(self):
        return (self.x2 - self.x1) * (self.y2 - self.y1)

    def __array__(self) -> np.ndarray:
        return np.array([self.x1, self.y1, self.x2, self.y2])

    def scale(self, factor: float):
        self.x1 *= factor
        self.x2 *= factor
        self.y1 *= factor
        self.y2 *= factor
        return self

    @staticmethod
    def fromarray(arr: np.ndarray):
        return AnnotationRect(arr[0], arr[1], arr[2], arr[3])


def read_groundtruth_file(path: str) -> List[AnnotationRect]:
    """Exercise 3.1b"""
    annotationRects = []
    with open(path, "r") as file:
        for line in file:
            if line.strip():
                values = line.strip().split()
                annotationRects.append(
                    AnnotationRect(
                        float(values[0]),
                        float(values[1]),
                        float(values[2]),
                        float(values[3]),
                    )
                )
    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()