pipeline.py

# pipeline.py

from __future__ import annotations

import uuid
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Optional, Sequence, List

import numpy as np
import torch
from PIL import Image

from utils.detector import load_detector, run_detector
from utils.classifier import load_wbc_classifier, classify_wbc_crop
from utils.analysis import CLASS_NAMES, map_age_to_group, pick_gender_for_group
from utils.report import build_api_response

# -------------------------------------------------
# PATHS & DEVICE
# -------------------------------------------------
REPO_ROOT = Path(__file__).resolve().parent

DETECTOR_WEIGHTS = REPO_ROOT / "yolov8_detector" / "best.pt"
CLASSIFIER_WEIGHTS = REPO_ROOT / "wbc_classifier" / "best_model_checkpoint.pth"
REF_CSV = REPO_ROOT / "data" / "WBC differential references.csv"

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")


class SmartCBC:
    """
    Main SmartCBC pipeline orchestrator.

    Usage (single FOV):

        from pipeline import SmartCBC
        cbc = SmartCBC()
        result = cbc.analyze(image, age=32, gender="M")

    Usage (multiple FOVs):

        result = cbc.analyze_batch([img1, img2, img3], age=32, gender="M")

    You can also pass a list directly to `analyze()` and it will auto-route:

        result = cbc.analyze([img1, img2, img3], age=32, gender="M")

    `result` is a dict ready for API / UI usage:
        - patient_id
        - timestamp
        - fovs_analyzed
        - coarse_counts         (RBC/WBC/Platelet)
        - wbc_subtypes          (raw counts per subtype)
        - wbc_percentages       (percent per subtype)
        - report_text           (plain text report)
        - calibration           (placeholders for now)
    """

    def __init__(
        self,
        detector_weights: Optional[str | Path] = None,
        classifier_weights: Optional[str | Path] = None,
        ref_csv: Optional[str | Path] = None,
        conf_thres: float = 0.25,
        imgsz: int = 512,
    ) -> None:
        self.detector_weights = str(detector_weights or DETECTOR_WEIGHTS)
        self.classifier_weights = str(classifier_weights or CLASSIFIER_WEIGHTS)
        self.ref_csv = str(ref_csv or REF_CSV)

        self.conf_thres = conf_thres
        self.imgsz = imgsz
        self.device = DEVICE

        # Load models once
        self.detector = load_detector(self.detector_weights)
        self.classifier = load_wbc_classifier(self.classifier_weights)

    # -------------------------------------------------
    # PUBLIC ENTRYPOINT (single OR multi)
    # -------------------------------------------------
    def analyze(
        self,
        image: Any | Sequence[Any],
        age: Optional[float] = None,
        gender: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Run SmartCBC analysis.

        If `image` is:
          - a single image (PIL / np.ndarray / path) -> analyze one FOV
          - a list/tuple of images                      -> aggregate over multiple FOVs
        """

        # If multiple FOVs are provided, delegate to analyze_batch()
        if isinstance(image, (list, tuple)):
            return self.analyze_batch(list(image), age=age, gender=gender)

        # Single-image path (current behavior)
        pil_img = self._ensure_pil(image)

        age_years, age_group, gender = self._resolve_age_gender(age, gender)

        coarse_counts, subtype_counts = self._run_models_on_image(pil_img)

        ai_result = {
            "patient_id": f"PAT-{uuid.uuid4().hex[:8].upper()}",
            "timestamp": datetime.now().isoformat(timespec="seconds"),
            "fovs_analyzed": 1,
            "coarse_counts": coarse_counts,
            "wbc_subtypes": subtype_counts,
            "calibration": {
                "fov_area_mm2": None,
                "calibration_constant": None,
            },
        }

        response = build_api_response(
            ai_result=ai_result,
            age_group=age_group,
            gender=gender,
            reference_csv=self.ref_csv,
            overlay_image=None,
        )

        response["age_years"] = age_years
        response["age_group"] = age_group
        response["gender"] = gender

        return response

    # -------------------------------------------------
    # NEW: MULTI-FOV ANALYSIS
    # -------------------------------------------------
    def analyze_batch(
        self,
        images: Sequence[Any],
        age: Optional[float] = None,
        gender: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Run SmartCBC analysis over MULTIPLE FOV images.

        Parameters
        ----------
        images : list/tuple of PIL / np.ndarray / path
        age    : float (years), optional
        gender : "M" | "F", optional

        Returns
        -------
        Aggregated result dict:
            - fovs_analyzed = len(images)
            - coarse_counts (sum over all FOVs)
            - wbc_subtypes (sum over all FOVs)
            - wbc_percentages, report_text, etc.
        """

        if not images:
            raise ValueError("analyze_batch() received an empty images list.")

        age_years, age_group, gender = self._resolve_age_gender(age, gender)

        # Initialize aggregate counts
        agg_coarse: Dict[str, int] = {"WBC": 0, "RBC": 0, "Platelet": 0}
        agg_subtypes: Dict[str, int] = {name: 0 for name in CLASS_NAMES}

        fov_count = 0

        for img in images:
            pil_img = self._ensure_pil(img)
            coarse_counts, subtype_counts = self._run_models_on_image(pil_img)

            # Aggregate coarse counts
            for k, v in coarse_counts.items():
                agg_coarse[k] = agg_coarse.get(k, 0) + v

            # Aggregate subtype counts
            for k, v in subtype_counts.items():
                agg_subtypes[k] = agg_subtypes.get(k, 0) + v

            fov_count += 1

        ai_result = {
            "patient_id": f"PAT-{uuid.uuid4().hex[:8].upper()}",
            "timestamp": datetime.now().isoformat(timespec="seconds"),
            "fovs_analyzed": fov_count,
            "coarse_counts": agg_coarse,
            "wbc_subtypes": agg_subtypes,
            "calibration": {
                "fov_area_mm2": None,
                "calibration_constant": None,
            },
        }

        response = build_api_response(
            ai_result=ai_result,
            age_group=age_group,
            gender=gender,
            reference_csv=self.ref_csv,
            overlay_image=None,
        )

        response["age_years"] = age_years
        response["age_group"] = age_group
        response["gender"] = gender

        return response

    # -------------------------------------------------
    # INTERNAL HELPERS
    # -------------------------------------------------
    def _resolve_age_gender(
        self,
        age: Optional[float],
        gender: Optional[str],
    ) -> tuple[float, str, Optional[str]]:
        """
        Compute age_years, age_group, gender with defaults and CSV-based inference.
        """
        age_years = float(age) if age is not None else 30.0
        age_group = map_age_to_group(age_years)

        if gender is None or str(gender).strip() == "":
            gender = pick_gender_for_group(age_group, csv_path=self.ref_csv)
        gender = None if gender is None else str(gender).upper()

        return age_years, age_group, gender

    def _ensure_pil(self, image: Any) -> Image.Image:
        """
        Convert various input types to a PIL.Image in RGB mode.
        """
        if isinstance(image, Image.Image):
            return image.convert("RGB")
        if isinstance(image, np.ndarray):
            if image.ndim == 2:
                image = np.stack([image] * 3, axis=-1)
            return Image.fromarray(image).convert("RGB")
        if isinstance(image, (str, Path)):
            return Image.open(image).convert("RGB")
        raise TypeError(f"Unsupported image type: {type(image)}")

    def _run_models_on_image(
        self,
        img: Image.Image,
    ) -> tuple[Dict[str, int], Dict[str, int]]:
        """
        Run YOLO detector + WBC classifier on a single image.

        Returns
        -------
        coarse_counts : {"WBC": int, "RBC": int, "Platelet": int}
        subtype_counts: {subtype_name: int}
        """
        coarse_counts: Dict[str, int] = {"WBC": 0, "RBC": 0, "Platelet": 0}
        subtype_counts: Dict[str, int] = {name: 0 for name in CLASS_NAMES}

        detections = run_detector(
            model=self.detector,
            image=img,
            imgsz=self.imgsz,
            conf_thres=self.conf_thres,
        )

        w, h = img.size

        for det in detections:
            x1, y1, x2, y2 = det["box"]
            label = det["label"]  # "RBC", "WBC", "Platelet"

            x1 = max(int(x1), 0)
            y1 = max(int(y1), 0)
            x2 = min(int(x2), w)
            y2 = min(int(y2), h)

            if label in coarse_counts:
                coarse_counts[label] += 1
            else:
                coarse_counts[label] = 1

            if label == "WBC" and x2 > x1 and y2 > y1:
                crop = img.crop((x1, y1, x2, y2))
                subtype = classify_wbc_crop(self.classifier, crop)
                if subtype in subtype_counts:
                    subtype_counts[subtype] += 1
                else:
                    subtype_counts[subtype] = 1

        return coarse_counts, subtype_counts