98 lines
2.7 KiB
Python
98 lines
2.7 KiB
Python
from typing import Sequence
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw
|
|
import math
|
|
|
|
from ..a3.annotation import AnnotationRect, read_groundtruth_file
|
|
from .anchor_grid import get_anchor_grid
|
|
|
|
|
|
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 idx in np.ndindex(anchor_grid.shape[:-1]):
|
|
for gt in gts:
|
|
calculated_iou = iou(AnnotationRect.fromarray(anchor_grid[idx]), gt)
|
|
label_grid[idx] = False
|
|
if calculated_iou >= min_iou:
|
|
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()
|