image-filter/examples/medical_io.py

"""Example-only local medical I/O; the filter library remains NumPy-only."""
from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import nibabel as nib
import numpy as np

SAMPLES = Path(__file__).resolve().parents[1] / 'data/samples'
CASE_NAMES = ('mri', 'pet', 'ct', 'spect', 'ultrasound', 'mr-dicom', 'ct-series')


@dataclass
class MedicalImage:
    pixels: np.ndarray
    modality: str
    format: str
    source: Path
    orientation: str
    files: int = 1


def bundled_cases() -> dict[str, tuple[Path, str]]:
    cases = {name.lower(): (SAMPLES / f'{name}.nii.gz', name) for name in ['MRI', 'PET', 'CT', 'SPECT', 'Ultrasound']}
    try:
        import pydicom
        root = Path(pydicom.__file__).parent / 'data/test_files'
        cases['mr-dicom'] = (root / 'MR_small.dcm', 'MR')
        cases['ct-series'] = (root / 'dicomdirtests/98892001/CT5N', 'CT')
    except ImportError:
        pass
    return cases


def _validate(pixels: np.ndarray) -> np.ndarray:
    # Only singleton extras are removed; time/channel dimensions are rejected.
    pixels = np.squeeze(pixels)
    if pixels.ndim not in (2, 3) or pixels.dtype.kind not in 'buif':
        raise ValueError('Only scalar 2D/3D medical images are supported; select time/channel externally')
    if not np.isfinite(pixels).all():
        raise ValueError('Medical image contains nonfinite values')
    return pixels


def load_medical(path: Path, modality: str = 'Unknown', series_uid: str | None = None) -> MedicalImage:
    path = path.resolve()
    if not path.exists():
        raise FileNotFoundError(f'Local medical fixture unavailable: {path}; no download attempted')
    if path.name.lower().endswith(('.nii', '.nii.gz')):
        source = nib.as_closest_canonical(nib.load(path))
        pixels = _validate(source.get_fdata())
        return MedicalImage(pixels, modality, 'NIfTI', path, 'RAS voxel axes')
    if path.is_dir() or path.suffix.lower() in ('.nrrd', '.nhdr'):
        import SimpleITK as sitk
        count = 1
        if path.is_dir():
            ids = sitk.ImageSeriesReader.GetGDCMSeriesIDs(str(path)) or ()
            if not ids:
                raise ValueError(f'No DICOM series in {path.name}')
            if series_uid is None and len(ids) != 1:
                raise ValueError('Multiple DICOM series found; select one with --series-uid')
            uid = series_uid or ids[0]
            if uid not in ids:
                raise ValueError('--series-uid not found in directory')
            files = sitk.ImageSeriesReader.GetGDCMSeriesFileNames(str(path), uid)
            if len(files) < 2:
                raise ValueError('DICOM series directory must contain at least two slices')
            import pydicom
            modality = str(pydicom.dcmread(files[0], stop_before_pixels=True).get('Modality', modality))
            reader = sitk.ImageSeriesReader()
            reader.SetFileNames(files)
            source = reader.Execute()
            count = len(files)
            format_name = 'DICOM series'
        else:
            source = sitk.ReadImage(str(path))
            format_name = 'NRRD'
        if source.GetNumberOfComponentsPerPixel() != 1:
            raise ValueError('Vector/color images are not supported in this medical demo')
        if source.GetDimension() == 3:
            source = sitk.DICOMOrient(source, 'RAS')
        pixels = sitk.GetArrayFromImage(source)
        # SimpleITK arrays are z,y,x; expose x,y,z like canonical nibabel.
        pixels = pixels.transpose(tuple(reversed(range(pixels.ndim))))
        return MedicalImage(_validate(pixels), modality, format_name, path, 'RAS voxel axes' if source.GetDimension() == 3 else 'native axes', count)
    if path.suffix.lower() == '.dcm':
        import pydicom
        from pydicom.pixels import apply_modality_lut
        source = pydicom.dcmread(path)
        if int(source.get('SamplesPerPixel', 1)) != 1 or int(source.get('NumberOfFrames', 1)) != 1:
            raise ValueError('Single DICOM accepts one monochrome frame; use a series directory for volumes')
        pixels = apply_modality_lut(source.pixel_array, source).astype(np.float64)
        if source.get('PhotometricInterpretation') == 'MONOCHROME1':
            pixels = pixels.max() + pixels.min() - pixels
        return MedicalImage(_validate(pixels), str(source.get('Modality', modality)), 'DICOM', path, 'native row/column')
    raise ValueError('Expected .nii/.nii.gz, .nrrd/.nhdr, .dcm, or a DICOM series directory')


def representative_plane(case: MedicalImage, axis: int = 2, index: int | None = None) -> tuple[np.ndarray, dict]:
    pixels = case.pixels
    selected = None
    if pixels.ndim == 3:
        if index is not None and not 0 <= index < pixels.shape[axis]:
            raise ValueError(f'Slice index outside axis length {pixels.shape[axis]}')
        if index is None:
            # Ignore border slices; maximize contrast among the central 80%.
            lo = int(pixels.shape[axis] * .1)
            hi = max(lo + 1, int(np.ceil(pixels.shape[axis] * .9)))
            scores = np.std(pixels, axis=tuple(i for i in range(3) if i != axis))
            selected = lo + int(np.argmax(scores[lo:hi]))
        else:
            selected = index
        plane = np.rot90(np.take(pixels, selected, axis=axis))
    else:
        if index is not None:
            raise ValueError('--slice-index requires a volume')
        plane = pixels
    low, high = np.percentile(plane, (1, 99))
    if high <= low:
        low, high = float(plane.min()), float(plane.max())
    if high <= low:
        raise ValueError('Selected slice has no intensity variation')
    image = np.clip((plane - low) / (high - low), 0, 1).astype(np.float64)
    info = {'modality': case.modality, 'format': case.format, 'source': case.source.name,
            'volume_shape': list(pixels.shape), 'orientation': case.orientation,
            'slice_axis': axis if selected is not None else None, 'slice_index': selected,
            'window': [float(low), float(high)], 'files': case.files,
            'input_shape': list(image.shape), 'processing': 'all filters independently on selected 2D plane'}
    return image, info