hypertower/scripts/legacy/eval_papila_classifier.py

"""Evaluate REFUGE-trained classifier on Papila images using UNet crops."""

from __future__ import annotations

import argparse
import csv
from pathlib import Path
from typing import Dict, List, Optional, Sequence, Set

import numpy as np
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm
from PIL import Image, ImageDraw

import sys

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from classes.refuge_preprocessing import RefugePreprocessing, RefugeSample
from classes.refuge_segmentation import RefugeSegmentation
from classes.refuge_classification import (
    RefugeClassification,
    RefugeClassificationDataset,
    RefugeClassificationRecord,
    UNetGeometryProvider,
    _default_image_transform,
    _geometry_from_mask,
)
from classes.backbones import BACKBONES, load_backbone_weights
from classes.unet_segmenter import UNetSegmenter
from classes.papila_builders import build_papila_clinical


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Evaluate classifier on Papila with UNet crops")
    parser.add_argument("--filtered-metrics", type=Path, required=True, help="CSV of Papila samples with acceptable Dice")
    parser.add_argument("--segmenter-manifest", type=Path, required=True, help="Manifest used to train the UNet segmenter")
    parser.add_argument("--segmenter-weights", type=Path, required=True, help="Path to trained UNet weights (best.pt)")
    parser.add_argument("--classifier-weights", type=Path, required=False, help="Path to classifier checkpoint (refuge_classifier_best.pt)")
    parser.add_argument("--refuge-root", type=Path, default=Path("REFUGE"))
    parser.add_argument("--image-dir", type=Path, default=Path("Papila/FundusImages"))
    parser.add_argument("--clinical-dir", type=Path, default=Path("Papila/ClinicalData"))
    parser.add_argument("--label-col", type=str, default="Diagnosis", help="Column name holding Papila labels")
    parser.add_argument(
        "--positive-labels",
        nargs="*",
        default=["glaucoma", "glaucoma suspect", "suspect"],
        help="Values treated as glaucoma-positive when labels are non-numeric",
    )
    parser.add_argument("--dice-threshold", type=float, default=0.01, help="Minimum Dice (disc or cup) to keep a sample")
    parser.add_argument("--segmenter-threshold", type=float, default=0.5, help="Probability threshold for UNet geometry")
    parser.add_argument("--segmenter-normalize", choices=["none", "imagenet", "per_image"], default="per_image")
    parser.add_argument("--segmenter-tta", action="store_true", help="Enable TTA (H/V flips) when deriving geometry")
    parser.add_argument("--crop-scale", type=float, default=2.5)
    parser.add_argument("--crop-size", type=int, default=224)
    parser.add_argument("--batch-size", type=int, default=32)
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    parser.add_argument("--output", type=Path, default=None, help="Optional CSV to store per-sample probabilities")
    parser.add_argument(
        "--cache-dir",
        type=Path,
        default=Path("analysis_data/classifier_cache"),
        help="Directory to reuse classifier preprocessing cache",
    )
    parser.add_argument(
        "--use-gt-masks",
        action="store_true",
        help="Use ground truth Papila contours instead of UNet predictions",
    )
    parser.add_argument(
        "--gt-contours-dir",
        type=Path,
        default=Path("Papila/ExpertsSegmentations/Contours"),
        help="Directory containing Papila contour text files",
    )
    parser.add_argument(
        "--backbone",
        type=str,
        default=None,
        help="Optional backbone name (e.g. inception_v3, densenet121). Requires matching classifier weights.",
    )
    return parser.parse_args()


def load_allowed_ids(path: Path, dice_threshold: float) -> Set[str]:
    allowed: Set[str] = set()
    with path.open(newline="") as fp:
        reader = csv.DictReader(fp)
        for row in reader:
            sample_id = row.get("sample_id")
            if not sample_id or sample_id == "__mean__":
                continue
            try:
                disc = float(row.get("dice_disc", "nan"))
                cup = float(row.get("dice_cup", "nan"))
            except ValueError:
                continue
            if disc < dice_threshold and cup < dice_threshold:
                continue
            allowed.add(sample_id)
    return allowed


def build_papila_samples(
    image_dir: Path,
    clinical_dir: Path,
    label_col: str,
    positive_labels: Sequence[str],
    allowed_ids: Set[str],
) -> List[RefugeSample]:
    clinical = build_papila_clinical(
        image_dir=str(image_dir),
        clinical_dir=str(clinical_dir),
        label_col=label_col,
        cat_cols=[],
    )
    positives = {lbl.lower() for lbl in positive_labels}
    samples: Dict[str, RefugeSample] = {}
    for _, row in clinical.df.iterrows():
        image_path = clinical.get_image_path(row)
        sample_id = f"papila_{image_path.stem}"
        if sample_id not in allowed_ids or sample_id in samples:
            continue
        value = row.get(label_col)
        if value is None or (isinstance(value, float) and np.isnan(value)):
            continue
        label: Optional[int]
        try:
            label_int = int(value)
            if label_int == 2:
                continue
            label = 1 if label_int > 0 else 0
        except (TypeError, ValueError):
            label = 1 if str(value).strip().lower() in positives else 0
        samples[sample_id] = RefugeSample(
            sample_id=sample_id,
            dataset="papila",
            split="eval",
            image_path=Path(image_path),
            label=label,
            device=None,
            mask_path=None,
            fovea_coord=None,
        )
    return list(samples.values())


