diff --git a/mmp/a4/label_grid.py b/mmp/a4/label_grid.py index ebc3d12..d3644b8 100644 --- a/mmp/a4/label_grid.py +++ b/mmp/a4/label_grid.py @@ -1,7 +1,10 @@ from typing import Sequence import numpy as np +from PIL import Image, ImageDraw +import math -from ..a3.annotation import AnnotationRect +from ..a3.annotation import AnnotationRect, read_groundtruth_file +from .anchor_grid import get_anchor_grid def iou(rect1: AnnotationRect, rect2: AnnotationRect) -> float: @@ -24,11 +27,71 @@ 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 idx in np.ndindex(anchor_grid.shape[:-1]): for gt in gts: - calculated_iou = iou(item, gt) - label_grid[width, ratio, row, col] = False + calculated_iou = iou(AnnotationRect.fromarray(anchor_grid[idx]), gt) + label_grid[idx] = False if calculated_iou >= min_iou: - label_grid[width, ratio, row, col] = True + 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() diff --git a/output.jpg b/output.jpg new file mode 100644 index 0000000..6dc66a9 Binary files /dev/null and b/output.jpg differ diff --git a/output2.jpg b/output2.jpg new file mode 100644 index 0000000..003b8d6 Binary files /dev/null and b/output2.jpg differ