diff --git a/README.md b/README.md index 0281fb7..bd8d307 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,7 @@ a LangGraph-based multi-agent pipeline for automated classification of neuroimag | `CNNClassifier` | VGG16 / DenseNet169 / ResNet101 | Task-specific classification | | `SAM3Tool` | SAM3 frozen backbone + linear probe | Lesion segmentation (Dice = 0.836) | | `BiomedCLIPTool` | microsoft/BiomedCLIP (ViT-B/16, layer 6) | Zero-shot re-ranking for ambiguous multiclass cases | +| `SiibraAtlasTool` | EBRAINS Julich-Brain v2.9 | Anatomical region assignment via MNI152 coordinates | **Pipeline flow** (linear - every node runs for every image): @@ -47,7 +48,8 @@ MultiAgentMedClassifier/ │ ├── medgemma_agent.py # MedGemma: triage, bbox diagnosis, report │ ├── cnn_tool.py # CNN classifier (VGG16 / DenseNet / ResNet) │ ├── sam3_tool.py # SAM3 segmentation + linear probe head -│ └── biomedclip_tool.py # BiomedCLIP zero-shot / linear probe +│ ├── biomedclip_tool.py # BiomedCLIP zero-shot / linear probe +│ └── sibra_tool.py # EBRAINS siibra: lesion centroid → Julich-Brain region ├── pipeline/ │ ├── graph.py # LangGraph StateGraph assembly │ ├── nodes.py # Node factory functions @@ -64,9 +66,12 @@ MultiAgentMedClassifier/ ├── prompts/ │ ├── system_prompt.txt # MedGemma radiologist persona + JSON schema │ └── system_prompt_bbox.txt # Same schema, bbox-overlay context +├── tests/ +│ └── test_atlas_enrichment.py # Standalone test for atlas node on a single image ├── checkpoints/ # PyTorch state dicts: {arch}_{dataset}_final.pt ├── outputs/ │ ├── explainability/ # Saliency maps: gradcam_pp_*.png, ig_*.png +│ ├── segmentation/ # SAM3 binary masks (mask_*.png) and bbox overlays (guided_*.png) │ ├── fhir/ # FHIR R4 bundles: fhir_.json │ └── eval/ # comparison_summary.csv, _tumor_eval.jsonl ├── app.py # Gradio web GUI (python app.py → http://localhost:7860) @@ -210,12 +215,25 @@ zero-shot mode (BiomedCLIP) instead of attempting a download. ## Usage -**Single image:** +**Single image (PNG/JPEG):** ```bash python run_pipeline.py --image scan.png --task binary_tumor ``` +**Single DICOM slice:** + +```bash +python run_pipeline.py --image scan.dcm --task binary_tumor +``` + +**DICOM series directory:** + +```bash +python run_pipeline.py --image /path/to/dicom_series --task binary_tumor +``` +Raw DICOM inputs are converted to `outputs/preprocessed/*.png` before routing. For series inputs, the middle slice is selected as the representative 2D image and the chosen slice path is preserved in `metadata["dicom_path"]` for atlas enrichment. + **With explainability (Grad-CAM++ + Integrated Gradients):** ```bash @@ -304,6 +322,37 @@ The report is returned in `state["final_report"]` (plain text, ≤150 words). Th | `explainability_result` | Paths to `gradcam_pp_*.png` and `ig_*.png` (if enabled) | | `verification_result` | MedGemma post-hoc agreement check against Grad-CAM++ saliency map (if explainability enabled) | | `fhir_report` | FHIR R4 DiagnosticReport dict; saved to `outputs/fhir/fhir_.json` | +| `atlas_enrichment` | EBRAINS atlas assignment: `assigned_region`, `hemisphere`, `mni_coords`, `assignment_scores` (only on `sam3_then_cnn` path) | + +## Atlas Enrichment (EBRAINS / siibra) + +On the `sam3_then_cnn` path, after SAM3 produces a binary mask, the pipeline runs an optional anatomical assignment step using [siibra-python](https://siibra-python.readthedocs.io) and the [EBRAINS Julich-Brain Cytoarchitectonic Atlas v2.9](https://search.kg.ebrains.eu). + +The SAM3 mask centroid is mapped to MNI152 coordinates, which are then assigned to the nearest cytoarchitectonic region via a statistical probability map. + +**Coordinate accuracy (best → worst):** + +| Input | How it's used | +|---|---| +| `metadata["nifti_path"]` | NIfTI affine → exact MNI coords | +| `metadata["dicom_path"]` | `ImagePositionPatient` + `PixelSpacing` → scanner-space coords (≈ MNI for pre-registered datasets) | +| PNG only (default) | Pixel centroid normalised to MNI152 range; z locked to axial midplane (0 mm) | + +Pass coordinate metadata at inference time: +```python +initial_state("scan.png", "binary_tumor", metadata={"nifti_path": "scan.nii.gz"}) +``` + +If the main input itself is DICOM, `initial_state()` now injects `metadata["dicom_path"]` automatically after preprocessing the slice or series into a PNG. + +**Test the node standalone (without running the full pipeline):** +```bash +python tests/test_atlas_enrichment.py --image data/scan.png --mask outputs/segmentation/mask_abc123.png +``` + +The `atlas_enrichment` result flows into the MedGemma report prompt and is serialised as an `ebrains-atlas-assignment` extension in the FHIR bundle. + +**Note:** z=0.0 (axial midplane) is assumed when no NIfTI or DICOM is available. Region assignments are anatomically meaningful but spatially approximate. Scores are Julich-Brain probability map values; scores < 0.1 occur near region boundaries or at the midplane. ## Explainability Methods diff --git a/agents/dicom_tool.py b/agents/dicom_tool.py new file mode 100644 index 0000000..540f3f7 --- /dev/null +++ b/agents/dicom_tool.py @@ -0,0 +1,212 @@ +# agents/dicom_tool.py +""" +DICOM ingestion and preprocessing for the neuroimaging pipeline. + +Handles: + - Single .dcm file (one slice) + - DICOM series directory (multi-slice → selects representative slice) + - Passthrough for PNG/JPEG (no-op) + +Output: + - PNG saved to outputs/preprocessed/_.png (input to existing pipeline) + - Metadata dict injected into NeuroimagingState +""" + +import pydicom +import numpy as np +from PIL import Image +from pathlib import Path +import hashlib + +from pydicom.misc import is_dicom + + +class DICOMPreprocessor: + + def __init__(self, output_dir: str = "outputs/preprocessed"): + self.output_dir = Path(output_dir) + self.output_dir.mkdir(parents=True, exist_ok=True) + + def prepare(self, input_path: str) -> dict: + """ + Main entry point. Accepts a .dcm file, a directory of .dcm files, + or a PNG/JPEG (passthrough). + + Returns: + { + "image_path": str, # path to PNG ready for CNN/MedGemma + "nifti_path": None, # populated if you add NIfTI conversion (see below) + "dicom_path": str | None, + "dicom_metadata": dict | None + } + """ + p = Path(input_path) + + if p.is_dir(): + return self._process_series(p) + elif self._is_dicom_file(p): + return self._process_single(p) + else: + # Already a PNG/JPEG — passthrough, no metadata + return { + "image_path": str(p), + "nifti_path": None, + "dicom_path": None, + "dicom_metadata": None, + } + + # ── Single slice ────────────────────────────────────────────────────────── + + def _process_single(self, dcm_path: Path) -> dict: + ds = pydicom.dcmread(str(dcm_path)) + metadata = self._extract_metadata(ds) + png_path = self._pixel_array_to_png(ds, self._output_stem(dcm_path)) + return { + "image_path": str(png_path), + "nifti_path": None, + "dicom_path": str(dcm_path), + "dicom_metadata": metadata, + } + + # ── Series (folder of slices) ───────────────────────────────────────────── + + def _process_series(self, series_dir: Path) -> dict: + """ + Load a DICOM series, sort slices by InstanceNumber, + select the middle slice as the representative 2D image, + and extract metadata from that selected slice. + """ + dcm_files = sorted(f for f in series_dir.iterdir() if self._is_dicom_file(f)) + if not dcm_files: + raise ValueError(f"No DICOM files found in {series_dir}") + + # Sort by InstanceNumber (slice position) + slices = [] + for f in dcm_files: + ds = pydicom.dcmread(str(f), stop_before_pixels=False) + instance = int(getattr(ds, "InstanceNumber", 0)) + slices.append((instance, f, ds)) + slices.sort(key=lambda x: x[0]) + + # Representative slice: middle of the series + # For tumour tasks the middle axial slice is usually most informative + # You could also select by maximum lesion area if a rough threshold is known + mid_idx = len(slices) // 2 + _, mid_path, mid_ds = slices[mid_idx] + + metadata = self._extract_metadata(mid_ds) + metadata["total_slices"] = len(slices) + metadata["selected_slice_idx"] = mid_idx + + png_path = self._pixel_array_to_png(mid_ds, self._output_stem(series_dir)) + + return { + "image_path": str(png_path), + "nifti_path": None, # see _series_to_nifti() below if needed + "dicom_path": str(mid_path), + "dicom_metadata": metadata, + } + + # ── Pixel array → normalised PNG ───────────────────────────────────────── + + def _pixel_array_to_png(self, ds: "pydicom.Dataset", stem: str) -> Path: + """ + Convert DICOM pixel array to a normalised 8-bit grayscale PNG. + + Handles: + - Modality LUT (rescale slope/intercept for CT Hounsfield units) + - Window/level for soft tissue vs bone CT windows + - Inversion for some MRI sequences + """ + pixel_array = ds.pixel_array.astype(np.float32) + + # Apply RescaleSlope / RescaleIntercept (standard for CT) + slope = float(getattr(ds, "RescaleSlope", 1.0)) + intercept = float(getattr(ds, "RescaleIntercept", 0.0)) + pixel_array = pixel_array * slope + intercept + + # Apply window/level if present (typical for CT brain window) + window_center = getattr(ds, "WindowCenter", None) + window_width = getattr(ds, "WindowWidth", None) + if window_center and window_width: + wc = float(window_center[0] if hasattr(window_center, "__iter__") else window_center) + ww = float(window_width[0] if hasattr(window_width, "__iter__") else window_width) + lo = wc - ww / 2 + hi = wc + ww / 2 + pixel_array = np.clip(pixel_array, lo, hi) + + if getattr(ds, "PhotometricInterpretation", "") == "MONOCHROME1": + pixel_array = pixel_array.max() - pixel_array + + # Min-max normalise to 0–255 + lo, hi = pixel_array.min(), pixel_array.max() + if hi > lo: + pixel_array = (pixel_array - lo) / (hi - lo) * 255.0 + pixel_array = pixel_array.astype(np.uint8) + + # Convert to RGB (your CNNs and MedGemma expect 3-channel input) + img = Image.fromarray(pixel_array, mode="L").convert("RGB") + + out_path = self.output_dir / f"{stem}.png" + img.save(str(out_path)) + return out_path + + @staticmethod + def _is_dicom_file(path: Path) -> bool: + if not path.is_file(): + return False + if path.suffix.lower() in {".dcm", ".dicom", ".ima"}: + return True + try: + return bool(is_dicom(str(path))) + except Exception: + return False + + @staticmethod + def _output_stem(path: Path) -> str: + digest = hashlib.sha1(str(path.resolve()).encode("utf-8")).hexdigest()[:8] + return f"{path.stem}_{digest}" if path.is_file() else f"{path.name}_{digest}" + + # ── Metadata extraction ─────────────────────────────────────────────────── + + @staticmethod + def _extract_metadata(ds: "pydicom.Dataset") -> dict: + """ + Pull the clinically and technically relevant tags from the DICOM header. + These are injected into NeuroimagingState and passed to MedGemma's prompt. + """ + def safe(tag): + val = getattr(ds, tag, None) + if val is None: + return None + # pydicom sequences are not JSON-serialisable — convert to str + return str(val) if not isinstance(val, (int, float, str)) else val + + return { + # Patient context (anonymised in real deployments) + "patient_age": safe("PatientAge"), # e.g. "068Y" + "patient_sex": safe("PatientSex"), # "M" | "F" | "O" + + # Acquisition parameters — critical for modality routing + "modality": safe("Modality"), # "MR" | "CT" + "series_description": safe("SeriesDescription"),# "T2 FLAIR AX" + "sequence_name": safe("SequenceVariant"), + "scanning_sequence": safe("ScanningSequence"), # "EP" | "GR" | "SE" + "field_strength": safe("MagneticFieldStrength"), # 1.5 | 3.0 + + # Scanner provenance + "manufacturer": safe("Manufacturer"), + "manufacturer_model": safe("ManufacturerModelName"), + "institution": safe("InstitutionName"), + + # Geometric parameters (needed for siibra coordinate mapping) + "slice_thickness_mm": safe("SliceThickness"), + "pixel_spacing": safe("PixelSpacing"), # [row_mm, col_mm] + "image_orientation": safe("ImageOrientationPatient"), + "image_position": safe("ImagePositionPatient"), # MNI registration anchor + + # Study context + "study_description": safe("StudyDescription"), + "body_part": safe("BodyPartExamined"), + "accession_number": safe("AccessionNumber"), # links to RIS/PACS + } diff --git a/agents/medgemma_agent.py b/agents/medgemma_agent.py index 721658a..373ea24 100644 --- a/agents/medgemma_agent.py +++ b/agents/medgemma_agent.py @@ -166,6 +166,9 @@ def diagnosis_to_routing( Task: {task} Route: {routing_path} +Known scan metadata (DICOM/header context, if available): +{metadata_context} + Initial MedGemma triage diagnosis: {initial_medgemma_dx} @@ -190,6 +193,10 @@ def diagnosis_to_routing( Spatial alignment — GradCAM++ ∩ SAM3 mask IoU: {saliency_iou} (IoU < 0.3 suggests the CNN attended to background rather than the lesion) +EBRAINS atlas assignment (Julich-Brain parcellation): +{atlas_enrichment} +(If available, use the assigned_region and hemisphere to contextualise the finding anatomically) + Your response MUST contain exactly two sections in this order: FINDINGS: @@ -343,25 +350,35 @@ def _move_inputs_to_model_device(self, inputs): # ── Primary triage (raw image) ──────────────────────────────────────────── - def diagnose(self, image_path: str) -> tuple[MedicalDiagnosis, RoutingDecision]: + def diagnose( + self, image_path: str, metadata: Optional[dict] = None + ) -> tuple[MedicalDiagnosis, RoutingDecision]: """ Run system_prompt.txt on the raw image. Returns (structured diagnosis, derived routing decision). """ image = Image.open(image_path).convert("RGB") - dx = self._run_diagnostic_prompt(image, SYSTEM_PROMPT) + dx = self._run_diagnostic_prompt( + image, + self._prepend_metadata_context(SYSTEM_PROMPT, metadata), + ) routing = diagnosis_to_routing(dx, self.routing_cfg) return dx, routing # ── SAM3-guided diagnosis (bbox overlay image) ──────────────────────────── - def diagnose_with_bbox(self, guided_image_path: str) -> MedicalDiagnosis: + def diagnose_with_bbox( + self, guided_image_path: str, metadata: Optional[dict] = None + ) -> MedicalDiagnosis: """ Run system_prompt_bbox.txt on the SAM3 bbox-overlay image. Used on the sam3_then_cnn path after segmentation. """ image = Image.open(guided_image_path).convert("RGB") - return self._run_diagnostic_prompt(image, SYSTEM_PROMPT_BBOX) + return self._run_diagnostic_prompt( + image, + self._prepend_metadata_context(SYSTEM_PROMPT_BBOX, metadata), + ) # ── Report generator ────────────────────────────────────────────────────── @@ -378,6 +395,8 @@ def generate_report( explainability_result: Optional[dict] = None, verification_result: Optional[dict] = None, saliency_iou: Optional[float] = None, + atlas_enrichment: Optional[dict] = None, + metadata: Optional[dict] = None, ) -> tuple[str, Optional[MedicalDiagnosis]]: def fmt(d) -> str: if d is None: @@ -388,9 +407,20 @@ def fmt(d) -> str: iou_str = f"{saliency_iou:.3f}" if saliency_iou is not None else "Not computed." + if atlas_enrichment: + atlas_str = ( + f"Region: {atlas_enrichment.get('assigned_region', 'unassigned')} | " + f"Hemisphere: {atlas_enrichment.get('hemisphere', 'unknown')} | " + f"MNI: {atlas_enrichment.get('mni_coords', [])} | " + f"Top candidates: {atlas_enrichment.get('assignment_scores', [])[:3]}" + ) + else: + atlas_str = "Not available (SAM3 path not taken or siibra query failed)." + prompt = REPORT_PROMPT_TEMPLATE.format( task=task, routing_path=" → ".join(routing_path), + metadata_context=self._format_metadata_context(metadata), initial_medgemma_dx=fmt(initial_medgemma_dx), sam3_medgemma_dx=fmt(sam3_medgemma_dx), cnn_result=fmt(cnn_result), @@ -399,6 +429,7 @@ def fmt(d) -> str: explainability_result=fmt(explainability_result), verification_result=fmt(verification_result), saliency_iou=iou_str, + atlas_enrichment=atlas_str, ) image_paths = [image_path] if sam3_result and sam3_result.get("guided_image_path"): @@ -455,6 +486,46 @@ def _run_diagnostic_prompt( severity_confidence=None, ) + def _prepend_metadata_context(self, prompt: str, metadata: Optional[dict]) -> str: + context = self._format_metadata_context(metadata) + if context == "Not available.": + return prompt + return ( + "Known scan metadata (DICOM/header context; may be incomplete):\n" + f"{context}\n\n" + f"{prompt}" + ) + + @staticmethod + def _format_metadata_context(metadata: Optional[dict]) -> str: + if not metadata: + return "Not available." + + fields = [ + ("Modality", "modality"), + ("Series description", "series_description"), + ("Scanning sequence", "scanning_sequence"), + ("Sequence variant", "sequence_name"), + ("Field strength", "field_strength"), + ("Patient age", "patient_age"), + ("Patient sex", "patient_sex"), + ("Slice thickness (mm)", "slice_thickness_mm"), + ("Pixel spacing", "pixel_spacing"), + ("Body part", "body_part"), + ] + lines = [] + for label, key in fields: + value = metadata.get(key) + if value in (None, "", [], {}): + continue + if key == "modality" and value == "MR": + value = "MRI" + if isinstance(value, (list, dict)): + value = json.dumps(value) + lines.append(f"- {label}: {value}") + + return "\n".join(lines) if lines else "Not available." + def _generate( self, image: Image.Image, diff --git a/agents/sibra_tool.py b/agents/sibra_tool.py new file mode 100644 index 0000000..85e537f --- /dev/null +++ b/agents/sibra_tool.py @@ -0,0 +1,335 @@ +# agents/siibra_tool.py +""" +siibra-python integration for anatomical region assignment. + +Requires: + pip install siibra + EBRAINS account for Knowledge Graph queries (optional — atlas queries work without auth, + KG feature queries require a token: https://ebrains.eu/register) + +siibra queries the EBRAINS Human Brain Atlas (Julich-Brain parcellation by default) +to assign a lesion centroid to a named brain region and retrieve multimodal features. +""" + +import tempfile + +import siibra +import nibabel as nib +import numpy as np +from pathlib import Path +from typing import Optional + + +class SiibraAtlasTool: + """ + Wraps siibra-python for anatomical assignment from lesion centroids. + + Usage in the pipeline: + 1. SAM3 outputs a binary mask and bbox in pixel space + 2. NIfTI affine converts pixel centroid → MNI152 mm coordinates + 3. siibra assigns those coordinates to a named parcellation region + 4. Optionally fetches receptor density / connectivity features for that region + + siibra 1.x API note: + Assignment requires a Map object (siibra.get_map), not parcellation.assign(). + Julich Brain uses STATISTICAL (probability) maps — LABELLED maps are not available + for the top-level parcellation and will raise NoMapAvailableError. + """ + + def __init__( + self, + parcellation: str = "julich 2.9", # v2.9 has full STATISTICAL MNI152 coverage + space: str = "mni152", + fetch_features: bool = False, # KG feature queries — requires EBRAINS token + auto_register: bool = False, # register NIfTI → MNI152 before siibra lookup + registration_type: str = "Affine", # "Affine" (fast, ~5-15s) or "SyN" (accurate, ~90s) + ): + self.parcellation = siibra.parcellations[parcellation] + self.space = siibra.spaces[space] + self.fetch_features = fetch_features + self.auto_register = auto_register + self.registration_type = registration_type + # siibra 1.x: assignment is done on a Map, not the parcellation directly + self._pmap = siibra.get_map( + parcellation=self.parcellation, + space=self.space, + maptype=siibra.MapType.STATISTICAL, + ) + # Lazily populated by _get_mni152_template() + self._mni152_template_path: Optional[str] = None + # Cache: nifti_path → invtransforms list (avoids re-registering same scan) + self._registration_cache: dict[str, list[str]] = {} + print(f"[SiibraAtlasTool] parcellation={self.parcellation.name}, map={self._pmap}, auto_register={auto_register}") + + # ── Main entry point ────────────────────────────────────────────────────── + + def assign_lesion( + self, + mask: np.ndarray, # binary mask from SAM3, shape (H, W) + nifti_path: Optional[str] = None, # NIfTI for affine (most accurate) + dicom_path: Optional[str] = None, # DICOM for scanner-space coords + voxel_size_mm: float = 1.0, + ) -> dict: + """ + Assign a SAM3 lesion mask centroid to a brain atlas region. + + Args: + mask: binary 2D mask (from SAM3Tool output) + nifti_path: path to the NIfTI scan the mask was derived from + voxel_size_mm: fallback voxel size if no NIfTI affine is available + + Returns dict with: + mni_coords: [x, y, z] in MNI152 mm + assigned_region: name of the Julich-Brain region + region_description: brief text description (if available) + hemisphere: "left" | "right" | "bilateral" + features: dict of multimodal features (if fetch_features=True) + assignment_scores: probability scores per candidate region + """ + # Step 1: pixel centroid from SAM3 mask + pixel_centroid = self._mask_centroid(mask) + + # Step 2: pixel → MNI mm via affine (NIfTI > DICOM > pixel fallback) + mni_coords = self._pixel_to_mni( + pixel_centroid, nifti_path, voxel_size_mm, + image_size=mask.shape[:2], + dicom_path=dicom_path, + ) + + # Step 3: siibra anatomical assignment via Map (siibra 1.x API) + # Map.assign() returns a DataFrame with columns: + # 'input structure', 'centroid', 'fragment', 'map value', 'region' + point = siibra.Point(tuple(float(v) for v in mni_coords), space=self.space) + try: + assignments = self._pmap.assign(point) + except (IndexError, ValueError) as e: + # Coordinates outside atlas bounds — most often caused by a NIfTI that is + # in native/SRI24 scanner space rather than MNI152. Set auto_register=True + # to run ANTsPy registration before siibra lookup. + print(f"[SiibraAtlasTool] coords {[round(float(v),1) for v in mni_coords]} " + f"outside atlas bounds ({e}). Use auto_register=True for non-MNI152 data.") + return { + "mni_coords": [round(float(v), 2) for v in mni_coords], + "assigned_region": "unassigned", + "hemisphere": "unknown", + "assignment_scores": [], + "error": "coordinates_out_of_bounds", + } + + empty = assignments is None or (hasattr(assignments, "empty") and assignments.empty) + if empty: + return { + "mni_coords": [round(float(v), 2) for v in mni_coords], + "assigned_region": "unassigned", + "hemisphere": "unknown", + "assignment_scores": [], + } + + # Sort by probability score descending + top5 = assignments.nlargest(5, "map value") + best = top5.iloc[0] + region_name = str(best["region"]) + + result = { + "mni_coords": [round(float(v), 2) for v in mni_coords], + "assigned_region": region_name, + "hemisphere": self._hemisphere(region_name), + "assignment_scores": [ + { + "region": str(row["region"]), + "score": round(float(row["map value"]), 4), + } + for _, row in top5.iterrows() + ], + } + + # Step 4: optional KG feature queries + if self.fetch_features: + result["features"] = self._fetch_regional_features(best["region"]) + + return result + + # ── Helpers ─────────────────────────────────────────────────────────────── + + @staticmethod + def _mask_centroid(mask: np.ndarray) -> tuple[float, float]: + """Compute (row, col) centroid of a binary mask.""" + coords = np.argwhere(mask > 0) + if len(coords) == 0: + h, w = mask.shape + return (h / 2, w / 2) + return tuple(coords.mean(axis=0)) + + def _pixel_to_mni( + self, + pixel_centroid: tuple, + nifti_path: Optional[str], + voxel_size_mm: float, + image_size: tuple = (224, 224), # (H, W) of the mask + dicom_path: Optional[str] = None, + ) -> np.ndarray: + """ + Convert 2D pixel centroid to 3D MNI mm coordinates. + + Priority: + 1. NIfTI affine → native-space coords; if auto_register=True, ANTsPy + then maps native → MNI152 via a cached registration transform. + 2. DICOM ImagePositionPatient + PixelSpacing (scanner space, ≈MNI for + pre-registered datasets such as BraTS) + 3. Normalised pixel fallback (approximate — axial centre slice only) + + For the DICOM path, pass the original .dcm file path via + state["metadata"]["dicom_path"]; the main pipeline scan can remain a PNG. + """ + row, col = pixel_centroid + h, w = image_size + + if nifti_path and Path(nifti_path).exists(): + img = nib.load(nifti_path) + affine = img.affine + z_slice = img.shape[2] // 2 + voxel = np.array([col, row, z_slice, 1.0]) + native_coords = (affine @ voxel)[:3] + if self.auto_register: + return self._transform_point_to_mni152(native_coords, nifti_path) + return native_coords + + if dicom_path and Path(dicom_path).exists(): + try: + import pydicom + dcm = pydicom.dcmread(dicom_path) + ipp = np.array(dcm.ImagePositionPatient, dtype=float) # [x,y,z] mm + iop = np.array(dcm.ImageOrientationPatient, dtype=float) + spacing = np.array(dcm.PixelSpacing, dtype=float) # [row_sp, col_sp] + F = iop.reshape(2, 3).T # 3×2 direction-cosine matrix + # Scanner-space coords (≈ MNI for pre-registered BraTS/ADNI scans) + return ipp + F[:, 0] * col * spacing[1] + F[:, 1] * row * spacing[0] + except Exception as e: + print(f"[SiibraAtlasTool] DICOM affine failed ({e}), using pixel fallback") + + # Fallback: normalise pixel position into MNI152 1mm range + # Maps (0,0)→(−91,−109) and (W,H)→(+91,+109) regardless of image resolution + mni_x = (col / w * 182 - 91) * voxel_size_mm + mni_y = (row / h * 218 - 109) * voxel_size_mm + return np.array([mni_x, mni_y, 0.0]) + + def _get_mni152_template(self) -> str: + """ + Return a path to the MNI152 1mm template NIfTI, downloading via nilearn + on first call and caching for the lifetime of this object. + """ + if self._mni152_template_path and Path(self._mni152_template_path).exists(): + return self._mni152_template_path + + from nilearn import datasets as nilearn_datasets + mni_img = nilearn_datasets.load_mni152_template(resolution=1) + tmp = tempfile.NamedTemporaryFile( + suffix=".nii.gz", delete=False, prefix="mni152_template_" + ) + tmp.close() + nib.save(mni_img, tmp.name) + self._mni152_template_path = tmp.name + print(f"[SiibraAtlasTool] MNI152 template saved to {tmp.name}") + return self._mni152_template_path + + def _transform_point_to_mni152( + self, + native_coords: np.ndarray, + nifti_path: str, + ) -> np.ndarray: + """ + Register a NIfTI scan to MNI152 and map a single native-space point + through the resulting transform. + + Uses an image-based approach to avoid apply_transforms_to_points + convention ambiguity: creates a single-voxel indicator volume at the + native centroid, warps it forward to MNI152 via apply_transforms, then + reads the peak voxel location in MNI152 mm coordinates. + + Registration (fwdtransforms) is cached per nifti_path. + """ + import ants + + if nifti_path not in self._registration_cache: + print(f"[SiibraAtlasTool] Registering {Path(nifti_path).name} → MNI152 " + f"({self.registration_type}) …") + template_path = self._get_mni152_template() + fixed = ants.image_read(template_path) + moving = ants.image_read(nifti_path) + reg = ants.registration( + fixed=fixed, + moving=moving, + type_of_transform=self.registration_type, + ) + self._registration_cache[nifti_path] = reg["fwdtransforms"] + print(f"[SiibraAtlasTool] Registration complete.") + + fwd_transforms = self._registration_cache[nifti_path] + template_path = self._get_mni152_template() + fixed = ants.image_read(template_path) + moving = ants.image_read(nifti_path) + + # Convert native mm coords → voxel indices in moving space + img = nib.load(nifti_path) + inv_affine = np.linalg.inv(img.affine) + vox = np.round((inv_affine @ np.append(native_coords, 1))[:3]).astype(int) + vox = np.clip(vox, 0, np.array(img.shape) - 1) + + # Build a single-voxel indicator volume in moving space and warp it forward + indicator = np.zeros(img.shape, dtype=np.float32) + indicator[vox[0], vox[1], vox[2]] = 1.0 + indicator_ants = moving.new_image_like(indicator) + warped = ants.apply_transforms( + fixed=fixed, + moving=indicator_ants, + transformlist=fwd_transforms, + interpolator="linear", + ) + + # Peak voxel in MNI152 space → MNI152 mm via template affine + warped_data = warped.numpy() + if warped_data.max() < 1e-6: + print(f"[SiibraAtlasTool] Warped indicator is empty, returning native coords") + return native_coords + peak_vox = np.unravel_index(warped_data.argmax(), warped_data.shape) + mni_affine = nib.load(template_path).affine + return (mni_affine @ np.array([*peak_vox, 1]))[:3] + + @staticmethod + def _hemisphere(region_name: str) -> str: + name = region_name.lower() + if "_l" in name or "left" in name: + return "left" + elif "_r" in name or "right" in name: + return "right" + return "bilateral" + + def _fetch_regional_features(self, region) -> dict: + """ + Fetch multimodal features linked to this region from EBRAINS KG. + Requires EBRAINS authentication token. + """ + features = {} + try: + receptor_features = siibra.get_features( + region, siibra.features.molecular.ReceptorDensityFingerprint + ) + if receptor_features: + features["receptor_densities"] = { + (f.receptors[0] if f.receptors else "unknown"): + round(float(np.mean(list(f.data.values()))), 4) + for f in receptor_features[:3] + } + except Exception: + pass + + try: + conn_features = siibra.get_features( + region, siibra.features.connectivity.StreamlineCounts + ) + if conn_features: + features["structural_connectivity_available"] = True + except Exception: + pass + + return features \ No newline at end of file diff --git a/eval/evaluate.py b/eval/evaluate.py index 12049bd..a9269d2 100644 --- a/eval/evaluate.py +++ b/eval/evaluate.py @@ -146,7 +146,7 @@ def compute_oracle_routing(cnn_result: dict, routing_cfg=None) -> str: def load_test_split(dataset_dir: str, task: str) -> list[dict]: """ Load test split from a directory structured as: - // + // Returns list of {"image_path": str, "label": str, "task": str} """ @@ -155,22 +155,32 @@ def load_test_split(dataset_dir: str, task: str) -> list[dict]: for class_dir in sorted(dataset_path.iterdir()): if not class_dir.is_dir(): continue - for img_file in class_dir.glob("*.png"): - samples.append( - { - "image_path": str(img_file), - "label": class_dir.name, - "task": task, - } - ) - for img_file in class_dir.glob("*.jpg"): - samples.append( - { - "image_path": str(img_file), - "label": class_dir.name, - "task": task, - } - ) + for item in sorted(class_dir.iterdir()): + if item.is_file() and item.suffix.lower() in { + ".png", + ".jpg", + ".jpeg", + ".dcm", + ".dicom", + ".ima", + }: + samples.append( + { + "image_path": str(item), + "label": class_dir.name, + "task": task, + } + ) + elif item.is_dir() and any( + f.suffix.lower() in {".dcm", ".dicom", ".ima"} for f in item.iterdir() + ): + samples.append( + { + "image_path": str(item), + "label": class_dir.name, + "task": task, + } + ) return samples diff --git a/outputs/siibra_test/100/100.jpg b/outputs/siibra_test/100/100.jpg new file mode 100644 index 0000000..bccaa15 Binary files /dev/null and b/outputs/siibra_test/100/100.jpg differ diff --git a/outputs/siibra_test/100/100_mask.png b/outputs/siibra_test/100/100_mask.png new file mode 100644 index 0000000..e0b1b3b Binary files /dev/null and b/outputs/siibra_test/100/100_mask.png differ diff --git a/outputs/siibra_test/100/100_result.json b/outputs/siibra_test/100/100_result.json new file mode 100644 index 0000000..ff93ac8 --- /dev/null +++ b/outputs/siibra_test/100/100_result.json @@ -0,0 +1,24 @@ +{ + "image": "data/processed/1/100.jpg", + "mask": "data/processed/1_mask/100_mask.png", + "mask_synthetic": false, + "nifti": null, + "dicom": null, + "auto_register": false, + "registration_type": "Affine", + "atlas_enrichment": { + "mni_coords": [ + -20.82, + -14.39, + 0.0 + ], + "assigned_region": "CGL (Metathalamus) right", + "hemisphere": "right", + "assignment_scores": [ + { + "region": "CGL (Metathalamus) right", + "score": 0.0 + } + ] + } +} \ No newline at end of file diff --git a/outputs/siibra_test/tumor_slice/tumor_mask.png b/outputs/siibra_test/tumor_slice/tumor_mask.png new file mode 100644 index 0000000..00dd518 Binary files /dev/null and b/outputs/siibra_test/tumor_slice/tumor_mask.png differ diff --git a/outputs/siibra_test/tumor_slice/tumor_slice.png b/outputs/siibra_test/tumor_slice/tumor_slice.png new file mode 100644 index 0000000..e8d2ba0 Binary files /dev/null and b/outputs/siibra_test/tumor_slice/tumor_slice.png differ diff --git a/outputs/siibra_test/tumor_slice/tumor_slice_result.json b/outputs/siibra_test/tumor_slice/tumor_slice_result.json new file mode 100644 index 0000000..d90b2ee --- /dev/null +++ b/outputs/siibra_test/tumor_slice/tumor_slice_result.json @@ -0,0 +1,28 @@ +{ + "image": "outputs/siibra_test/tumor_slice/tumor_slice.png", + "mask": "outputs/siibra_test/tumor_slice/tumor_mask.png", + "mask_synthetic": false, + "nifti": "data/BraTS2021/BraTS2021_Training_Data/BraTS2021_00000/BraTS2021_00000_t1ce.nii.gz", + "dicom": null, + "auto_register": true, + "registration_type": "Affine", + "atlas_enrichment": { + "mni_coords": [ + 42.0, + -33.0, + 9.0 + ], + "assigned_region": "Area TE 2.2 (STG) right", + "hemisphere": "right", + "assignment_scores": [ + { + "region": "Area TE 2.2 (STG) right", + "score": 0.2295 + }, + { + "region": "Area TE 1.1 (HESCHL) right", + "score": 0.0022 + } + ] + } +} diff --git a/outputs/siibra_test/volume_100_slice_83_t1ce/volume_100_slice_83_mask.png b/outputs/siibra_test/volume_100_slice_83_t1ce/volume_100_slice_83_mask.png new file mode 100644 index 0000000..bd66827 Binary files /dev/null and b/outputs/siibra_test/volume_100_slice_83_t1ce/volume_100_slice_83_mask.png differ diff --git a/outputs/siibra_test/volume_100_slice_83_t1ce/volume_100_slice_83_t1ce.png b/outputs/siibra_test/volume_100_slice_83_t1ce/volume_100_slice_83_t1ce.png new file mode 100644 index 0000000..462347e Binary files /dev/null and b/outputs/siibra_test/volume_100_slice_83_t1ce/volume_100_slice_83_t1ce.png differ diff --git a/outputs/siibra_test/volume_100_slice_83_t1ce/volume_100_slice_83_t1ce_result.json b/outputs/siibra_test/volume_100_slice_83_t1ce/volume_100_slice_83_t1ce_result.json new file mode 100644 index 0000000..6557e2c --- /dev/null +++ b/outputs/siibra_test/volume_100_slice_83_t1ce/volume_100_slice_83_t1ce_result.json @@ -0,0 +1,28 @@ +{ + "image": "outputs/siibra_test/volume_100_slice_83_t1ce.png", + "mask": "outputs/siibra_test/volume_100_slice_83_mask.png", + "mask_synthetic": false, + "nifti": "outputs/siibra_test/volume_100.nii.gz", + "dicom": null, + "auto_register": true, + "registration_type": "Affine", + "atlas_enrichment": { + "mni_coords": [ + -52.23, + 44.1, + 5.0 + ], + "assigned_region": "Frontal-I (GapMap) left", + "hemisphere": "left", + "assignment_scores": [ + { + "region": "Frontal-I (GapMap) left", + "score": 0.6053 + }, + { + "region": "Area 45 (IFG) left", + "score": 0.0827 + } + ] + } +} \ No newline at end of file diff --git a/pipeline/fhir_output.py b/pipeline/fhir_output.py index b93d3b1..e9dadd9 100644 --- a/pipeline/fhir_output.py +++ b/pipeline/fhir_output.py @@ -60,8 +60,16 @@ def _build_patient(patient_id: str) -> dict: def _build_imaging_study(study_id: str, patient_id: str, state: dict) -> dict: """ImagingStudy representing the MRI/CT scan.""" - modality_code = DICOM_MODALITY.get( - (state.get("medgemma_diagnosis") or {}).get("modality", "MRI"), "MR" + metadata = state.get("metadata") or {} + modality_value = metadata.get("modality") or (state.get("medgemma_diagnosis") or {}).get( + "modality", "MRI" + ) + modality_code = DICOM_MODALITY.get(modality_value, modality_value if modality_value in {"MR", "CT"} else "MR") + description = ( + metadata.get("study_description") + or metadata.get("series_description") + or metadata.get("source_image_path") + or state.get("image_path", "unknown") ) return { "resourceType": "ImagingStudy", @@ -72,7 +80,7 @@ def _build_imaging_study(study_id: str, patient_id: str, state: dict) -> dict: "system": "http://dicom.nema.org/resources/ontology/DCM", "code": modality_code, }], - "description": state.get("image_path", "unknown"), + "description": description, } @@ -205,6 +213,19 @@ def _build_diagnostic_report( ], }) + atlas = state.get("atlas_enrichment") + if atlas: + report["extension"].append({ + "url": "https://example.org/fhir/StructureDefinition/ebrains-atlas-assignment", + "extension": [ + {"url": "parcellation", "valueString": "Julich-Brain 3.0"}, + {"url": "assigned-region", "valueString": atlas.get("assigned_region", "")}, + {"url": "hemisphere", "valueString": atlas.get("hemisphere", "")}, + {"url": "mni-coordinates", "valueString": str(atlas.get("mni_coords", []))}, + {"url": "top-candidates", "valueString": str(atlas.get("assignment_scores", [])[:3])}, + ], + }) + saliency = state.get("explainability_result") or {} if saliency: report["presentedForm"] = [ diff --git a/pipeline/graph.py b/pipeline/graph.py index 4257568..b311bda 100644 --- a/pipeline/graph.py +++ b/pipeline/graph.py @@ -7,6 +7,8 @@ cnn_classify │ sam3_segment + | + atlas_enrichment │ biomedclip │ @@ -27,8 +29,11 @@ from agents.cnn_tool import CNNClassifier from agents.medgemma_agent import MedGemmaAgent from agents.sam3_tool import SAM3Tool +from agents.sibra_tool import SiibraAtlasTool from config import DEFAULT_CONFIG, PipelineConfig from pipeline.nodes import ( + human_review_node, + make_atlas_enrichment_node, make_biomedclip_node, make_cnn_node, make_explainability_node, @@ -60,11 +65,13 @@ def build_pipeline(cfg: PipelineConfig = None): cnn = CNNClassifier(cfg.model, cfg.preprocess) sam3 = SAM3Tool(cfg.model, output_dir=f"{cfg.output_dir}/segmentation") clip = BiomedCLIPTool(cfg.model, cfg.preprocess) + siibra = SiibraAtlasTool() # ── Create node functions via factories ─────────────────────────────────── triage_fn = make_triage_node(medgemma, cfg.routing) cnn_fn = make_cnn_node(cnn) sam3_fn = make_sam3_node(sam3) + atlas_fn = make_atlas_enrichment_node(siibra) biomedclip_fn = make_biomedclip_node(clip, cfg.routing) report_fn = make_report_node(medgemma, cfg.routing, skip_report=cfg.skip_report) verification_fn = make_verification_node(medgemma) @@ -76,6 +83,7 @@ def build_pipeline(cfg: PipelineConfig = None): workflow.add_node("triage", triage_fn) workflow.add_node("cnn_classify", cnn_fn) workflow.add_node("sam3_segment", sam3_fn) + workflow.add_node("atlas_enrichment", atlas_fn) workflow.add_node("biomedclip", biomedclip_fn) workflow.add_node("verification", verification_fn) workflow.add_node("report", report_fn) @@ -91,7 +99,8 @@ def build_pipeline(cfg: PipelineConfig = None): workflow.add_edge("triage", "cnn_classify") workflow.add_edge("cnn_classify", "sam3_segment") - workflow.add_edge("sam3_segment", "biomedclip") + workflow.add_edge("sam3_segment", "atlas_enrichment") + workflow.add_edge("atlas_enrichment", "biomedclip") workflow.add_edge("biomedclip", "explainability") workflow.add_edge("explainability", "verification") workflow.add_edge("verification", "report") diff --git a/pipeline/nodes.py b/pipeline/nodes.py index 10818e4..f9c3a20 100644 --- a/pipeline/nodes.py +++ b/pipeline/nodes.py @@ -9,6 +9,7 @@ import time from pathlib import Path +from agents.sibra_tool import SiibraAtlasTool import cv2 import numpy as np import torchvision.transforms as T @@ -48,7 +49,9 @@ def make_triage_node(agent, routing_cfg: RoutingConfig = None): """ def triage_node(state: NeuroimagingState) -> dict: t0 = _log_node_start("triage", state) - dx, routing = agent.diagnose(state["image_path"]) + dx, routing = agent.diagnose( + state["image_path"], metadata=state.get("metadata") + ) _log_node_done("triage", state, t0) return { @@ -121,8 +124,14 @@ def cnn_with_mask_node(state: NeuroimagingState) -> dict: # If SAM3 is unavailable/skipped, avoid a duplicate MedGemma call on the # original image; the initial triage already covered that view. if agent is not None and seg_valid and seg.get("guided_image_path"): - overlay_path = seg["guided_image_path"] - bbox_dx = agent.diagnose_with_bbox(overlay_path) + overlay_path = ( + seg["guided_image_path"] + if seg_valid and seg.get("guided_image_path") + else state["image_path"] + ) + bbox_dx = agent.diagnose_with_bbox( + overlay_path, metadata=state.get("metadata") + ) updates["medgemma_bbox_diagnosis"] = bbox_dx.model_dump() _log_node_done("cnn_with_mask", state, t0) @@ -177,6 +186,8 @@ def report_node(state: NeuroimagingState) -> dict: explainability_result=state.get("explainability_result"), verification_result=state.get("verification_result"), saliency_iou=saliency_iou, + atlas_enrichment=state.get("atlas_enrichment"), + metadata=state.get("metadata"), ) # Determine final prediction from MedGemma's fused diagnosis when available. @@ -449,3 +460,52 @@ def fhir_node(state: NeuroimagingState) -> dict: def route_from_triage(state: NeuroimagingState) -> str: """Legacy helper retained for compatibility with older experiments.""" return state["routing_decision"] + +def make_atlas_enrichment_node(siibra_tool: SiibraAtlasTool): + """ + Runs after SAM3 segmentation (sam3_then_cnn path only). + Loads the SAM3 binary mask from disk, maps the lesion centroid to a + Julich-Brain region via siibra, and stores the result in atlas_enrichment. + + Coordinate accuracy (best → worst): + NIfTI affine — pass state["metadata"]["nifti_path"] + DICOM affine — pass state["metadata"]["dicom_path"] + pixel fallback — approximate, axial centre slice only + """ + def atlas_enrichment_node(state: NeuroimagingState) -> dict: + seg_result = state.get("segmentation_result") + + if not seg_result or not seg_result.get("mask_path"): + return { + "atlas_enrichment": None, + "routing_path": state["routing_path"] + ["atlas_enrichment"], + } + + mask = np.array( + Image.open(seg_result["mask_path"]).convert("L") + ) > 127 # binary uint8-like bool array + + meta = state.get("metadata") or {} + nifti_path = meta.get("nifti_path") + dicom_path = meta.get("dicom_path") + + try: + atlas_result = siibra_tool.assign_lesion( + mask=mask.astype(np.uint8), + nifti_path=nifti_path, + dicom_path=dicom_path, + ) + print( + f"[atlas_enrichment] lesion → {atlas_result['assigned_region']} " + f"({atlas_result['hemisphere']}) MNI {atlas_result['mni_coords']}" + ) + except Exception as e: + print(f"[atlas_enrichment] siibra query failed: {e}") + atlas_result = None + + return { + "atlas_enrichment": atlas_result, + "routing_path": state["routing_path"] + ["atlas_enrichment"], + } + + return atlas_enrichment_node diff --git a/pipeline/state.py b/pipeline/state.py index 80ff500..22e44bf 100644 --- a/pipeline/state.py +++ b/pipeline/state.py @@ -2,8 +2,14 @@ LangGraph state definition for the neuroimaging multi-agent pipeline. """ +from pathlib import Path from typing import Optional, TypedDict +from agents.dicom_tool import DICOMPreprocessor + + +_DICOM_PREPROCESSOR: Optional[DICOMPreprocessor] = None + class SegmentationResult(TypedDict): mask_path: str # Path to saved binary mask PNG @@ -53,6 +59,9 @@ class NeuroimagingState(TypedDict): saliency_sam3_iou: Optional[float] # IoU between GradCAM++ heatmap and SAM3 mask (None when SAM3 mask is empty) sam3_mask_empty: bool # True when SAM3 predicted no lesion pixels + # ── Atlas enrichment (siibra) ───────────────────────────────────────────── + atlas_enrichment: Optional[dict] # assigned_region, mni_coords, hemisphere, scores + # ── Final output ────────────────────────────────────────────────────────── final_report: Optional[str] final_predicted_class: Optional[str] @@ -68,11 +77,30 @@ class NeuroimagingState(TypedDict): def initial_state( image_path: str, task: str, metadata: dict = None ) -> NeuroimagingState: - """Create a blank state for a new image.""" + """ + Create a blank state for a new image. + + Raw DICOM inputs are converted to a PNG in `outputs/preprocessed/` so the + existing PNG-based agents can run unchanged. Relevant DICOM header fields + are merged into `metadata`, and the original DICOM slice path is preserved + as `metadata["dicom_path"]` for atlas coordinate mapping. + """ + prepared = _prepare_input_image(image_path) + prepared_metadata = dict(prepared.get("dicom_metadata") or {}) + if prepared.get("dicom_path"): + prepared_metadata["dicom_path"] = prepared["dicom_path"] + if prepared.get("nifti_path"): + prepared_metadata["nifti_path"] = prepared["nifti_path"] + prepared_metadata["source_image_path"] = str(Path(image_path)) + + merged_metadata = dict(prepared_metadata) + if metadata: + merged_metadata.update(metadata) + return NeuroimagingState( - image_path=image_path, + image_path=prepared["image_path"], task=task, - metadata=metadata or {}, + metadata=merged_metadata, routing_decision=None, routing_confidence=0.0, routing_reasoning="", @@ -91,6 +119,14 @@ def initial_state( final_confidence=0.0, requires_human_review=False, verification_result=None, + atlas_enrichment=None, fhir_report=None, routing_path=[], ) + + +def _prepare_input_image(image_path: str) -> dict: + global _DICOM_PREPROCESSOR + if _DICOM_PREPROCESSOR is None: + _DICOM_PREPROCESSOR = DICOMPreprocessor() + return _DICOM_PREPROCESSOR.prepare(image_path) diff --git a/requirements.txt b/requirements.txt index 4865011..bd8a81b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -50,6 +50,15 @@ scikit-learn # FHIR output fhir.resources>=7.0.0 +# EBRAINS atlas (siibra) +siibra>=1.0a4 +nibabel>=5.0.0 +pydicom>=2.4.0 + +# MNI152 registration (ANTsPy + nilearn template) +antspyx>=0.5.0 +nilearn>=0.10.0 + # Code quality black>=24.0.0 isort>=5.13.0 diff --git a/run_pipeline.py b/run_pipeline.py index fc1df10..db90363 100644 --- a/run_pipeline.py +++ b/run_pipeline.py @@ -2,8 +2,11 @@ Entry point for the multi-agent neuroimaging pipeline. Usage examples: - # Single image: - python run_pipeline.py --image data/processed/1/2.jpg --task binary_tumor + # Single image (PNG/JPEG or a single DICOM slice): + python run_pipeline.py --image image.jpg --task binary_tumor + + # DICOM series directory (selects the middle slice automatically): + python run_pipeline.py --image path/to/dicom_series --task binary_tumor # Full evaluation across all datasets: python run_pipeline.py --eval \ @@ -62,7 +65,11 @@ def parse_args(): # Mode mode = p.add_mutually_exclusive_group(required=True) - mode.add_argument("--image", type=str, help="Path to a single input image") + mode.add_argument( + "--image", + type=str, + help="Path to PNG/JPEG, a single DICOM file, or a directory of DICOM slices", + ) mode.add_argument("--eval", action="store_true", help="Run full evaluation") mode.add_argument( "--tumor_eval", diff --git a/tests/test_atlas_enrichment.py b/tests/test_atlas_enrichment.py new file mode 100644 index 0000000..29dc3d8 --- /dev/null +++ b/tests/test_atlas_enrichment.py @@ -0,0 +1,181 @@ +""" +Quick test for the atlas enrichment node on a single image. + +Usage: + python tests/test_atlas_enrichment.py --image path/to/scan.png + python tests/test_atlas_enrichment.py --image scan.png --mask mask.png + python tests/test_atlas_enrichment.py --image scan.png --dicom scan.dcm + python tests/test_atlas_enrichment.py --image scan.png --nifti scan.nii.gz + + # For data NOT already in MNI152 space (BraTS/SRI24, raw scanner DICOMs): + python tests/test_atlas_enrichment.py --image scan.png --nifti scan.nii.gz --register + python tests/test_atlas_enrichment.py --image scan.png --nifti scan.nii.gz --register --registration_type SyN + +Results are saved to outputs/siibra_test/ by default. + +If --mask is omitted, a synthetic centre-blob mask is generated from the image +dimensions (simulates a lesion in the left-centre of the brain). + +Requires: + pip install siibra nibabel pydicom antspyx nilearn +""" + +import argparse +import json +import sys +import tempfile +from pathlib import Path + +import numpy as np +from PIL import Image + +# ── resolve project root so imports work from any cwd ──────────────────────── +ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(ROOT)) + + +def make_synthetic_mask(image_path: str) -> str: + """ + Generate a binary mask PNG with an ellipse blob in the left-centre of the + image (rough approximation of a left-hemisphere tumour location). + Saved to a temp file; caller is responsible for cleanup. + """ + img = Image.open(image_path).convert("L") + w, h = img.size + + mask = np.zeros((h, w), dtype=np.uint8) + cy, cx = h // 2, w // 3 # left-centre + ry, rx = h // 6, w // 8 # ellipse radii + + Y, X = np.ogrid[:h, :w] + blob = ((X - cx) / rx) ** 2 + ((Y - cy) / ry) ** 2 <= 1.0 + mask[blob] = 255 + + tmp = tempfile.NamedTemporaryFile(suffix="_mask.png", delete=False) + Image.fromarray(mask).save(tmp.name) + return tmp.name + + +def run( + image_path: str, + mask_path: str | None, + nifti_path: str | None, + dicom_path: str | None, + output_path: str | None = None, + auto_register: bool = False, + registration_type: str = "Affine", +): + from agents.sibra_tool import SiibraAtlasTool + from pipeline.nodes import make_atlas_enrichment_node + from pipeline.state import initial_state + + # ── Build a minimal fake state ──────────────────────────────────────────── + synthetic_mask = False + if mask_path is None: + print("[test] No mask supplied — generating synthetic centre-blob mask.") + mask_path = make_synthetic_mask(image_path) + synthetic_mask = True + + state = initial_state(image_path=image_path, task="binary_tumor", metadata={}) + state["segmentation_result"] = { + "mask_path": mask_path, + "bbox": [0, 0, 224, 224], + "guided_image_path": image_path, + "dice_estimate": 0.0, + } + if nifti_path: + state["metadata"]["nifti_path"] = nifti_path + if dicom_path: + state["metadata"]["dicom_path"] = dicom_path + state["routing_path"] = ["triage", "sam3_segment"] + + # ── Run tool directly ───────────────────────────────────────────────────── + print("\n─── SiibraAtlasTool.assign_lesion ───────────────────────────────────") + tool = SiibraAtlasTool( + fetch_features=False, + auto_register=auto_register, + registration_type=registration_type, + ) + + mask_arr = np.array(Image.open(mask_path).convert("L")) > 127 + atlas_result = tool.assign_lesion( + mask=mask_arr.astype(np.uint8), + nifti_path=nifti_path, + dicom_path=dicom_path, + ) + print(json.dumps(atlas_result, indent=2, default=str)) + + # ── Run through the node ────────────────────────────────────────────────── + print("\n─── make_atlas_enrichment_node (node output) ────────────────────────") + node = make_atlas_enrichment_node(tool) + node_output = node(state) + print(json.dumps(node_output, indent=2, default=str)) + + # ── Summary ────────────────────────────────────────────────────────────── + enrichment = node_output.get("atlas_enrichment") or {} + print("\n─── Summary ─────────────────────────────────────────────────────────") + print(f" Image : {image_path}") + print(f" Mask : {mask_path}{' (synthetic)' if synthetic_mask else ''}") + print(f" MNI coords : {enrichment.get('mni_coords', 'n/a')}") + print(f" Region : {enrichment.get('assigned_region', 'n/a')}") + print(f" Hemisphere : {enrichment.get('hemisphere', 'n/a')}") + top = enrichment.get("assignment_scores", [])[:3] + if top: + print(" Top regions :") + for c in top: + print(f" {c['score']:.4f} {c['region']}") + print(f" Routing path : {node_output.get('routing_path')}") + + if output_path: + record = { + "image": image_path, + "mask": mask_path, + "mask_synthetic": synthetic_mask, + "nifti": nifti_path, + "dicom": dicom_path, + "auto_register": auto_register, + "registration_type": registration_type, + "atlas_enrichment": enrichment, + } + Path(output_path).parent.mkdir(parents=True, exist_ok=True) + with open(output_path, "w") as fh: + json.dump(record, fh, indent=2, default=str) + print(f"\n Saved to : {output_path}") + + if synthetic_mask: + Path(mask_path).unlink(missing_ok=True) + + +def main(): + parser = argparse.ArgumentParser(description="Test atlas enrichment node on a single image.") + parser.add_argument("--image", required=True, help="Path to input brain scan PNG") + parser.add_argument("--mask", default=None, help="Binary mask PNG (white = lesion). Auto-generated if omitted.") + parser.add_argument("--nifti", default=None, help="NIfTI file for accurate MNI affine") + parser.add_argument("--dicom", default=None, help="DICOM file for scanner-space coords") + parser.add_argument("--output", default=None, help="Save result JSON to this path (default: outputs/siibra_test/_result.json)") + parser.add_argument("--register", action="store_true", + help="Register NIfTI to MNI152 via ANTsPy before siibra lookup. " + "Required for non-MNI152 data (BraTS/SRI24, raw scanner space).") + parser.add_argument("--registration_type", default="Affine", + choices=["Affine", "SyN"], + help="ANTsPy registration type: Affine (~5-15s) or SyN (~90s, more accurate). Default: Affine") + args = parser.parse_args() + + if not Path(args.image).exists(): + print(f"ERROR: image not found: {args.image}") + sys.exit(1) + + stem = Path(args.image).stem + output_path = args.output or f"outputs/siibra_test/{stem}/{stem}_result.json" + + try: + run(args.image, args.mask, args.nifti, args.dicom, output_path, + auto_register=args.register, registration_type=args.registration_type) + except ModuleNotFoundError as e: + print(f"\nERROR: missing dependency — {e}") + print("Run: pip install siibra nibabel pydicom antspyx nilearn") + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/utils/convert_h5.py b/utils/convert_h5.py new file mode 100644 index 0000000..88dd63e --- /dev/null +++ b/utils/convert_h5.py @@ -0,0 +1,166 @@ +""" +Convert BraTS2020 per-slice h5 files to PNG + NIfTI for pipeline testing. + +The BraTS2020 training h5 files have the structure: + image: (240, 240, 4) float64 — channels: T1, T1ce, T2, FLAIR + mask: (240, 240, 3) uint8 — channels: NCR/NET, ED, ET (tumour sub-regions) + +Usage: + # Convert volume 100, auto-pick best slice, build NIfTI + python utils/convert_h5.py --volume 100 + + # Specific slice + python utils/convert_h5.py --volume 100 --slice 80 + + # Custom data dir and output dir + python utils/convert_h5.py --volume 100 --data_dir data/BraTS2020/BraTS2020_training_data/data --out_dir outputs/brats_test + +Output files: + volume_{id}_slice_{n}_t1ce.png — T1ce channel greyscale PNG (pipeline input) + volume_{id}_slice_{n}_mask.png — binary tumour mask PNG (white = tumour) + volume_{id}.nii.gz — full 3D T1ce volume in MNI152 space (accurate siibra coords) + +Then test the atlas tool: + python tests/test_atlas_enrichment.py \\ + --image outputs/brats_test/volume_100_slice_80_t1ce.png \\ + --mask outputs/brats_test/volume_100_slice_80_mask.png \\ + --nifti outputs/brats_test/volume_100.nii.gz +""" + +import argparse +import os +from pathlib import Path + +import h5py +import nibabel as nib +import numpy as np +from PIL import Image + +# BraTS2020 is pre-registered to MNI152 1mm isotropic space. +# Standard affine maps voxel (i,j,k) → MNI152 mm. +BRATS_MNI152_AFFINE = np.array([ + [-1, 0, 0, 90], + [ 0, -1, 0, 126], + [ 0, 0, 1, -72], + [ 0, 0, 0, 1], +], dtype=np.float32) + +CHANNEL_NAMES = ["t1", "t1ce", "t2", "flair"] + + +def load_slice(h5_path: Path) -> tuple[np.ndarray, np.ndarray]: + with h5py.File(h5_path, "r") as f: + image = f["image"][()] # (240, 240, 4) float64 + mask = f["mask"][()] # (240, 240, 3) uint8 + return image, mask + + +def find_slices(data_dir: Path, volume_id: int) -> list[Path]: + prefix = f"volume_{volume_id}_slice_" + slices = sorted( + data_dir.glob(f"{prefix}*.h5"), + key=lambda p: int(p.stem.split("_slice_")[-1]), + ) + if not slices: + raise FileNotFoundError(f"No h5 files found for volume {volume_id} in {data_dir}") + return slices + + +def best_slice(slice_paths: list[Path]) -> tuple[int, Path]: + """Return (slice_idx, path) for the slice with the most tumour voxels.""" + best_idx, best_path, best_count = 0, slice_paths[0], 0 + for path in slice_paths: + _, mask = load_slice(path) + count = int((mask > 0).sum()) + if count > best_count: + best_count = count + best_path = path + best_idx = int(path.stem.split("_slice_")[-1]) + return best_idx, best_path + + +def normalise_to_uint8(arr: np.ndarray) -> np.ndarray: + lo, hi = arr.min(), arr.max() + if hi == lo: + return np.zeros_like(arr, dtype=np.uint8) + return ((arr - lo) / (hi - lo) * 255).astype(np.uint8) + + +def save_png(arr_2d: np.ndarray, path: Path) -> None: + Image.fromarray(arr_2d).save(path) + + +def build_nifti(slice_paths: list[Path], channel: int = 1) -> nib.Nifti1Image: + """Stack all slices into a 3D volume and return a NIfTI with BraTS MNI152 affine.""" + slices = [] + for p in slice_paths: + img, _ = load_slice(p) + slices.append(img[:, :, channel]) # (240, 240) + + volume = np.stack(slices, axis=2).astype(np.float32) # (240, 240, N_slices) + return nib.Nifti1Image(volume, affine=BRATS_MNI152_AFFINE) + + +def convert(volume_id: int, data_dir: Path, out_dir: Path, slice_idx: int | None = None, channel: int = 1): + out_dir.mkdir(parents=True, exist_ok=True) + + slice_paths = find_slices(data_dir, volume_id) + print(f"Found {len(slice_paths)} slices for volume {volume_id}") + + if slice_idx is None: + slice_idx, chosen_path = best_slice(slice_paths) + print(f"Auto-selected slice {slice_idx} (most tumour voxels)") + else: + chosen_path = data_dir / f"volume_{volume_id}_slice_{slice_idx}.h5" + if not chosen_path.exists(): + raise FileNotFoundError(f"Slice not found: {chosen_path}") + + image, mask = load_slice(chosen_path) + ch_name = CHANNEL_NAMES[channel] + + # PNG: selected MRI channel + png_path = out_dir / f"volume_{volume_id}_slice_{slice_idx}_{ch_name}.png" + save_png(normalise_to_uint8(image[:, :, channel]), png_path) + print(f"Saved image PNG : {png_path}") + + # Binary mask PNG: union of all 3 tumour sub-regions + binary_mask = (mask.sum(axis=2) > 0).astype(np.uint8) * 255 + mask_path = out_dir / f"volume_{volume_id}_slice_{slice_idx}_mask.png" + save_png(binary_mask, mask_path) + print(f"Saved mask PNG : {mask_path}") + + # NIfTI: full 3D volume with MNI152 affine + nifti_path = out_dir / f"volume_{volume_id}.nii.gz" + nii = build_nifti(slice_paths, channel=channel) + nib.save(nii, nifti_path) + print(f"Saved NIfTI : {nifti_path}") + + print(f"\nRun atlas test with:") + print(f" python tests/test_atlas_enrichment.py \\") + print(f" --image {png_path} \\") + print(f" --mask {mask_path} \\") + print(f" --nifti {nifti_path}") + + return png_path, mask_path, nifti_path + + +def main(): + parser = argparse.ArgumentParser(description="Convert BraTS2020 h5 slices to PNG + NIfTI") + parser.add_argument("--volume", type=int, required=True, help="Volume ID (e.g. 100)") + parser.add_argument("--slice", type=int, default=None, help="Slice index (default: auto, picks slice with most tumour)") + parser.add_argument("--channel", type=int, default=1, help="MRI channel: 0=T1 1=T1ce 2=T2 3=FLAIR (default: 1=T1ce)") + parser.add_argument("--data_dir", default="data/BraTS2020/BraTS2020_training_data/data") + parser.add_argument("--out_dir", default="outputs/brats_test") + args = parser.parse_args() + + convert( + volume_id=args.volume, + data_dir=Path(args.data_dir), + out_dir=Path(args.out_dir), + slice_idx=args.slice, + channel=args.channel, + ) + + +if __name__ == "__main__": + main() diff --git a/utils/convert.py b/utils/convert_mat.py similarity index 100% rename from utils/convert.py rename to utils/convert_mat.py