def load_contour(path: Path) -> np.ndarray:
    coords = np.loadtxt(path)
    if coords.ndim == 1:
        coords = coords.reshape(-1, 2)
    return coords


def contour_to_mask(coords: np.ndarray, size: Sequence[int]) -> np.ndarray:
    if coords is None or coords.size == 0:
        return np.zeros((size[1], size[0]), dtype=np.uint8)
    img = Image.new("L", size, 0)
    draw = ImageDraw.Draw(img)
    points = [tuple(map(float, pt)) for pt in coords]
    draw.polygon(points, outline=1, fill=1)
    return np.array(img, dtype=np.uint8)


class PapilaGTGeometryProvider:
    def __init__(self, contours_dir: Path) -> None:
        self.contours_dir = contours_dir

    def _pick(self, base: str, kind: str) -> Optional[Path]:
        for exp in ("exp2", "exp1"):
            cand = self.contours_dir / f"{base}_{kind}_{exp}.txt"
            if cand.exists():
                return cand
        return None

    def __call__(self, sample: RefugeSample, scale: float):
        base = Path(sample.image_path).stem
        disc_path = self._pick(base, "disc")
        cup_path = self._pick(base, "cup")
        if disc_path is None or cup_path is None:
            raise RuntimeError(f"Missing ground-truth contours for {sample.sample_id}")

        image = Image.open(sample.image_path).convert("RGB")
        disc_coords = load_contour(disc_path)
        cup_coords = load_contour(cup_path)
        disc_mask = contour_to_mask(disc_coords, image.size)
        cup_mask = contour_to_mask(cup_coords, image.size)
        cup_mask = ((cup_mask > 0) & (disc_mask > 0)).astype(np.uint8)
        geom = _geometry_from_mask(disc_mask, scale)
        return geom, disc_mask.astype(np.uint8), cup_mask.astype(np.uint8)


def build_backbone(name: Optional[str]) -> Optional[torch.nn.Module]:
    if not name:
        return None
    key = name.lower()
    if key not in BACKBONES:
        raise ValueError(f"Unknown backbone '{name}'. Available: {', '.join(sorted(BACKBONES.keys()))}")
    spec = BACKBONES[key]
    model = spec.ctor(weights=spec.weights_default)
    out_dim, model = spec.strip(model)
    setattr(model, "_feature_dim", out_dim)
    if key == "refugelike":
        load_backbone_weights(key, model)
    return model


def evaluate_records(
    clf: RefugeClassification,
    records: Sequence[RefugeClassificationRecord],
    device: str,
    batch_size: int,
) -> Dict[str, float]:
    dataset = RefugeClassificationDataset(
        records,
        transform=clf.eval_transform,
        polar_transform=clf.polar_transform,
        size=clf.crop_size,
    )
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=0)
    clf.backbone.to(device).eval()
    clf.classifier_head.to(device).eval()
    preds: List[float] = []
    targets: List[int] = []
    with torch.no_grad():
        for batch in tqdm(loader, desc="Papila Eval", leave=False, unit="batch"):
            images = batch["image"].to(device)
            polars = batch["polar"].to(device)
            extra_feats = batch["features"].to(device)
            labels = batch["label"].cpu().numpy().tolist()
            feats_img = clf.backbone(images)
            feats = feats_img
            if clf.use_polar:
                feats_polar = clf.backbone(polars)
                feats = torch.cat([feats, feats_polar], dim=1)
            if clf.extra_feature_dim > 0:
                feats = torch.cat([feats, extra_feats], dim=1)
            logits = clf.classifier_head(feats)
            probs = torch.softmax(logits, dim=1)[:, 1].cpu().numpy().tolist()
            preds.extend(probs)
            targets.extend(labels)
    metrics: Dict[str, float] = {"count": float(len(targets))}
    unique_labels = set(targets)
    if len(unique_labels) >= 2:
        metrics["auc"] = float(torchmetrics_auc(targets, preds))
    else:
        metrics["auc"] = float("nan")
    preds_bin = [1 if p >= 0.5 else 0 for p in preds]
    accuracy = sum(int(p == t) for p, t in zip(preds_bin, targets)) / max(1, len(targets))
    metrics["accuracy"] = float(accuracy)
    metrics["mean_prob"] = float(np.mean(preds)) if preds else float("nan")
    metrics["labels_pos"] = float(sum(targets))
    if preds:
        metrics["probs_std"] = float(np.std(preds))
    return metrics


