mmp_wise2526_franksim/mmp/a4/anchor_grid.py

from typing import Sequence
import numpy as np


def get_anchor_grid(
    num_rows: int,
    num_cols: int,
    scale_factor: float,
    anchor_widths: Sequence[float],
    aspect_ratios: Sequence[float],
) -> np.ndarray:
    anchor_grid = np.empty(
        [len(anchor_widths), len(aspect_ratios), num_rows, num_cols, 4], dtype=float
    )
    for width_idx, ratio_idx, row, col in np.ndindex(anchor_grid.shape[:-1]):
        anchor_point = (
            col * scale_factor + scale_factor / 2,
            row * scale_factor + scale_factor / 2,
        )
        width = anchor_widths[width_idx]
        ratio = aspect_ratios[ratio_idx]
        anchor_grid[width_idx, ratio_idx, row, col] = get_box(
            width, ratio, anchor_point
        )

    return anchor_grid


def get_box(
    width: float, ratio: float, anchor_point: tuple[float, float]
) -> np.ndarray:
    box = np.empty(4, dtype=float)
    box[0] = anchor_point[0] - (width / 2)
    box[1] = anchor_point[1] - (width * ratio / 2)
    box[2] = anchor_point[0] + (width / 2)
    box[3] = anchor_point[1] + (width * ratio / 2)
    return box