diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d1a0ffcd..10a85cd7 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -36,7 +36,7 @@ jobs: - name: Install dependencies run: | uv pip install \ - ".[bayesian,generative,zarr,dev]" \ + ".[bayesian,generative,zarr,qc,dev]" \ monai \ pyro-ppl diff --git a/nobrainer/cli/main.py b/nobrainer/cli/main.py index 3b6ea348..4416d6bb 100644 --- a/nobrainer/cli/main.py +++ b/nobrainer/cli/main.py @@ -668,6 +668,79 @@ def zarr_suggest_shards(n_volumes, volume_shape, dtype, n_input_files, levels): click.echo(json.dumps(result, indent=2)) +# --------------------------------------------------------------------------- +# qc subcommands +# --------------------------------------------------------------------------- + + +@cli.group() +def qc(): + """Quality control tools for brain MRI.""" + + +@qc.command() +@click.option( + "--input-dir", + required=True, + type=click.Path(exists=True), + help="Directory containing reference .nii.gz files.", +) +@click.option( + "--output-dir", + required=True, + type=click.Path(), + help="Root directory for corrupted outputs.", +) +@click.option( + "--corruptions", + default=None, + help="Comma-separated corruption names (default: all).", +) +@click.option( + "--severities", + default="1,2,3,4,5", + help="Comma-separated severity levels.", + **_option_kwds, +) +@click.option("--resume/--no-resume", default=True, **_option_kwds) +@click.option("--dry-run", is_flag=True, help="Print plan without writing files.") +def corrupt(*, input_dir, output_dir, corruptions, severities, resume, dry_run): + """Generate corrupted brain MRI scans for QC benchmarking.""" + from pathlib import Path + + from ..qc.corrupt import generate_corrupted_dataset + + corruption_list = corruptions.split(",") if corruptions else None + severity_list = [int(s) for s in severities.split(",")] + + metadata = generate_corrupted_dataset( + input_dir=Path(input_dir), + output_dir=Path(output_dir), + corruptions=corruption_list, + severities=severity_list, + resume=resume, + dry_run=dry_run, + ) + click.echo(f"Generated {len(metadata)} corrupted scans.") + + +@qc.command() +@click.argument("scan_path", type=click.Path(exists=True)) +@click.option( + "--seg-path", + default=None, + type=click.Path(exists=True), + help="Path to SynthSeg segmentation for CNR/CJV metrics.", +) +def iqms(scan_path, seg_path): + """Extract image quality metrics from a scan.""" + from ..qc.metrics import extract_iqms + + result = extract_iqms(scan_path, seg_path=seg_path) + for key, val in result.items(): + click.echo(f" {key}: {val:.4f}" if val == val else f" {key}: NaN") + + # For debugging only. if __name__ == "__main__": cli() diff --git a/nobrainer/qc/__init__.py b/nobrainer/qc/__init__.py new file mode 100644 index 00000000..b095e958 --- /dev/null +++ b/nobrainer/qc/__init__.py @@ -0,0 +1,30 @@ +"""Quality control module for brain MRI. + +Provides: +- Severity-calibrated corruption generation for QC benchmarking +- Signal-based image quality metric (IQM) extraction +- Machine preference scoring via downstream task degradation +- 3D → 2D slice extraction strategies for VLM evaluation +- Pipeline gating logic (accept/reject/review) +""" + +from nobrainer.qc.corrupt import generate_corrupted_dataset, generate_corrupted_scan +from nobrainer.qc.corruption_configs import CorruptionConfig, get_corruption_configs +from nobrainer.qc.evaluate import QC_PROMPT, parse_qc_response +from nobrainer.qc.gate import QCGate +from nobrainer.qc.metrics import extract_iqms +from nobrainer.qc.preference import compute_dice_preference +from nobrainer.qc.slice_extractor import extract_slices + +__all__ = [ + "CorruptionConfig", + "get_corruption_configs", + "generate_corrupted_scan", + "generate_corrupted_dataset", + "extract_iqms", + "compute_dice_preference", + "extract_slices", + "QC_PROMPT", + "parse_qc_response", + "QCGate", +] diff --git a/nobrainer/qc/corrupt.py b/nobrainer/qc/corrupt.py new file mode 100644 index 00000000..30d7747e --- /dev/null +++ b/nobrainer/qc/corrupt.py @@ -0,0 +1,232 @@ +"""Controlled corruption generation for QC evaluation. + +Generates corrupted versions of brain MRI scans with fixed seeds +and full metadata tracking for reproducible QC benchmarking. +""" + +from __future__ import annotations + +import json +import logging +from pathlib import Path + +import nibabel as nib +import torch +import torchio as tio +from tqdm import tqdm + +from nobrainer.qc.corruption_configs import CorruptionConfig, get_corruption_configs + +logger = logging.getLogger(__name__) + + +def _resolve_device(device: str) -> str: + """Map ``"auto"`` to ``"cuda"`` when available, else ``"cpu"``. + + Other values are passed through verbatim so callers can pin to a + specific device (``"cuda:0"``, ``"cpu"``). + """ + if device == "auto": + return "cuda" if torch.cuda.is_available() else "cpu" + return device + + +def generate_corrupted_scan( + input_path: Path, + output_path: Path, + config: CorruptionConfig, + severity: int, + seed: int, + device: str = "cpu", +) -> dict: + """Apply a corruption to a single scan with fixed seed. + + Parameters: + input_path: Path to reference NIfTI file. + output_path: Path to save corrupted NIfTI file. + config: Corruption configuration. + severity: Severity level (1-5). + seed: Random seed for reproducibility. + device: Torch device on which to run the transform. ``"cpu"`` (default), + ``"cuda"``, or ``"auto"`` (selects ``"cuda"`` when available, else + ``"cpu"``). FFT-based corruptions (motion / ghosting / spike) gain + ~100-300x speedup on a modern CUDA GPU; element-wise corruptions + (noise, blur) gain ~10-30x. Some transforms (e.g. bias_field) fall + back internally to CPU; the wrapper catches and retries on CPU. + + Returns: + Metadata describing the corruption applied. + """ + torch.manual_seed(seed) + resolved_device = _resolve_device(device) + if resolved_device.startswith("cuda"): + torch.cuda.manual_seed_all(seed) + + subject = tio.Subject(t1=tio.ScalarImage(str(input_path))) + if resolved_device != "cpu": + # Move the underlying tensor to GPU so TorchIO's k-space FFTs run + # there. ScalarImage.set_data preserves affine + spatial metadata. + subject.t1.set_data(subject.t1.data.to(resolved_device)) + + transform = config.get_transform(severity) + try: + corrupted = transform(subject) + except (RuntimeError, NotImplementedError, TypeError) as exc: + # Some TorchIO transforms drop to SimpleITK internally + # (RandomMotion's rigid-body resampling) or NumPy-only paths + # (bias_field's polynomial fit) and can't accept CUDA tensors. + # Retry on CPU rather than failing the run. TypeError is the + # specific failure mode for the SITK path: "can't convert + # cuda:0 device type tensor to numpy". + if resolved_device != "cpu": + logger.warning( + "Transform %s failed on %s (%s); retrying on CPU", + config.name, + resolved_device, + exc, + ) + subject = tio.Subject(t1=tio.ScalarImage(str(input_path))) + corrupted = transform(subject) + else: + raise + + # Save using nibabel to preserve original affine and header + orig_nii = nib.load(str(input_path)) + corrupted_data = corrupted.t1.data.squeeze() + if corrupted_data.is_cuda: + corrupted_data = corrupted_data.cpu() + # nibabel requires numpy; this is the one acceptable numpy conversion + corrupted_nii = nib.Nifti1Image( + corrupted_data.numpy(), orig_nii.affine, orig_nii.header + ) + + output_path.parent.mkdir(parents=True, exist_ok=True) + nib.save(corrupted_nii, str(output_path)) + + # Save JSON sidecar with full provenance + metadata = { + "ref_path": str(input_path), + "cor_path": str(output_path), + "corruption_type": config.name, + "corruption_domain": config.domain, + "severity": severity, + "seed": seed, + "transform_params": config.severity_params[severity], + } + sidecar_path = output_path.with_suffix(".json") + sidecar_path.write_text(json.dumps(metadata, indent=2, default=str)) + + return metadata + + +def generate_corrupted_dataset( + input_dir: Path, + output_dir: Path, + corruptions: list[str] | None = None, + severities: list[int] | None = None, + resume: bool = True, + dry_run: bool = False, + device: str = "cpu", +) -> list[dict]: + """Generate corrupted versions of all scans in input_dir. + + Parameters: + input_dir: Directory containing reference .nii.gz files. + output_dir: Root directory for corrupted outputs. + corruptions: Corruption names to apply. None = all 8 types. + severities: Severity levels to generate. None = [1, 2, 3, 4, 5]. + resume: Skip files whose output already exists. + dry_run: Print what would be generated without writing files. + device: Forwarded to :func:`generate_corrupted_scan`. ``"cpu"`` (default), + ``"cuda"``, or ``"auto"``. CPU is the safe default; bench shows + ~100-300x speedup on A100 for FFT-based corruptions. + + Returns: + Metadata dicts, one per corrupted scan. + """ + all_configs = get_corruption_configs() + if corruptions is not None: + unknown = set(corruptions) - set(all_configs) + if unknown: + raise ValueError(f"Unknown corruptions: {unknown}") + all_configs = {k: v for k, v in all_configs.items() if k in corruptions} + if severities is None: + severities = [1, 2, 3, 4, 5] + + input_files = sorted(input_dir.glob("*.nii.gz")) + if not input_files: + raise FileNotFoundError(f"No .nii.gz files found in {input_dir}") + + total = len(input_files) * len(all_configs) * len(severities) + logger.info( + "Corruption plan: %d scans x %d types x %d severities = %d total", + len(input_files), + len(all_configs), + len(severities), + total, + ) + + if dry_run: + logger.info("Dry run — no files will be written.") + return [] + + resolved_device = _resolve_device(device) + if resolved_device != "cpu": + logger.info("Phase 02 corruptions running on device=%s", resolved_device) + + all_metadata: list[dict] = [] + with tqdm(total=total, desc="Generating corruptions") as pbar: + for input_path in input_files: + for config_name, config in all_configs.items(): + for sev in severities: + output_path = ( + output_dir / config_name / f"severity_{sev}" / input_path.name + ) + # Treat zero-byte placeholder files as incomplete: a prior + # interrupted run may have left the path as a truncated + # write. Re-do those rather than skipping with stale data. + if ( + resume + and output_path.exists() + and output_path.stat().st_size > 0 + ): + pbar.update(1) + continue + if output_path.exists() and output_path.stat().st_size == 0: + try: + output_path.unlink() + except OSError: + pass + + seed = hash(f"{input_path.name}_{config_name}_{sev}") % (2**32) + + try: + meta = generate_corrupted_scan( + input_path, + output_path, + config, + sev, + seed, + device=resolved_device, + ) + all_metadata.append(meta) + except Exception as exc: + logger.warning( + "Failed: %s %s sev%d: %s", + input_path.name, + config_name, + sev, + exc, + ) + all_metadata.append( + { + "ref_path": str(input_path), + "cor_path": str(output_path), + "corruption_type": config_name, + "severity": sev, + "error": str(exc), + } + ) + pbar.update(1) + + return all_metadata diff --git a/nobrainer/qc/corruption_configs.py b/nobrainer/qc/corruption_configs.py new file mode 100644 index 00000000..ce9c983b --- /dev/null +++ b/nobrainer/qc/corruption_configs.py @@ -0,0 +1,163 @@ +"""Corruption configurations with severity levels. + +Each corruption type has 5 severity levels with calibrated parameters. +Used by nobrainer.qc.corrupt for systematic QC evaluation. + +K-space corruptions (motion, ghosting, spike) use TorchIO because +MONAI does not simulate k-space physics. + +Image-space corruptions (noise, bias, blur, downsample, gamma) also +use TorchIO for consistency — the QC corruption pipeline operates on +tio.Subject objects, not MONAI dictionary transforms. + +This does NOT duplicate nobrainer.augmentation, which uses MONAI +dictionary transforms for on-the-fly training augmentation. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torchio as tio + + +@dataclass(frozen=True) +class CorruptionConfig: + """Configuration for a single corruption type with severity levels. + + Parameters: + name: Human-readable corruption name. + domain: "kspace" or "image". + transform_class: TorchIO transform class. + severity_params: Mapping from severity level (1-5) to transform kwargs. + """ + + name: str + domain: str + transform_class: type + severity_params: dict[int, dict[str, Any]] + + def get_transform(self, severity: int) -> tio.Transform: + """Create a TorchIO transform for the given severity level. + + Parameters: + severity: Severity level (1-5). + + Returns: + Configured TorchIO transform instance. + """ + if severity not in self.severity_params: + valid = sorted(self.severity_params.keys()) + raise ValueError( + f"Severity {severity} not defined for '{self.name}'. Valid: {valid}" + ) + return self.transform_class(**self.severity_params[severity]) + + +def get_corruption_configs() -> dict[str, CorruptionConfig]: + """Return all corruption configurations with 5 severity levels. + + Returns: + Mapping from corruption name to its configuration. + """ + return { + # ── K-space corruptions (TorchIO only; MONAI has no equivalent) ── + "motion": CorruptionConfig( + name="motion", + domain="kspace", + transform_class=tio.RandomMotion, + severity_params={ + 1: {"degrees": 2, "translation": 1, "num_transforms": 2}, + 2: {"degrees": 4, "translation": 3, "num_transforms": 4}, + 3: {"degrees": 6, "translation": 5, "num_transforms": 6}, + 4: {"degrees": 8, "translation": 7, "num_transforms": 8}, + 5: {"degrees": 10, "translation": 10, "num_transforms": 10}, + }, + ), + "ghosting": CorruptionConfig( + name="ghosting", + domain="kspace", + transform_class=tio.RandomGhosting, + severity_params={ + 1: {"intensity": 0.2, "num_ghosts": 2}, + 2: {"intensity": 0.4, "num_ghosts": 4}, + 3: {"intensity": 0.6, "num_ghosts": 6}, + 4: {"intensity": 0.8, "num_ghosts": 8}, + 5: {"intensity": 1.0, "num_ghosts": 10}, + }, + ), + "spike": CorruptionConfig( + name="spike", + domain="kspace", + transform_class=tio.RandomSpike, + severity_params={ + 1: {"num_spikes": 1, "intensity": (0.5, 1.0)}, + 2: {"num_spikes": 2, "intensity": (1.0, 2.0)}, + 3: {"num_spikes": 3, "intensity": (2.0, 3.0)}, + 4: {"num_spikes": 4, "intensity": (3.0, 4.0)}, + 5: {"num_spikes": 5, "intensity": (4.0, 5.0)}, + }, + ), + # ── Image-space corruptions (TorchIO for pipeline consistency) ── + "noise": CorruptionConfig( + name="noise", + domain="image", + transform_class=tio.RandomNoise, + severity_params={ + 1: {"std": (0.005, 0.01)}, + 2: {"std": (0.01, 0.03)}, + 3: {"std": (0.03, 0.06)}, + 4: {"std": (0.06, 0.10)}, + 5: {"std": (0.10, 0.15)}, + }, + ), + "bias_field": CorruptionConfig( + name="bias_field", + domain="image", + transform_class=tio.RandomBiasField, + severity_params={ + 1: {"coefficients": 0.1}, + 2: {"coefficients": 0.3}, + 3: {"coefficients": 0.5}, + 4: {"coefficients": 0.7}, + 5: {"coefficients": 1.0}, + }, + ), + "blur": CorruptionConfig( + name="blur", + domain="image", + transform_class=tio.RandomBlur, + severity_params={ + 1: {"std": (0.25, 0.5)}, + 2: {"std": (0.5, 1.0)}, + 3: {"std": (1.0, 2.0)}, + 4: {"std": (2.0, 3.0)}, + 5: {"std": (3.0, 4.0)}, + }, + ), + "downsample": CorruptionConfig( + name="downsample", + domain="image", + transform_class=tio.RandomAnisotropy, + severity_params={ + 1: {"downsampling": (1.5, 2.0)}, + 2: {"downsampling": (2.0, 3.0)}, + 3: {"downsampling": (3.0, 4.0)}, + 4: {"downsampling": (4.0, 5.0)}, + 5: {"downsampling": (5.0, 6.0)}, + }, + ), + "gamma": CorruptionConfig( + name="gamma", + domain="image", + transform_class=tio.RandomGamma, + severity_params={ + 1: {"log_gamma": (-0.1, 0.1)}, + 2: {"log_gamma": (-0.2, 0.2)}, + 3: {"log_gamma": (-0.3, 0.3)}, + 4: {"log_gamma": (-0.5, 0.5)}, + 5: {"log_gamma": (-0.7, 0.7)}, + }, + ), + } diff --git a/nobrainer/qc/evaluate.py b/nobrainer/qc/evaluate.py new file mode 100644 index 00000000..890ab897 --- /dev/null +++ b/nobrainer/qc/evaluate.py @@ -0,0 +1,143 @@ +"""VLM-based quality evaluation wrapper. + +Shared prompts and response parsers for quality scoring with +vision-language models. Model loading and per-architecture inference +logic intentionally lives in caller code, not in this module — the +prompts and parsers here are stable across consumers, while the HF / +PEFT / quantisation surface those consumers depend on changes more +frequently. Two prompt + parser pairs are provided: + +* ``QC_PROMPT`` and ``parse_qc_response`` — single-head Likert score + in the format ``"SCORE: N\\nREASON: ..."``. +* ``QC_DUAL_PROMPT`` and ``parse_dual_qc_response`` — two-head Likert + scores in the format ``"Quality: N Thickness: M"``, suitable for + any downstream task that wants a dual-axis quality assessment in a + single forward pass. +""" + +from __future__ import annotations + +import logging +import re + +logger = logging.getLogger(__name__) + +QC_PROMPT = ( + "Assess the quality of this brain MRI scan. Rate it from 1 (unusable, " + "severe artifacts making it unsuitable for analysis) to 5 (excellent " + "quality, no visible artifacts). Describe any quality issues you observe " + "including: motion artifacts, noise, blurring, intensity inhomogeneity, " + "ghosting, or resolution problems.\n" + "Output your response in this exact format:\n" + "SCORE: [integer 1-5]\n" + "REASON: [one paragraph description]" +) + +QC_DUAL_PROMPT = ( + "Assess this brain MRI scan on two axes:\n" + "1. Segmentation quality (1=unusable, 5=excellent) — how well do brain " + "structures segment under SynthSeg-style whole-brain parcellation?\n" + "2. Cortical thickness reliability (1=highly distorted, 5=stable) — how " + "reliable are per-region cortical thickness measurements?\n" + "Output your response in this exact format on one line:\n" + "Quality: [integer 1-5] Thickness: [integer 1-5]" +) + + +def parse_qc_response(text: str) -> dict[str, int | str | bool]: + """Parse a VLM response into structured score and reason. + + Parameters: + text: Raw VLM output text. + + Returns: + Keys: "score" (int or None), "reason" (str), "parse_success" (bool). + """ + result: dict[str, int | str | bool] = { + "score": None, + "reason": "", + "parse_success": False, + } + + # Try structured format first + score_match = re.search(r"SCORE:\s*(\d)", text) + if score_match: + score = int(score_match.group(1)) + if 1 <= score <= 5: + result["score"] = score + result["parse_success"] = True + + # Fallback: extract any digit 1-5 + if result["score"] is None: + digits = re.findall(r"\b([1-5])\b", text) + if digits: + result["score"] = int(digits[0]) + result["parse_success"] = False # mark as fallback + + # Extract reason + reason_match = re.search(r"REASON:\s*(.+)", text, re.DOTALL) + if reason_match: + result["reason"] = reason_match.group(1).strip() + else: + result["reason"] = text.strip() + + return result + + +def parse_dual_qc_response(text: str) -> dict[str, int | None | bool]: + """Parse dual-head QC output ``Quality: N Thickness: M`` into bucket scores. + + Pairs with :data:`QC_DUAL_PROMPT`. A two-head Likert assessment of a + brain MRI scan: one integer 1-5 for segmentation quality, one for + cortical-thickness reliability, emitted in a single forward pass. + The parser extracts both integers; downstream code uses them as + bucketed labels, regression targets, or whatever the consumer needs. + + Tolerant to: case variation (``quality`` / ``Quality`` / ``QUALITY``), + whitespace and punctuation between key and value (``Quality: 4`` / + ``Quality=4`` / ``Quality:4``), trailing decimals (``Quality: 4.0`` + truncates via ``int(float(...))``), and key-value separators on the + same or different lines. + + Out-of-range scores (< 1 or > 5) parse to ``None`` rather than being + clamped, so downstream code can distinguish "model emitted nonsense" + from "model emitted a valid bucket". + + Parameters: + text: Raw VLM output text. + + Returns: + Dict with keys ``"quality"`` (int 1-5 or None), ``"thickness"`` + (int 1-5 or None), and ``"parse_success"`` (bool — True iff + BOTH heads parsed in range). + """ + result: dict[str, int | None | bool] = { + "quality": None, + "thickness": None, + "parse_success": False, + } + + quality_match = re.search(r"quality\s*[:=]?\s*(\d+(?:\.\d+)?)", text, re.IGNORECASE) + if quality_match: + try: + score = int(float(quality_match.group(1))) + if 1 <= score <= 5: + result["quality"] = score + except ValueError: + pass + + thickness_match = re.search( + r"thickness\s*[:=]?\s*(\d+(?:\.\d+)?)", text, re.IGNORECASE + ) + if thickness_match: + try: + score = int(float(thickness_match.group(1))) + if 1 <= score <= 5: + result["thickness"] = score + except ValueError: + pass + + result["parse_success"] = ( + result["quality"] is not None and result["thickness"] is not None + ) + return result diff --git a/nobrainer/qc/gate.py b/nobrainer/qc/gate.py new file mode 100644 index 00000000..2bd2c443 --- /dev/null +++ b/nobrainer/qc/gate.py @@ -0,0 +1,96 @@ +"""Pipeline gating logic for automated QC decisions. + +Given a quality score (from VLM or IQMs), decide whether a scan +should be accepted, rejected, or flagged for manual review. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import logging + +logger = logging.getLogger(__name__) + + +@dataclass +class GateDecision: + """Result of a QC gating decision. + + Parameters: + action: "accept", "reject", or "review". + score: The quality score that triggered the decision. + reason: Human-readable explanation. + """ + + action: str + score: float + reason: str + + +class QCGate: + """Threshold-based QC gating. + + Parameters: + accept_threshold: Scores >= this are accepted. + reject_threshold: Scores <= this are rejected. + score_range: Expected (min, max) for scores. Default (1.0, 5.0) + for VLM scores. + """ + + def __init__( + self, + accept_threshold: float = 3.5, + reject_threshold: float = 2.0, + score_range: tuple[float, float] = (1.0, 5.0), + ) -> None: + if reject_threshold >= accept_threshold: + raise ValueError( + f"reject_threshold ({reject_threshold}) must be < " + f"accept_threshold ({accept_threshold})" + ) + self.accept_threshold = accept_threshold + self.reject_threshold = reject_threshold + self.score_range = score_range + + def decide(self, score: float) -> GateDecision: + """Make a gating decision for a single scan. + + Parameters: + score: Quality score. + + Returns: + The accept/reject/review decision. + """ + if score != score: # NaN check + return GateDecision( + action="review", + score=score, + reason="Score is NaN; manual review required.", + ) + + if score >= self.accept_threshold: + return GateDecision( + action="accept", + score=score, + reason=( + f"Score {score:.2f} >= " f"accept threshold {self.accept_threshold}" + ), + ) + + if score <= self.reject_threshold: + return GateDecision( + action="reject", + score=score, + reason=( + f"Score {score:.2f} <= " f"reject threshold {self.reject_threshold}" + ), + ) + + return GateDecision( + action="review", + score=score, + reason=( + f"Score {score:.2f} between reject ({self.reject_threshold}) " + f"and accept ({self.accept_threshold}) thresholds" + ), + ) diff --git a/nobrainer/qc/metrics.py b/nobrainer/qc/metrics.py new file mode 100644 index 00000000..9e7d0aa5 --- /dev/null +++ b/nobrainer/qc/metrics.py @@ -0,0 +1,189 @@ +"""Signal-based image quality metric extraction. + +Computes mriqc-style IQMs (SNR, CNR, EFC, FBER, CJV) using PyTorch +without requiring the full mriqc BIDS-app pipeline. + +Performance (2026-04-27 patch): + Pre-patch ``_get_tissue_masks`` iterated ``_GM_LABELS`` (72 labels under + ``--parc``) in Python with full-volume OR-merging — same bottleneck + pattern as :mod:`nobrainer.qc.preference`. The current implementation + uses :func:`torch.isin` (single C++ vectorised call) and moves volume + + seg to CUDA when available. Numerics may differ at float32 ULP level + between CPU and GPU; the algebraic definitions are unchanged. +""" + +from __future__ import annotations + +import logging +from pathlib import Path + +import nibabel as nib +import nibabel.processing +import torch + +logger = logging.getLogger(__name__) + +# FreeSurfer label IDs for tissue classification. +# +# GM is the union of the coarse FreeSurfer cortex labels (3 lh, 42 rh) AND +# the Desikan-Killiany parcellation labels (1001-1035 lh, 2001-2035 rh). +# SynthSeg run with ``--parc`` REPLACES the coarse cortex labels with DK; +# without ``--parc`` only 3/42 are emitted. Taking the union keeps CNR and +# CJV defined in either mode — a seg that has any cortical voxel gets a GM +# mask, which both metrics require. Parallel to the cortex-label widening +# in :mod:`nobrainer.qc.preference`. +# +# WM labels are not affected by ``--parc``. +_WM_LABELS: list[int] = [2, 41] +_DK_LH_LABELS: list[int] = list(range(1001, 1036)) +_DK_RH_LABELS: list[int] = list(range(2001, 2036)) +_GM_LABELS: list[int] = [3, 42, *_DK_LH_LABELS, *_DK_RH_LABELS] + + +def _get_tissue_masks( + volume: torch.Tensor, + seg: torch.Tensor | None = None, +) -> dict[str, torch.Tensor]: + """Derive brain/background/WM/GM masks. + + Vectorised: a single :func:`torch.isin` per tissue class replaces the + per-label Python OR-loop. Operates on whichever device the input + tensors live on. + + Parameters: + volume: 3D intensity volume. + seg: SynthSeg segmentation (integer labels). If ``None``, uses Otsu + thresholding as fallback for brain/background only. + + Returns: + Boolean masks: ``"brain"``, ``"background"``, and (when ``seg`` is + provided) ``"wm"``, ``"gm"``. + """ + masks: dict[str, torch.Tensor] = {} + + if seg is not None: + masks["brain"] = seg > 0 + masks["background"] = seg == 0 + wm_t = torch.tensor(_WM_LABELS, dtype=seg.dtype, device=seg.device) + gm_t = torch.tensor(_GM_LABELS, dtype=seg.dtype, device=seg.device) + masks["wm"] = torch.isin(seg, wm_t) + masks["gm"] = torch.isin(seg, gm_t) + return masks + + # Otsu fallback (kept identical to pre-patch; runs on whatever device + # ``volume`` is on). + flat = volume[volume > 0].flatten() + if flat.numel() < 10: + masks["brain"] = volume > 0 + masks["background"] = volume <= 0 + return masks + + hist = torch.histc(flat, bins=256) + bin_edges = torch.linspace( + flat.min().item(), flat.max().item(), 257, device=volume.device + ) + bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2 + + total = hist.sum() + cumsum = torch.cumsum(hist, dim=0) + cumsum_val = torch.cumsum(hist * bin_centers, dim=0) + + w0 = cumsum / total + w1 = 1 - w0 + mean0 = cumsum_val / (cumsum + 1e-8) + mean1 = (cumsum_val[-1] - cumsum_val) / (total - cumsum + 1e-8) + + inter_class_var = w0 * w1 * (mean0 - mean1) ** 2 + threshold = bin_centers[torch.argmax(inter_class_var)] + + masks["brain"] = volume > threshold + masks["background"] = ~masks["brain"] + return masks + + +def extract_iqms( + scan_path: str | Path, + seg_path: str | Path | None = None, +) -> dict[str, float]: + """Extract image quality metrics from a single scan. + + Auto-uses CUDA when available — uploads volume + seg once, runs all + metrics on GPU. Falls back to CPU otherwise. + + Parameters: + scan_path: Path to NIfTI scan. + seg_path: Path to SynthSeg segmentation. If provided, used for tissue + classification. If ``None``, falls back to Otsu thresholding + (only brain/background available; CNR and CJV will be NaN). + + Returns: + Keys: ``"snr"``, ``"cnr"``, ``"efc"``, ``"fber"``, ``"cjv"``. + """ + nii = nib.load(str(scan_path)) + + seg_nii = None + if seg_path is not None: + seg_nii = nib.load(str(seg_path)) + # SynthSeg writes at 1 mm isotropic even when the input scan is + # anisotropic (e.g. FastMRI 0.6875 x 0.6875 x 5 mm), so seg and scan + # frequently have different shapes. Resample the seg into the scan's + # space with nearest-neighbour so downstream mask indexing works and + # scan intensity statistics (SNR, CNR, FBER, CJV) are computed at + # native resolution without interpolation smoothing the noise floor. + if seg_nii.shape != nii.shape: + seg_nii = nibabel.processing.resample_from_to(seg_nii, nii, order=0) + + device = "cuda" if torch.cuda.is_available() else "cpu" + volume = torch.from_numpy(nii.get_fdata()).float().to(device) + + seg = None + if seg_nii is not None: + seg = torch.from_numpy(seg_nii.get_fdata()).long().to(device) + + masks = _get_tissue_masks(volume, seg) + + brain = volume[masks["brain"]] + bg = volume[masks["background"]] + + bg_std = bg.std() if bg.numel() > 1 else torch.tensor(1e-8, device=device) + + # SNR: mean(brain) / std(background) + snr = (brain.mean() / (bg_std + 1e-8)).item() + + # FBER: mean(foreground_energy) / mean(background_energy) + fg_energy = (brain**2).mean() + bg_energy = (bg**2).mean() if bg.numel() > 0 else torch.tensor(1e-8, device=device) + fber = (fg_energy / (bg_energy + 1e-8)).item() + + # EFC: entropy focus criterion + # -sum(p * log(p)) where p = normalized gradient magnitude histogram + grad_x = torch.diff(volume, dim=0).abs() + grad_y = torch.diff(volume, dim=1).abs() + grad_z = torch.diff(volume, dim=2).abs() + # Use mean gradient magnitude (trim to common shape) + min_shape = [ + min(grad_x.shape[i], grad_y.shape[i], grad_z.shape[i]) for i in range(3) + ] + grad_mag = ( + grad_x[: min_shape[0], : min_shape[1], : min_shape[2]] + + grad_y[: min_shape[0], : min_shape[1], : min_shape[2]] + + grad_z[: min_shape[0], : min_shape[1], : min_shape[2]] + ) + grad_flat = grad_mag.flatten() + grad_sum = grad_flat.sum() + 1e-8 + p = grad_flat / grad_sum + p = p[p > 0] + efc = -(p * p.log()).sum().item() + + # CNR and CJV require WM/GM segmentation + cnr = float("nan") + cjv = float("nan") + if "wm" in masks and "gm" in masks: + wm = volume[masks["wm"]] + gm = volume[masks["gm"]] + if wm.numel() > 1 and gm.numel() > 1: + mean_diff = (wm.mean() - gm.mean()).abs() + cnr = (mean_diff / (bg_std + 1e-8)).item() + cjv = ((wm.std() + gm.std()) / (mean_diff + 1e-8)).item() + + return {"snr": snr, "cnr": cnr, "efc": efc, "fber": fber, "cjv": cjv} diff --git a/nobrainer/qc/preference.py b/nobrainer/qc/preference.py new file mode 100644 index 00000000..1a6b2915 --- /dev/null +++ b/nobrainer/qc/preference.py @@ -0,0 +1,132 @@ +"""Machine preference scoring via downstream task degradation. + +Computes per-structure Dice scores between reference and corrupted +SynthSeg segmentations. The Dice degradation between the two +segmentations serves as a machine-derived preference signal: where the +downstream pipeline's output disagrees most with itself under +corruption is a measure of how much the corruption damaged the inputs. + +Performance (2026-04-27 patch): + Pre-patch ``_dice_for_labels`` iterated ``label_ids`` in Python and + OR-merged full-volume boolean masks. For cortex (72 labels under + ``--parc``) on 16M-voxel volumes that's 144 full-volume torch ops on + CPU → ~70 s/pair. The current implementation uses a single + :func:`torch.isin` call (C++ vectorised) and moves both segs to CUDA + when available — ~0.5-1 s/pair on A100. Math is unchanged: + ``Dice = 2|A∩B| / (|A|+|B|)``. +""" + +from __future__ import annotations + +import logging +from pathlib import Path + +import nibabel as nib +import torch + +logger = logging.getLogger(__name__) + +# Key FreeSurfer structures and their label IDs. +# +# Cortex is the union of the coarse FreeSurfer labels (3 lh, 42 rh) AND the +# Desikan-Killiany parcellation labels (1001-1035 lh, 2001-2035 rh). SynthSeg +# with ``--parc`` REPLACES the coarse cortex labels with DK; without ``--parc`` +# only 3/42 are emitted. Taking the union makes cortex_dice defined in either +# mode — a seg that has any cortical voxel gets a value, not NaN. +# +# Subcortical structures stay as single-label pairs. On scans with limited +# FOV (e.g. FastMRI axial slabs that miss brainstem/cerebellum) those Dice +# values are NaN, which is correct — absent-from-seg is not a zero-overlap +# failure, it's a "structure not in the image" statement. mean_dice collapses +# to the average over structures that were present on both sides. +_DK_LH_LABELS: list[int] = list(range(1001, 1036)) +_DK_RH_LABELS: list[int] = list(range(2001, 2036)) + +STRUCTURE_LABELS: dict[str, list[int]] = { + "hippocampus": [17, 53], + "cortex": [3, 42, *_DK_LH_LABELS, *_DK_RH_LABELS], + "ventricle": [4, 43], + "thalamus": [10, 49], + "caudate": [11, 50], + "putamen": [12, 51], + "brainstem": [16], + "cerebellum": [8, 47], +} + + +def _dice_for_labels( + ref: torch.Tensor, + cor: torch.Tensor, + label_ids: list[int], +) -> float: + """Compute Dice for a set of merged label IDs. + + Vectorised: a single :func:`torch.isin` call replaces the per-label + Python OR-loop. For cortex (72 labels) that drops the per-pair work + from 144 full-volume ops to 2. + + Parameters: + ref: Reference segmentation (integer labels, on any torch device). + cor: Corrupted segmentation (integer labels, on any torch device). + label_ids: FreeSurfer label IDs to merge into a single binary mask. + + Returns: + Dice coefficient in [0, 1]. NaN if the structure is absent from both + segmentations (one-sided absence yields 0, by Sørensen-Dice). + """ + if not label_ids: + return float("nan") + label_tensor = torch.tensor(label_ids, dtype=ref.dtype, device=ref.device) + ref_mask = torch.isin(ref, label_tensor) + cor_mask = torch.isin(cor, label_tensor) + + ref_sum = ref_mask.sum() + cor_sum = cor_mask.sum() + + if ref_sum == 0 and cor_sum == 0: + return float("nan") + + intersection = (ref_mask & cor_mask).sum().float() + return (2.0 * intersection / (ref_sum + cor_sum).float()).item() + + +def compute_dice_preference( + ref_seg_path: str | Path, + cor_seg_path: str | Path, +) -> dict[str, float]: + """Compute per-structure Dice between reference and corrupted segmentations. + + Auto-uses CUDA when available — uploads each seg once and runs all 8 + structure compares on GPU. Falls back to CPU otherwise (still ~10× faster + than the pre-patch OR-loop because the per-label work collapses to one + :func:`torch.isin` call regardless of device). + + Parameters: + ref_seg_path: Path to reference SynthSeg segmentation NIfTI. + cor_seg_path: Path to corrupted SynthSeg segmentation NIfTI. + + Returns: + Keys: ``"mean_dice"``, plus ``"{structure}_dice"`` for each + structure in :data:`STRUCTURE_LABELS`. + """ + ref_nii = nib.load(str(ref_seg_path)) + cor_nii = nib.load(str(cor_seg_path)) + + device = "cuda" if torch.cuda.is_available() else "cpu" + ref = torch.from_numpy(ref_nii.get_fdata()).long().to(device) + cor = torch.from_numpy(cor_nii.get_fdata()).long().to(device) + + results: dict[str, float] = {} + dice_values: list[float] = [] + + for name, label_ids in STRUCTURE_LABELS.items(): + d = _dice_for_labels(ref, cor, label_ids) + results[f"{name}_dice"] = d + if not (d != d): # not NaN + dice_values.append(d) + + results["mean_dice"] = ( + sum(dice_values) / len(dice_values) if dice_values else float("nan") + ) + + return results diff --git a/nobrainer/qc/slice_extractor.py b/nobrainer/qc/slice_extractor.py new file mode 100644 index 00000000..ae89ee02 --- /dev/null +++ b/nobrainer/qc/slice_extractor.py @@ -0,0 +1,78 @@ +"""3D → 2D slice extraction strategies for VLM evaluation. + +Extracts representative 2D slices from 3D brain MRI volumes +for input to 2D vision-language models. +""" + +from __future__ import annotations + +import logging +from pathlib import Path + +import nibabel as nib +import torch + +logger = logging.getLogger(__name__) + + +def _normalize_to_uint8(slice_2d: torch.Tensor) -> torch.Tensor: + """Clip to [p2, p98] and normalize to [0, 255].""" + flat = slice_2d.flatten() + if flat.numel() == 0: + return slice_2d.byte() + p2 = torch.quantile(flat.float(), 0.02) + p98 = torch.quantile(flat.float(), 0.98) + clipped = slice_2d.float().clamp(p2, p98) + if p98 - p2 < 1e-8: + return torch.zeros_like(clipped).byte() + normalized = ((clipped - p2) / (p98 - p2) * 255).byte() + return normalized + + +def extract_slices( + scan_path: str | Path, + method: str = "mid", + orientations: list[str] | None = None, +) -> dict[str, torch.Tensor]: + """Extract 2D slices from a 3D volume. + + Parameters: + scan_path: Path to NIfTI volume. + method: Extraction strategy: "mid" (center slices) or "max_info" + (slices with maximum non-zero voxel count). + orientations: Which orientations to extract. + Default: ["axial", "coronal", "sagittal"]. + + Returns: + Keys: "{method}_{orientation}", values: uint8 tensors (H, W). + """ + if orientations is None: + orientations = ["axial", "coronal", "sagittal"] + + nii = nib.load(str(scan_path)) + volume = torch.from_numpy(nii.get_fdata()).float() + + # Axis mapping: orientation → dimension to slice along + axis_map = {"axial": 2, "coronal": 1, "sagittal": 0} + + results: dict[str, torch.Tensor] = {} + + for orient in orientations: + if orient not in axis_map: + raise ValueError(f"Unknown orientation '{orient}'. Use: {list(axis_map)}") + + axis = axis_map[orient] + + if method == "mid": + idx = volume.shape[axis] // 2 + elif method == "max_info": + # Select slice with maximum non-zero voxel count + counts = (volume > 0).sum(dim=[d for d in range(3) if d != axis]) + idx = counts.argmax().item() + else: + raise ValueError(f"Unknown method '{method}'. Use: 'mid', 'max_info'") + + slice_2d = volume.select(axis, idx) + results[f"{method}_{orient}"] = _normalize_to_uint8(slice_2d) + + return results diff --git a/nobrainer/tests/unit/test_qc_corrupt.py b/nobrainer/tests/unit/test_qc_corrupt.py new file mode 100644 index 00000000..4a08d8cc --- /dev/null +++ b/nobrainer/tests/unit/test_qc_corrupt.py @@ -0,0 +1,145 @@ +"""Unit tests for nobrainer.qc.corrupt.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import nibabel as nib +import numpy as np +import pytest + + +def _make_nifti(tmp_path: Path, name: str = "test.nii.gz", shape=(32, 32, 32)) -> Path: + """Create a synthetic NIfTI file with non-zero brain-like data.""" + arr = np.random.RandomState(42).normal(100, 20, shape).astype(np.float32) + arr = np.clip(arr, 0, 200) + path = tmp_path / name + nib.save(nib.Nifti1Image(arr, np.eye(4)), str(path)) + return path + + +class TestGenerateCorruptedScan: + def test_produces_output_file(self, tmp_path): + from nobrainer.qc.corrupt import generate_corrupted_scan + from nobrainer.qc.corruption_configs import get_corruption_configs + + input_path = _make_nifti(tmp_path, "ref.nii.gz") + output_path = tmp_path / "corrupted" / "out.nii.gz" + config = get_corruption_configs()["noise"] + + meta = generate_corrupted_scan( + input_path, output_path, config, severity=3, seed=42 + ) + + assert output_path.exists() + assert output_path.with_suffix(".json").exists() + assert meta["corruption_type"] == "noise" + assert meta["severity"] == 3 + + def test_preserves_affine(self, tmp_path): + from nobrainer.qc.corrupt import generate_corrupted_scan + from nobrainer.qc.corruption_configs import get_corruption_configs + + input_path = _make_nifti(tmp_path) + output_path = tmp_path / "out.nii.gz" + config = get_corruption_configs()["blur"] + + generate_corrupted_scan(input_path, output_path, config, severity=1, seed=0) + + orig = nib.load(str(input_path)) + corrupted = nib.load(str(output_path)) + np.testing.assert_array_almost_equal(orig.affine, corrupted.affine) + + def test_same_seed_same_output(self, tmp_path): + from nobrainer.qc.corrupt import generate_corrupted_scan + from nobrainer.qc.corruption_configs import get_corruption_configs + + input_path = _make_nifti(tmp_path) + config = get_corruption_configs()["noise"] + + out1 = tmp_path / "out1.nii.gz" + out2 = tmp_path / "out2.nii.gz" + + generate_corrupted_scan(input_path, out1, config, severity=2, seed=123) + generate_corrupted_scan(input_path, out2, config, severity=2, seed=123) + + data1 = nib.load(str(out1)).get_fdata() + data2 = nib.load(str(out2)).get_fdata() + np.testing.assert_array_equal(data1, data2) + + def test_different_severity_different_output(self, tmp_path): + from nobrainer.qc.corrupt import generate_corrupted_scan + from nobrainer.qc.corruption_configs import get_corruption_configs + + input_path = _make_nifti(tmp_path) + config = get_corruption_configs()["noise"] + + out1 = tmp_path / "s1.nii.gz" + out5 = tmp_path / "s5.nii.gz" + + generate_corrupted_scan(input_path, out1, config, severity=1, seed=42) + generate_corrupted_scan(input_path, out5, config, severity=5, seed=42) + + data1 = nib.load(str(out1)).get_fdata() + data5 = nib.load(str(out5)).get_fdata() + # Higher severity should produce larger deviation from original + orig = nib.load(str(input_path)).get_fdata() + diff1 = np.abs(data1 - orig).mean() + diff5 = np.abs(data5 - orig).mean() + assert diff5 > diff1 + + def test_sidecar_json_valid(self, tmp_path): + from nobrainer.qc.corrupt import generate_corrupted_scan + from nobrainer.qc.corruption_configs import get_corruption_configs + + input_path = _make_nifti(tmp_path) + output_path = tmp_path / "out.nii.gz" + config = get_corruption_configs()["gamma"] + + generate_corrupted_scan(input_path, output_path, config, severity=2, seed=0) + + sidecar = json.loads(output_path.with_suffix(".json").read_text()) + assert sidecar["corruption_type"] == "gamma" + assert sidecar["severity"] == 2 + assert "seed" in sidecar + + +class TestGenerateCorruptedDataset: + def test_resume_skips_existing(self, tmp_path): + from nobrainer.qc.corrupt import generate_corrupted_dataset + + input_dir = tmp_path / "refs" + input_dir.mkdir() + _make_nifti(input_dir, "scan1.nii.gz") + + output_dir = tmp_path / "corrupted" + # First run + meta1 = generate_corrupted_dataset( + input_dir, output_dir, corruptions=["noise"], severities=[1] + ) + # Second run (should skip) + meta2 = generate_corrupted_dataset( + input_dir, output_dir, corruptions=["noise"], severities=[1] + ) + assert len(meta1) == 1 + assert len(meta2) == 0 # all skipped + + def test_empty_dir_raises(self, tmp_path): + from nobrainer.qc.corrupt import generate_corrupted_dataset + + empty = tmp_path / "empty" + empty.mkdir() + with pytest.raises(FileNotFoundError): + generate_corrupted_dataset(empty, tmp_path / "out") + + def test_unknown_corruption_raises(self, tmp_path): + from nobrainer.qc.corrupt import generate_corrupted_dataset + + input_dir = tmp_path / "refs" + input_dir.mkdir() + _make_nifti(input_dir) + with pytest.raises(ValueError, match="Unknown corruptions"): + generate_corrupted_dataset( + input_dir, tmp_path / "out", corruptions=["nonexistent"] + ) diff --git a/nobrainer/tests/unit/test_qc_corruption_configs.py b/nobrainer/tests/unit/test_qc_corruption_configs.py new file mode 100644 index 00000000..e26be9a6 --- /dev/null +++ b/nobrainer/tests/unit/test_qc_corruption_configs.py @@ -0,0 +1,69 @@ +"""Unit tests for nobrainer.qc.corruption_configs.""" + +from __future__ import annotations + +import pytest + + +class TestCorruptionConfig: + def test_get_all_configs(self): + from nobrainer.qc.corruption_configs import get_corruption_configs + + configs = get_corruption_configs() + assert len(configs) == 8 + assert "motion" in configs + assert "noise" in configs + + def test_all_have_five_severities(self): + from nobrainer.qc.corruption_configs import get_corruption_configs + + for name, config in get_corruption_configs().items(): + assert set(config.severity_params.keys()) == { + 1, + 2, + 3, + 4, + 5, + }, f"{name} missing severities" + + def test_kspace_domain_count(self): + from nobrainer.qc.corruption_configs import get_corruption_configs + + kspace = [c for c in get_corruption_configs().values() if c.domain == "kspace"] + assert len(kspace) == 3 # motion, ghosting, spike + + def test_image_domain_count(self): + from nobrainer.qc.corruption_configs import get_corruption_configs + + image = [c for c in get_corruption_configs().values() if c.domain == "image"] + assert len(image) == 5 # noise, bias_field, blur, downsample, gamma + + def test_get_transform_returns_callable(self): + from nobrainer.qc.corruption_configs import get_corruption_configs + + config = get_corruption_configs()["noise"] + transform = config.get_transform(severity=1) + assert callable(transform) + + def test_invalid_severity_raises(self): + from nobrainer.qc.corruption_configs import get_corruption_configs + + config = get_corruption_configs()["motion"] + with pytest.raises(ValueError, match="Severity 0 not defined"): + config.get_transform(severity=0) + + def test_frozen_dataclass(self): + from nobrainer.qc.corruption_configs import get_corruption_configs + + config = get_corruption_configs()["motion"] + with pytest.raises(AttributeError): + config.name = "changed" + + def test_all_transforms_instantiate(self): + """Every config at every severity should produce a valid transform.""" + from nobrainer.qc.corruption_configs import get_corruption_configs + + for name, config in get_corruption_configs().items(): + for sev in range(1, 6): + transform = config.get_transform(sev) + assert callable(transform), f"{name} sev {sev} not callable" diff --git a/nobrainer/tests/unit/test_qc_evaluate.py b/nobrainer/tests/unit/test_qc_evaluate.py new file mode 100644 index 00000000..bab384fd --- /dev/null +++ b/nobrainer/tests/unit/test_qc_evaluate.py @@ -0,0 +1,138 @@ +"""Unit tests for nobrainer.qc.evaluate.""" + +from __future__ import annotations + + +class TestParseQcResponse: + def test_structured_format(self): + from nobrainer.qc.evaluate import parse_qc_response + + text = "SCORE: 4\nREASON: Good quality scan with minimal noise." + result = parse_qc_response(text) + assert result["score"] == 4 + assert result["parse_success"] is True + assert "Good quality" in result["reason"] + + def test_fallback_digit_extraction(self): + from nobrainer.qc.evaluate import parse_qc_response + + text = "This scan is rated 3 out of 5 due to moderate motion." + result = parse_qc_response(text) + assert result["score"] == 3 + assert result["parse_success"] is False + + def test_no_score_found(self): + from nobrainer.qc.evaluate import parse_qc_response + + text = "Unable to determine quality." + result = parse_qc_response(text) + assert result["score"] is None + assert result["parse_success"] is False + + def test_out_of_range_score_ignored(self): + from nobrainer.qc.evaluate import parse_qc_response + + text = "SCORE: 9\nREASON: Invalid score." + result = parse_qc_response(text) + # 9 is out of 1-5 range, should not be accepted as structured + assert result["score"] is None or result["parse_success"] is False + + def test_reason_without_score_prefix(self): + from nobrainer.qc.evaluate import parse_qc_response + + text = "The image has severe motion artifacts and is unusable." + result = parse_qc_response(text) + assert result["reason"] == text.strip() + + def test_multiline_reason(self): + from nobrainer.qc.evaluate import parse_qc_response + + text = ( + "SCORE: 2\n" + "REASON: Severe motion artifacts visible.\n" + "Additional ghosting in the posterior region." + ) + result = parse_qc_response(text) + assert result["score"] == 2 + assert result["parse_success"] is True + assert "ghosting" in result["reason"] + + def test_qc_prompt_is_string(self): + from nobrainer.qc.evaluate import QC_PROMPT + + assert isinstance(QC_PROMPT, str) + assert "SCORE" in QC_PROMPT + assert "REASON" in QC_PROMPT + + +class TestParseDualQcResponse: + def test_structured_format(self): + from nobrainer.qc.evaluate import parse_dual_qc_response + + result = parse_dual_qc_response("Quality: 4 Thickness: 3") + assert result["quality"] == 4 + assert result["thickness"] == 3 + assert result["parse_success"] is True + + def test_case_insensitive(self): + from nobrainer.qc.evaluate import parse_dual_qc_response + + result = parse_dual_qc_response("quality: 5 thickness: 2") + assert result["quality"] == 5 + assert result["thickness"] == 2 + assert result["parse_success"] is True + + def test_decimal_values_truncated(self): + from nobrainer.qc.evaluate import parse_dual_qc_response + + result = parse_dual_qc_response("Quality: 4.0 Thickness: 3.5") + assert result["quality"] == 4 + assert result["thickness"] == 3 + assert result["parse_success"] is True + + def test_multiline_format(self): + from nobrainer.qc.evaluate import parse_dual_qc_response + + result = parse_dual_qc_response("Quality: 2\nThickness: 4") + assert result["quality"] == 2 + assert result["thickness"] == 4 + assert result["parse_success"] is True + + def test_quality_out_of_range_rejected(self): + from nobrainer.qc.evaluate import parse_dual_qc_response + + result = parse_dual_qc_response("Quality: 9 Thickness: 3") + assert result["quality"] is None + assert result["thickness"] == 3 + assert result["parse_success"] is False + + def test_thickness_out_of_range_rejected(self): + from nobrainer.qc.evaluate import parse_dual_qc_response + + result = parse_dual_qc_response("Quality: 4 Thickness: 0") + assert result["quality"] == 4 + assert result["thickness"] is None + assert result["parse_success"] is False + + def test_only_quality_present(self): + from nobrainer.qc.evaluate import parse_dual_qc_response + + result = parse_dual_qc_response("Quality: 4") + assert result["quality"] == 4 + assert result["thickness"] is None + assert result["parse_success"] is False + + def test_no_match(self): + from nobrainer.qc.evaluate import parse_dual_qc_response + + result = parse_dual_qc_response("Unable to assess this scan.") + assert result["quality"] is None + assert result["thickness"] is None + assert result["parse_success"] is False + + def test_qc_dual_prompt_constant(self): + from nobrainer.qc.evaluate import QC_DUAL_PROMPT + + assert isinstance(QC_DUAL_PROMPT, str) + assert "Quality:" in QC_DUAL_PROMPT + assert "Thickness:" in QC_DUAL_PROMPT diff --git a/nobrainer/tests/unit/test_qc_gate.py b/nobrainer/tests/unit/test_qc_gate.py new file mode 100644 index 00000000..3a168590 --- /dev/null +++ b/nobrainer/tests/unit/test_qc_gate.py @@ -0,0 +1,69 @@ +"""Unit tests for nobrainer.qc.gate.""" + +from __future__ import annotations + +import pytest + + +class TestQCGate: + def test_accept(self): + from nobrainer.qc.gate import QCGate + + gate = QCGate(accept_threshold=3.5, reject_threshold=2.0) + decision = gate.decide(4.0) + assert decision.action == "accept" + + def test_reject(self): + from nobrainer.qc.gate import QCGate + + gate = QCGate(accept_threshold=3.5, reject_threshold=2.0) + decision = gate.decide(1.5) + assert decision.action == "reject" + + def test_review(self): + from nobrainer.qc.gate import QCGate + + gate = QCGate(accept_threshold=3.5, reject_threshold=2.0) + decision = gate.decide(2.5) + assert decision.action == "review" + + def test_nan_triggers_review(self): + from nobrainer.qc.gate import QCGate + + gate = QCGate() + decision = gate.decide(float("nan")) + assert decision.action == "review" + + def test_boundary_accept(self): + from nobrainer.qc.gate import QCGate + + gate = QCGate(accept_threshold=3.5, reject_threshold=2.0) + decision = gate.decide(3.5) + assert decision.action == "accept" + + def test_boundary_reject(self): + from nobrainer.qc.gate import QCGate + + gate = QCGate(accept_threshold=3.5, reject_threshold=2.0) + decision = gate.decide(2.0) + assert decision.action == "reject" + + def test_invalid_thresholds(self): + from nobrainer.qc.gate import QCGate + + with pytest.raises(ValueError, match="must be <"): + QCGate(accept_threshold=2.0, reject_threshold=3.0) + + def test_equal_thresholds_invalid(self): + from nobrainer.qc.gate import QCGate + + with pytest.raises(ValueError, match="must be <"): + QCGate(accept_threshold=3.0, reject_threshold=3.0) + + def test_decision_has_reason(self): + from nobrainer.qc.gate import QCGate + + gate = QCGate() + decision = gate.decide(4.0) + assert isinstance(decision.reason, str) + assert len(decision.reason) > 0 diff --git a/nobrainer/tests/unit/test_qc_metrics.py b/nobrainer/tests/unit/test_qc_metrics.py new file mode 100644 index 00000000..341a4c09 --- /dev/null +++ b/nobrainer/tests/unit/test_qc_metrics.py @@ -0,0 +1,83 @@ +"""Unit tests for nobrainer.qc.metrics.""" + +from __future__ import annotations + +from pathlib import Path + +import nibabel as nib +import numpy as np + + +def _make_scan(tmp_path: Path, shape=(32, 32, 32)) -> Path: + """Create a scan with brain-like contrast (brain center, zero background).""" + arr = np.zeros(shape, dtype=np.float32) + # Brain region in center + arr[8:24, 8:24, 8:24] = np.random.RandomState(0).normal(150, 20, (16, 16, 16)) + arr = np.clip(arr, 0, 300) + path = tmp_path / "scan.nii.gz" + nib.save(nib.Nifti1Image(arr, np.eye(4)), str(path)) + return path + + +def _make_seg(tmp_path: Path, shape=(32, 32, 32)) -> Path: + """Create a segmentation with WM (2), GM (3), CSF (4).""" + seg = np.zeros(shape, dtype=np.int32) + seg[10:22, 10:22, 10:22] = 2 # WM + seg[8:24, 8:24, 8:24][seg[8:24, 8:24, 8:24] == 0] = 3 # GM shell + seg[14:18, 14:18, 14:18] = 4 # CSF center + path = tmp_path / "seg.nii.gz" + nib.save(nib.Nifti1Image(seg, np.eye(4)), str(path)) + return path + + +class TestExtractIqms: + def test_returns_all_keys(self, tmp_path): + from nobrainer.qc.metrics import extract_iqms + + scan = _make_scan(tmp_path) + result = extract_iqms(scan) + assert set(result.keys()) == {"snr", "cnr", "efc", "fber", "cjv"} + + def test_snr_positive(self, tmp_path): + from nobrainer.qc.metrics import extract_iqms + + scan = _make_scan(tmp_path) + result = extract_iqms(scan) + assert result["snr"] > 0 + + def test_with_segmentation_enables_cnr(self, tmp_path): + from nobrainer.qc.metrics import extract_iqms + + scan = _make_scan(tmp_path) + seg = _make_seg(tmp_path) + result = extract_iqms(scan, seg_path=seg) + assert result["cnr"] == result["cnr"] # not NaN + + def test_without_segmentation_cnr_is_nan(self, tmp_path): + from nobrainer.qc.metrics import extract_iqms + + scan = _make_scan(tmp_path) + result = extract_iqms(scan, seg_path=None) + assert result["cnr"] != result["cnr"] # NaN + + def test_efc_is_finite(self, tmp_path): + from nobrainer.qc.metrics import extract_iqms + + scan = _make_scan(tmp_path) + result = extract_iqms(scan) + assert np.isfinite(result["efc"]) + + def test_fber_positive(self, tmp_path): + from nobrainer.qc.metrics import extract_iqms + + scan = _make_scan(tmp_path) + result = extract_iqms(scan) + assert result["fber"] > 0 + + def test_cjv_with_segmentation_is_finite(self, tmp_path): + from nobrainer.qc.metrics import extract_iqms + + scan = _make_scan(tmp_path) + seg = _make_seg(tmp_path) + result = extract_iqms(scan, seg_path=seg) + assert np.isfinite(result["cjv"]) diff --git a/nobrainer/tests/unit/test_qc_preference.py b/nobrainer/tests/unit/test_qc_preference.py new file mode 100644 index 00000000..c2baaf20 --- /dev/null +++ b/nobrainer/tests/unit/test_qc_preference.py @@ -0,0 +1,71 @@ +"""Unit tests for nobrainer.qc.preference.""" + +from __future__ import annotations + +from pathlib import Path + +import nibabel as nib +import numpy as np +import pytest + + +def _make_seg(tmp_path: Path, name: str, perturb: bool = False) -> Path: + """Create a segmentation. If perturb=True, shift some labels.""" + seg = np.zeros((32, 32, 32), dtype=np.int32) + seg[8:24, 8:24, 8:24] = 3 # cortex + seg[10:22, 10:22, 10:22] = 2 # WM + seg[14:18, 14:18, 14:18] = 4 # ventricle + seg[10:14, 10:14, 10:14] = 17 # L hippocampus + + if perturb: + # Shift hippocampus region → reduces Dice + seg[10:14, 10:14, 10:14] = 0 + seg[12:16, 12:16, 12:16] = 17 + + path = tmp_path / name + nib.save(nib.Nifti1Image(seg, np.eye(4)), str(path)) + return path + + +class TestComputeDicePreference: + def test_identical_segs_perfect_dice(self, tmp_path): + from nobrainer.qc.preference import compute_dice_preference + + seg = _make_seg(tmp_path, "seg.nii.gz") + result = compute_dice_preference(seg, seg) + assert result["mean_dice"] == pytest.approx(1.0) + assert result["hippocampus_dice"] == pytest.approx(1.0) + + def test_perturbed_lower_dice(self, tmp_path): + from nobrainer.qc.preference import compute_dice_preference + + ref = _make_seg(tmp_path, "ref.nii.gz", perturb=False) + cor = _make_seg(tmp_path, "cor.nii.gz", perturb=True) + result = compute_dice_preference(ref, cor) + assert result["hippocampus_dice"] < 1.0 + assert result["mean_dice"] < 1.0 + + def test_returns_all_structures(self, tmp_path): + from nobrainer.qc.preference import STRUCTURE_LABELS, compute_dice_preference + + seg = _make_seg(tmp_path, "seg.nii.gz") + result = compute_dice_preference(seg, seg) + for name in STRUCTURE_LABELS: + assert f"{name}_dice" in result + assert "mean_dice" in result + + def test_absent_structure_is_nan(self, tmp_path): + """Structures not present in either seg should return NaN.""" + from nobrainer.qc.preference import compute_dice_preference + + # Our _make_seg doesn't include brainstem (label 16) + seg = _make_seg(tmp_path, "seg.nii.gz") + result = compute_dice_preference(seg, seg) + assert result["brainstem_dice"] != result["brainstem_dice"] # NaN + + def test_cortex_dice_is_one_for_identical(self, tmp_path): + from nobrainer.qc.preference import compute_dice_preference + + seg = _make_seg(tmp_path, "seg.nii.gz") + result = compute_dice_preference(seg, seg) + assert result["cortex_dice"] == pytest.approx(1.0) diff --git a/nobrainer/tests/unit/test_qc_slice_extractor.py b/nobrainer/tests/unit/test_qc_slice_extractor.py new file mode 100644 index 00000000..41543b17 --- /dev/null +++ b/nobrainer/tests/unit/test_qc_slice_extractor.py @@ -0,0 +1,83 @@ +"""Unit tests for nobrainer.qc.slice_extractor.""" + +from __future__ import annotations + +from pathlib import Path + +import nibabel as nib +import numpy as np +import pytest +import torch + + +def _make_volume(tmp_path: Path, shape=(32, 40, 48)) -> Path: + arr = np.random.RandomState(0).normal(100, 30, shape).astype(np.float32) + arr[:5, :, :] = 0 # empty border + path = tmp_path / "vol.nii.gz" + nib.save(nib.Nifti1Image(arr, np.eye(4)), str(path)) + return path + + +class TestExtractSlices: + def test_mid_returns_three_orientations(self, tmp_path): + from nobrainer.qc.slice_extractor import extract_slices + + vol = _make_volume(tmp_path) + slices = extract_slices(vol, method="mid") + assert len(slices) == 3 + assert "mid_axial" in slices + assert "mid_coronal" in slices + assert "mid_sagittal" in slices + + def test_slice_shapes(self, tmp_path): + from nobrainer.qc.slice_extractor import extract_slices + + vol = _make_volume(tmp_path, shape=(32, 40, 48)) + slices = extract_slices(vol, method="mid") + # Axial slices along dim 2 → shape (32, 40) + assert slices["mid_axial"].shape == (32, 40) + # Coronal along dim 1 → shape (32, 48) + assert slices["mid_coronal"].shape == (32, 48) + # Sagittal along dim 0 → shape (40, 48) + assert slices["mid_sagittal"].shape == (40, 48) + + def test_uint8_range(self, tmp_path): + from nobrainer.qc.slice_extractor import extract_slices + + vol = _make_volume(tmp_path) + slices = extract_slices(vol, method="mid") + for s in slices.values(): + assert s.dtype == torch.uint8 + assert s.max() <= 255 + assert s.min() >= 0 + + def test_max_info_selects_non_empty(self, tmp_path): + from nobrainer.qc.slice_extractor import extract_slices + + vol = _make_volume(tmp_path) + slices = extract_slices(vol, method="max_info", orientations=["sagittal"]) + # Should NOT select slice 0-4 (all zeros) + # max_info sagittal = index along dim 0 with most non-zero voxels + assert "max_info_sagittal" in slices + + def test_unknown_method_raises(self, tmp_path): + from nobrainer.qc.slice_extractor import extract_slices + + vol = _make_volume(tmp_path) + with pytest.raises(ValueError, match="Unknown method"): + extract_slices(vol, method="invalid") + + def test_unknown_orientation_raises(self, tmp_path): + from nobrainer.qc.slice_extractor import extract_slices + + vol = _make_volume(tmp_path) + with pytest.raises(ValueError, match="Unknown orientation"): + extract_slices(vol, orientations=["diagonal"]) + + def test_single_orientation(self, tmp_path): + from nobrainer.qc.slice_extractor import extract_slices + + vol = _make_volume(tmp_path) + slices = extract_slices(vol, orientations=["axial"]) + assert len(slices) == 1 + assert "mid_axial" in slices diff --git a/pyproject.toml b/pyproject.toml index ecac034d..4f8d64f1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,9 +64,10 @@ zarr = ["zarr >= 3.0", "nifti-zarr", "ome-zarr >= 0.14.0", "dask[array]", "scipy croissant = ["mlcroissant"] versioning = ["datalad >= 0.19"] tfrecord = ["tfrecord >= 1.14"] +qc = ["torchio >= 1.0"] dev = ["pre-commit", "pytest", "pytest-cov", "scipy"] all = [ - "nobrainer[bayesian,generative,zarr,croissant,versioning,tfrecord,dev]", + "nobrainer[bayesian,generative,zarr,croissant,versioning,tfrecord,qc,dev]", ] [tool.hatch.version]