def torchmetrics_auc(targets: Sequence[int], preds: Sequence[float]) -> float:
    try:
        from sklearn.metrics import roc_auc_score
    except ImportError as exc:
        raise RuntimeError("scikit-learn is required to compute AUC") from exc

    return float(roc_auc_score(targets, preds))


def main() -> None:
    args = parse_args()
    device = args.device

    allowed_ids = load_allowed_ids(args.filtered_metrics, args.dice_threshold)
    if not allowed_ids:
        raise SystemExit("No Papila samples passed the Dice threshold.")

    papila_samples = build_papila_samples(
        args.image_dir,
        args.clinical_dir,
        args.label_col,
        args.positive_labels,
        allowed_ids,
    )
    if not papila_samples:
        raise SystemExit("No Papila samples with labels matched the filtered metrics.")

    cache_dir = args.cache_dir
    if args.use_gt_masks and cache_dir is not None:
        cache_dir = cache_dir / "gt"

    if args.use_gt_masks:
        geometry_provider = PapilaGTGeometryProvider(args.gt_contours_dir)
        segmenter = None
    else:
        segmenter = UNetSegmenter(
            manifest_path=args.segmenter_manifest,
            device=device,
            normalize=args.segmenter_normalize,
        )
        seg_state = torch.load(args.segmenter_weights, map_location=device)
        seg_state_dict = seg_state.get("model", seg_state)
        segmenter.model.load_state_dict(seg_state_dict)
        segmenter.model.to(device)
        geometry_provider = UNetGeometryProvider(
            segmenter=segmenter,
            threshold=args.segmenter_threshold,
            tta=args.segmenter_tta,
        )

    pre = RefugePreprocessing(args.refuge_root)
    dummy_seg = RefugeSegmentation(pre)
    backbone = build_backbone(args.backbone)
    clf = RefugeClassification(
        pre,
        dummy_seg,
        geometry_fn=geometry_provider,
        cache_dir=cache_dir,
        backbone=backbone,
    )
    clf.crop_scale = args.crop_scale
    clf.crop_size = args.crop_size
    clf.eval_transform = _default_image_transform(args.crop_size)
    clf.ttt_transform = clf.eval_transform

    if args.classifier_weights is not None:
        clf_state = torch.load(args.classifier_weights, map_location=device)
        clf.backbone.load_state_dict(clf_state["backbone"])
        clf.classifier_head.load_state_dict(clf_state["classifier"])
        clf.rotation_head.load_state_dict(clf_state["rotation"])
        if "feature_reg" in clf_state and getattr(clf, "feature_reg_head", None) is not None:
            clf.feature_reg_head.load_state_dict(clf_state["feature_reg"])

    records = clf.build_records_for_samples(
        papila_samples,
        crop_scale=args.crop_scale,
        progress_prefix="papila_eval",
    )
    if not records:
        raise SystemExit("Unable to build any records; check geometry predictions or labels.")

    metrics = evaluate_records(clf, records, device=device, batch_size=args.batch_size)
    print(f"Samples evaluated: {int(metrics['count'])}")
    print(f"AUC: {metrics['auc']:.4f}" if not np.isnan(metrics['auc']) else "AUC: NaN")
    print(f"Accuracy @0.5: {metrics['accuracy']:.4f}")
    print(f"Mean glaucoma prob: {metrics['mean_prob']:.4f}")

    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        with args.output.open("w", newline="") as fp:
            writer = csv.writer(fp)
            writer.writerow(["sample_id", "prob_glaucoma", "label"])
            clf.backbone.eval()
            clf.classifier_head.eval()
            dataset = RefugeClassificationDataset(
                records,
                transform=clf.eval_transform,
                polar_transform=clf.polar_transform,
                size=clf.crop_size,
            )
            loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, num_workers=0)
            with torch.no_grad():
                for batch in tqdm(loader, desc="Papila Output", leave=False, unit="batch"):
                    images = batch["image"].to(device)
                    polars = batch["polar"].to(device)
                    extra_feats = batch["features"].to(device)
                    ids = batch["sample_id"]
                    labels = batch["label"].tolist()
                    feats_img = clf.backbone(images)
                    feats = feats_img
                    if clf.use_polar:
                        feats_polar = clf.backbone(polars)
                        feats = torch.cat([feats, feats_polar], dim=1)
                    if clf.extra_feature_dim > 0:
                        feats = torch.cat([feats, extra_feats], dim=1)
                    logits = clf.classifier_head(feats)
                    probs = torch.softmax(logits, dim=1)[:, 1].cpu().numpy().tolist()
                    for sid, prob, label in zip(ids, probs, labels):
                        writer.writerow([sid, prob, label])
        print(f"Per-sample probabilities written to {args.output}")


if __name__ == "__main__":
    main()