mmp_wise2526_franksim/mmp/a4/label_grid.py

from typing import Sequence
import numpy as np

from ..a3.annotation import AnnotationRect


def iou(rect1: AnnotationRect, rect2: AnnotationRect) -> float:
    x_left = max(rect1.x1, rect2.x1)
    y_top = max(rect1.y1, rect2.y1)
    x_right = min(rect1.x2, rect2.x2)
    y_bottom = min(rect1.y2, rect2.y2)

    intersection_area = max(x_right - x_left, 0) * max(y_bottom - y_top, 0)

    rect1_area = rect1.area()
    rect2_area = rect2.area()

    union_area = rect1_area + rect2_area - intersection_area

    return intersection_area / union_area


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 gt in gts:
            calculated_iou = iou(item, gt)
            label_grid[width, ratio, row, col] = False
            if calculated_iou >= min_iou:
                label_grid[width, ratio, row, col] = True
                break
    return label_grid