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, )