35 lines
1.0 KiB
Python
35 lines
1.0 KiB
Python
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:
|
|
iou = iou(item, gt)
|
|
label_grid[width, ratio, row, col] = False
|
|
if iou >= min_iou:
|
|
label_grid[width, ratio, row, col] = True
|
|
break
|
|
return label_grid
|