src/pipeline.py

from __future__ import annotations

from collections.abc import Sequence

from PIL import Image

from src.interfaces.detector import Detector
from src.interfaces.reporter import ReportGenerator
from src.interfaces.segmenter import Segmenter
from src.schemas.detection import DetectionResult
from src.schemas.report import PipelineOutput, ReportRequest
from src.schemas.segmentation import MaskResult


class CXRPipeline:
    """Three-stage chest X-ray analysis pipeline.

    Stage 1 (required) – Detection & localization via Detector.
    Stage 2 (optional) – Segmentation via Segmenter.
    Stage 3 (optional) – Report generation via ReportGenerator.

    Pass segmenter=None or reporter=None to skip those stages.
    """

    def __init__(
        self,
        detector: Detector,
        segmenter: Segmenter | None = None,
        reporter: ReportGenerator | None = None,
    ) -> None:
        self.detector = detector
        self.segmenter = segmenter
        self.reporter = reporter

    # -- Stage 2 helper ------------------------------------------------------

    def _run_segmentation(
        self,
        detection: DetectionResult,
        processed_images: list[Image.Image],
    ) -> list[MaskResult]:
        assert self.segmenter is not None
        masks: list[MaskResult] = []

        for localized_finding in detection.findings:
            for view in localized_finding.localizations:
                if view.status != "localized" or not view.boxes:
                    continue
                image = processed_images[view.image_index]
                # Use the first accepted box per view (the model returns one box normally)
                bbox = view.boxes[0].box_2d
                mask_result = self.segmenter.segment(
                    image=image,
                    bbox_normalized=bbox,
                    finding_label=localized_finding.finding,
                    image_path=view.image_path,
                    image_index=view.image_index,
                )
                masks.append(mask_result)

                status_str = "ok" if mask_result.status == "success" else f"FAILED: {mask_result.error}"
                print(
                    f"[Stage 2] {localized_finding.finding} | "
                    f"view {view.image_index} | {status_str}"
                )

        return masks

    # -- public entry-point --------------------------------------------------

    def run(
        self,
        images: Sequence[Image.Image],
        input_images: list[str],
        case_id: str | None = None,
    ) -> PipelineOutput:
        # ── Stage 1 ──────────────────────────────────────────────────────────
        print("\n" + "=" * 60)
        print("STAGE 1 – Detection & Localization")
        print("=" * 60)
        detection, processed_images = self.detector.detect(
            images, input_images, case_id=case_id
        )

        # ── Stage 2 ──────────────────────────────────────────────────────────
        masks: list[MaskResult] = []
        if self.segmenter is not None:
            print("\n" + "=" * 60)
            print("STAGE 2 – Segmentation")
            print("=" * 60)
            masks = self._run_segmentation(detection, processed_images)

        # ── Stage 3 ──────────────────────────────────────────────────────────
        report = None
        if self.reporter is not None:
            print("\n" + "=" * 60)
            print("STAGE 3 – Report Generation")
            print("=" * 60)
            report_request = ReportRequest(
                case_id=case_id,
                input_images=input_images,
                images=list(processed_images),
                findings=detection.findings,
                masks=masks,
            )
            report = self.reporter.generate_report(report_request)

        return PipelineOutput(
            detection=detection,
            masks=masks,
            report=report,
            processed_images=processed_images,
        )