src/localization_quality.py

from __future__ import annotations

from pydantic import BaseModel, Field

from src.schemas.detection import BBox as LocalizationBox

BBOX_NORMALIZED_SCALE = 1000
BBOX_NORMALIZED_AREA = BBOX_NORMALIZED_SCALE * BBOX_NORMALIZED_SCALE


class LocalizationProfile(BaseModel):
    name: str
    max_area_ratio: float | None = None
    max_width_ratio: float | None = None
    max_height_ratio: float | None = None


FOCAL_PROFILE = LocalizationProfile(
    name="focal",
    max_area_ratio=0.12,
    max_width_ratio=0.60,
    max_height_ratio=0.60,
)

DIFFUSE_PROFILE = LocalizationProfile(
    name="diffuse",
    max_area_ratio=None,
    max_width_ratio=None,
    max_height_ratio=None,
)

FOCAL_FINDINGS = frozenset(
    {
        "nodule",
        "pulmonary nodule",
        "nodular opacity",
        "focal opacity",
        "focal airspace opacity",
        "fracture",
        "rib fracture",
    }
)

class LocalizationGateResult(BaseModel):
    accepted_boxes: list[LocalizationBox] = Field(default_factory=list)
    rejection_reasons: list[str] = Field(default_factory=list)


def calculate_bbox_area_ratio(box: LocalizationBox) -> float:
    y0, x0, y1, x1 = box.box_2d
    height = y1 - y0
    width = x1 - x0
    return (height * width) / BBOX_NORMALIZED_AREA


def calculate_bbox_width_ratio(box: LocalizationBox) -> float:
    _, x0, _, x1 = box.box_2d
    return (x1 - x0) / BBOX_NORMALIZED_SCALE


def calculate_bbox_height_ratio(box: LocalizationBox) -> float:
    y0, _, y1, _ = box.box_2d
    return (y1 - y0) / BBOX_NORMALIZED_SCALE


def get_localization_profile(finding_name: str) -> LocalizationProfile:
    normalized = finding_name.strip().lower()
    if normalized in FOCAL_FINDINGS:
        return FOCAL_PROFILE
    return DIFFUSE_PROFILE


def evaluate_localization_boxes(
    finding_name: str,
    boxes: list[LocalizationBox],
) -> LocalizationGateResult:
    profile = get_localization_profile(finding_name)
    accepted_boxes: list[LocalizationBox] = []
    rejection_reasons: list[str] = []

    for box in boxes:
        area_ratio = calculate_bbox_area_ratio(box)
        width_ratio = calculate_bbox_width_ratio(box)
        height_ratio = calculate_bbox_height_ratio(box)
        box_reasons: list[str] = []

        if (
            profile.max_area_ratio is not None
            and area_ratio > profile.max_area_ratio
        ):
            box_reasons.append(
                f"bbox_area_ratio={area_ratio:.4f} exceeds "
                f"{profile.name} max_area_ratio={profile.max_area_ratio:.4f}"
            )

        if (
            profile.max_width_ratio is not None
            and width_ratio > profile.max_width_ratio
        ):
            box_reasons.append(
                f"bbox_width_ratio={width_ratio:.4f} exceeds "
                f"{profile.name} max_width_ratio={profile.max_width_ratio:.4f}"
            )

        if (
            profile.max_height_ratio is not None
            and height_ratio > profile.max_height_ratio
        ):
            box_reasons.append(
                f"bbox_height_ratio={height_ratio:.4f} exceeds "
                f"{profile.name} max_height_ratio={profile.max_height_ratio:.4f}"
            )

        if box_reasons:
            rejection_reasons.extend(box_reasons)
        else:
            accepted_boxes.append(box)

    return LocalizationGateResult(
        accepted_boxes=accepted_boxes,
        rejection_reasons=rejection_reasons,
    )