panseg 0.1.0__py3-none-any.whl

panseg/artifacts.py

@@ -0,0 +1,109 @@
+"""Load timestamped .eval-artifacts dumps produced by panseg infer."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import numpy as np
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+DEFAULT_ARTIFACTS_DIR = REPO_ROOT / ".eval-artifacts"
+
+
+def list_eval_runs(root: Path = DEFAULT_ARTIFACTS_DIR) -> list[Path]:
+    if not root.is_dir():
+        return []
+    runs = [p for p in root.iterdir() if p.is_dir() and (p / "manifest.json").is_file()]
+    return sorted(runs, key=lambda p: p.name, reverse=True)
+
+
+def load_manifest(run_dir: Path) -> dict:
+    path = run_dir / "manifest.json"
+    if not path.is_file():
+        raise FileNotFoundError(f"No manifest.json in {run_dir}")
+    return json.loads(path.read_text())
+
+
+def as_3d(data: np.ndarray) -> np.ndarray:
+    while data.ndim > 3:
+        data = data[..., 0]
+    if data.ndim != 3:
+        raise ValueError(f"Expected a 3D volume, got shape {data.shape}")
+    return data
+
+
+def load_nifti_3d(path: Path) -> np.ndarray:
+    import nibabel as nib
+
+    image = nib.as_closest_canonical(nib.load(path))
+    return as_3d(np.asarray(image.dataobj))
+
+
+def load_eval_artifact(run_dir: Path) -> dict:
+    run_dir = Path(run_dir)
+    manifest = load_manifest(run_dir)
+    image_name = manifest.get("image")
+    if not image_name:
+        raise FileNotFoundError(f"{run_dir} manifest is missing image")
+    image_path = run_dir / image_name
+    if not image_path.is_file():
+        raise FileNotFoundError(f"Image not found: {image_path}")
+
+    gt_name = manifest.get("ground_truth")
+    gt_path = run_dir / gt_name if gt_name else None
+    if gt_path is not None and not gt_path.is_file():
+        gt_path = None
+
+    masks: list[dict] = []
+    for entry in manifest.get("models", []):
+        if entry.get("status") != "ok" or not entry.get("mask"):
+            continue
+        mask_path = run_dir / entry["mask"]
+        if mask_path.is_file():
+            masks.append({"name": entry.get("name", mask_path.stem), "path": mask_path})
+
+    return {
+        "run_dir": run_dir,
+        "manifest": manifest,
+        "image_path": image_path,
+        "ground_truth_path": gt_path,
+        "masks": masks,
+    }
+
+
+def window_slice(volume: np.ndarray, index: int) -> np.ndarray:
+    sl = volume[:, :, index].astype(np.float32)
+    lo, hi = np.percentile(sl, [0.5, 99.5])
+    if hi <= lo:
+        hi = lo + 1.0
+    scaled = np.clip((sl - lo) / (hi - lo), 0.0, 1.0)
+    return scaled.T
+
+
+def overlay_mask(base: np.ndarray, mask: np.ndarray, color: tuple[float, float, float], alpha: float = 0.45) -> np.ndarray:
+    rgb = np.stack([base, base, base], axis=-1)
+    return overlay_on_rgb(rgb, mask, color, alpha)
+
+
+def overlay_on_rgb(
+    rgb: np.ndarray,
+    mask: np.ndarray,
+    color: tuple[float, float, float],
+    alpha: float = 0.45,
+) -> np.ndarray:
+    out = np.array(rgb, dtype=np.float32, copy=True)
+    hit = mask.T > 0
+    tint = np.array(color, dtype=np.float32)
+    out[hit] = (1.0 - alpha) * out[hit] + alpha * tint
+    return np.clip(out, 0.0, 1.0)
+
+
+def dice_score(pred: np.ndarray, target: np.ndarray) -> float:
+    pred_b = pred > 0
+    target_b = target > 0
+    intersection = np.logical_and(pred_b, target_b).sum()
+    denom = pred_b.sum() + target_b.sum()
+    if denom == 0:
+        return float("nan")
+    return float(2.0 * intersection / denom)

panseg/cli.py

@@ -0,0 +1,176 @@
+import logging
+from pathlib import Path
+from typing import Literal
+
+import typer
+
+from panseg.infer import (
+    DEFAULT_PREDICTION_ARTIFACTS_DIR,
+    run_hf_inference,
+)
+from panseg.viewer import launch_viewer
+from panseg.pipeline import PansegLightningModule
+from panseg.utils import DATASET_PARQUETS, DATASETS, MODELS, build_model, download_dataset
+
+app = typer.Typer()
+
+
+@app.callback(invoke_without_command=True)
+def main(ctx: typer.Context) -> None:
+    if ctx.invoked_subcommand is None:
+        print("Hello, welcome to panseg!")
+
+
+@app.command()
+def download(
+    dataset: str = typer.Option(
+        ...,
+        "--dataset",
+        "-d",
+        help="Dataset name, e.g. medical-decathlon-task07",
+    ),
+    output_dir: Path = typer.Option(
+        ...,
+        "--output-dir",
+        "-o",
+        help="Directory to download the dataset into",
+    ),
+) -> None:
+    """Download a registered Kaggle dataset."""
+    try:
+        download_dataset(dataset, output_dir)
+    except ValueError as exc:
+        raise typer.BadParameter(str(exc), param_hint="--dataset") from exc
+    print(f"Downloaded '{dataset}' to {output_dir}")
+
+
+@app.command("eval")
+def evaluate(
+    parquet: Path | None = typer.Option(
+        None,
+        "--parquet",
+        help="Dataset manifest parquet (defaults to the dataset's parquet)",
+        file_okay=True,
+        dir_okay=False,
+        path_type=Path,
+    ),
+    modality: Literal["ct", "mri"] = typer.Option(
+        ...,
+        "--modality",
+        "-m",
+        help="Imaging modality",
+    ),
+    model: str = typer.Option(
+        ...,
+        "--model",
+        help="Model name, e.g. totalseg, dints, mrseg, or nnunet",
+    ),
+    dataset: str = typer.Option(
+        ...,
+        "--dataset",
+        help="Dataset name, e.g. msd-pancreas, amos-pancreas, amos-ct-pancreas, panseg-t1-pancreas, or panseg-t2-pancreas",
+    ),
+    checkpoint: Path | None = typer.Option(
+        None,
+        "--checkpoint",
+        help="nnU-Net checkpoint.pth or trainer folder (required for --model nnunet)",
+        exists=True,
+        file_okay=True,
+        dir_okay=True,
+        path_type=Path,
+    ),
+    output: Path | None = typer.Option(
+        None,
+        "--output",
+        help="Write one row of metrics per evaluated case to this CSV",
+        file_okay=True,
+        dir_okay=False,
+        path_type=Path,
+    ),
+) -> None:
+    """Run evaluation on the test split of a parquet manifest."""
+    if model not in MODELS:
+        known = ", ".join(sorted(MODELS))
+        raise typer.BadParameter(
+            f"Unknown model '{model}'. Known models: {known}",
+            param_hint="--model",
+        )
+    if dataset not in DATASETS:
+        known = ", ".join(sorted(DATASETS))
+        raise typer.BadParameter(
+            f"Unknown dataset '{dataset}'. Known datasets: {known}",
+            param_hint="--dataset",
+        )
+    try:
+        parquet_path = parquet or DATASET_PARQUETS[dataset]
+        module = PansegLightningModule(
+            build_model(model, modality, checkpoint_path=checkpoint)
+        )
+        module.run_eval(parquet_path, dataset, output_csv=output)
+    except ValueError as exc:
+        raise typer.BadParameter(str(exc)) from exc
+    except FileNotFoundError as exc:
+        raise typer.BadParameter(str(exc), param_hint="--parquet") from exc
+
+
+@app.command("infer")
+def infer(
+    input_path: Path = typer.Option(
+        ...,
+        "--input",
+        "-i",
+        help="Path to a NIfTI volume (.nii / .nii.gz)",
+        exists=True,
+        file_okay=True,
+        dir_okay=False,
+        path_type=Path,
+    ),
+    checkpoint: str | None = typer.Option(
+        None,
+        "--checkpoint",
+        help=(
+            "Optional Hugging Face file path in "
+            "dillondesilva/comp3888-capstone-panseg-models. "
+            "Defaults to the CT or MRI annealing checkpoint for the resolved modality."
+        ),
+    ),
+    modality: Literal["ct", "mri"] | None = typer.Option(
+        None,
+        "--modality",
+        "-m",
+        help="Override the intensity-based CT/MRI routing",
+    ),
+    accelerator: Literal["auto", "cpu", "cuda", "xpu", "mps"] = typer.Option(
+        "auto",
+        "--accelerator",
+        "-a",
+        help="Device used to run the checkpoint",
+    ),
+    artifacts_dir: Path = typer.Option(
+        DEFAULT_PREDICTION_ARTIFACTS_DIR,
+        "--artifacts-dir",
+        help="Root directory for timestamped prediction dumps",
+        file_okay=False,
+        dir_okay=True,
+        path_type=Path,
+    ),
+) -> None:
+    """Download an HF nnU-Net checkpoint, segment one volume, and write a prediction artifact."""
+    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
+    try:
+        run_dir = run_hf_inference(
+            input_path,
+            modality=modality,
+            checkpoint=checkpoint,
+            accelerator=accelerator,
+            artifacts_root=artifacts_dir,
+        )
+    except (FileNotFoundError, ValueError, ImportError, RuntimeError) as exc:
+        raise typer.BadParameter(str(exc)) from exc
+    print(f"saved prediction artifact: {run_dir}")
+    if typer.confirm("Launch napari on the prediction artifact?", default=True):
+        launch_viewer(run_dir)
+
+
+if __name__ == "__main__":
+    app()

panseg/constants.py

@@ -0,0 +1,9 @@
+"""Shared panseg constants."""
+
+# Air on the Hounsfield scale, and a bright-voxel floor that CT clears and typical MRI does not.
+CT_AIR_HU = -1000.0
+CT_P99_MIN_HU = 400.0
+
+HF_MODELS_REPO = "dillondesilva/comp3888-capstone-panseg-models"
+DEFAULT_HF_CT_CHECKPOINT = "ct/ct_final_ckpt_c2b_annealing_aggressive.pth"
+DEFAULT_HF_MRI_CHECKPOINT = "mri/mri_final_ckpt_c2b_annealing_aggressive.pth"

panseg/datasets.py

@@ -0,0 +1,162 @@
+import json
+from collections.abc import Sequence
+from pathlib import Path
+
+import nibabel as nib
+import pandas as pd
+import torch
+from torch.utils.data import Dataset
+
+_REPO_ROOT = Path(__file__).resolve().parents[1]
+DEFAULT_PARQUET = _REPO_ROOT / "dataset_generation" / "output" / "ct_pancreas.parquet"
+MRI_PARQUET = _REPO_ROOT / "dataset_generation" / "output" / "mri_pancreas.parquet"
+_REQUIRED_COLUMNS = {
+    "dataset",
+    "case_id",
+    "image_path",
+    "label_path",
+    "label_value",
+    "split",
+}
+
+
+def resolve_volume_path(rel: str, repo_root: Path) -> Path:
+    path = Path(repo_root) / rel
+    if path.is_file():
+        return path
+    if rel.startswith("datasets/"):
+        alt = Path(repo_root) / rel.removeprefix("datasets/")
+        if alt.is_file():
+            return alt
+    return path
+
+
+class ParquetPancreasDataset(Dataset):
+    def __init__(
+        self,
+        parquet_path: Path = DEFAULT_PARQUET,
+        source_dataset: str | Sequence[str] = "msd",
+        split: str | None = "test",
+        repo_root: Path = _REPO_ROOT,
+    ) -> None:
+        parquet_path = Path(parquet_path)
+        if not parquet_path.is_file():
+            raise FileNotFoundError(f"Parquet not found: {parquet_path}")
+
+        df = pd.read_parquet(parquet_path)
+        missing = _REQUIRED_COLUMNS - set(df.columns)
+        if missing:
+            raise ValueError(f"Parquet is missing columns: {sorted(missing)}")
+
+        if isinstance(source_dataset, str):
+            sources = (source_dataset,)
+        else:
+            sources = tuple(source_dataset)
+        df = df[df["dataset"].isin(sources)]
+        if split is not None:
+            df = df[df["split"] == split]
+        if df.empty:
+            qualifier = f" with split={split!r}" if split is not None else ""
+            label = ", ".join(sources)
+            raise FileNotFoundError(
+                f"No '{label}' rows{qualifier} in {parquet_path}"
+            )
+
+        self.cases: list[tuple[str, Path, Path, tuple[int, ...]]] = []
+        for row in df.itertuples(index=False):
+            image_path = resolve_volume_path(row.image_path, repo_root)
+            label_path = resolve_volume_path(row.label_path, repo_root)
+            if not image_path.is_file():
+                raise FileNotFoundError(f"Image not found: {image_path}")
+            if not label_path.is_file():
+                raise FileNotFoundError(f"Label not found: {label_path}")
+            extra = json.loads(getattr(row, "label_extra_values", None) or "[]")
+            values = (int(row.label_value), *(int(v) for v in extra))
+            self.cases.append((row.case_id, image_path, label_path, values))
+
+    def __len__(self) -> int:
+        return len(self.cases)
+
+    def __getitem__(self, index: int) -> dict[str, torch.Tensor | str]:
+        name, image_path, label_path, label_values = self.cases[index]
+        image = nib.load(image_path)
+        volume = torch.from_numpy(image.get_fdata().astype("float32"))
+        label = torch.from_numpy(nib.load(label_path).get_fdata())
+        mask = torch.isin(label, torch.tensor(label_values, dtype=label.dtype)).to(dtype=torch.float32)
+        affine = torch.tensor(image.affine, dtype=torch.float32)
+        return {"volume": volume, "mask": mask, "affine": affine, "name": name}
+
+
+class MSDPancreasDataset(ParquetPancreasDataset):
+    def __init__(
+        self,
+        parquet_path: Path = DEFAULT_PARQUET,
+        split: str | None = "test",
+        repo_root: Path = _REPO_ROOT,
+    ) -> None:
+        super().__init__(
+            parquet_path=parquet_path,
+            source_dataset="msd",
+            split=split,
+            repo_root=repo_root,
+        )
+
+
+class AMOSPancreasDataset(ParquetPancreasDataset):
+    def __init__(
+        self,
+        parquet_path: Path = MRI_PARQUET,
+        split: str | None = "test",
+        repo_root: Path = _REPO_ROOT,
+    ) -> None:
+        super().__init__(
+            parquet_path=parquet_path,
+            source_dataset="amos",
+            split=split,
+            repo_root=repo_root,
+        )
+
+
+class AMOSCTPancreasDataset(ParquetPancreasDataset):
+    def __init__(
+        self,
+        parquet_path: Path = DEFAULT_PARQUET,
+        split: str | None = "test",
+        repo_root: Path = _REPO_ROOT,
+    ) -> None:
+        super().__init__(
+            parquet_path=parquet_path,
+            source_dataset="amos",
+            split=split,
+            repo_root=repo_root,
+        )
+
+
+class PanSegT1Dataset(ParquetPancreasDataset):
+    def __init__(
+        self,
+        parquet_path: Path = MRI_PARQUET,
+        split: str | None = "test",
+        repo_root: Path = _REPO_ROOT,
+    ) -> None:
+        super().__init__(
+            parquet_path=parquet_path,
+            source_dataset="panseg_t1",
+            split=split,
+            repo_root=repo_root,
+        )
+
+
+class PanSegT2Dataset(ParquetPancreasDataset):
+    def __init__(
+        self,
+        parquet_path: Path = MRI_PARQUET,
+        split: str | None = "test",
+        repo_root: Path = _REPO_ROOT,
+    ) -> None:
+        super().__init__(
+            parquet_path=parquet_path,
+            source_dataset="panseg_t2",
+            split=split,
+            repo_root=repo_root,
+        )