src/schemas/detection.py

from __future__ import annotations

from typing import Literal

from pydantic import BaseModel, Field, field_validator, model_validator


Certainty = Literal["positive", "probable", "questionable"]

LocalizationStatus = Literal[
    "localized",
    "abstained",
    "rejected_by_quality_gate",
    "parser_error",
]


class BBox(BaseModel):
    """Bounding box in normalized [y0, x0, y1, x1] space on [0, 1000]."""

    box_2d: list[int]
    label: str = Field(min_length=1)

    @field_validator("box_2d")
    @classmethod
    def validate_box(cls, value: list[int]) -> list[int]:
        if len(value) != 4:
            raise ValueError("box_2d must contain exactly four values: [y0, x0, y1, x1].")
        if any(c < 0 or c > 1000 for c in value):
            raise ValueError("box_2d coordinates must be normalized to [0, 1000].")
        return value

    @model_validator(mode="after")
    def validate_order(self) -> "BBox":
        y0, x0, y1, x1 = self.box_2d
        if not (x0 < x1 and y0 < y1):
            raise ValueError("box_2d must satisfy x0 < x1 and y0 < y1.")
        return self


class Finding(BaseModel):
    finding: str = Field(min_length=1)
    anatomical_location: str = Field(min_length=1)
    certainty: Certainty


class FindingDiscovery(BaseModel):
    findings: list[Finding] = Field(default_factory=list)


class ViewLocalization(BaseModel):
    image_index: int
    image_path: str
    status: LocalizationStatus
    boxes: list[BBox] = Field(default_factory=list)
    candidate_boxes: list[BBox] = Field(default_factory=list)
    rejection_reasons: list[str] = Field(default_factory=list)


class LocalizedFinding(BaseModel):
    finding: str
    anatomical_location: str
    certainty: Certainty
    localizations: list[ViewLocalization] = Field(default_factory=list)


class DetectionResult(BaseModel):
    case_id: str | None = None
    input_images: list[str]
    findings: list[LocalizedFinding] = Field(default_factory=list)