intake.py

"""Upload intake: normalize NIfTI / NRRD / MetaImage / DICOM series to .nii.gz.

Downstream (RADAR_inference DataFolder/evaluate) only ever sees one .nii.gz,
so the vendored upstream chain stays untouched. No new dependencies:
SimpleITK and nibabel are already required.
"""

import os
VOLUME_SUFFIXES = (".nii", ".nii.gz", ".nrrd", ".mha")  # self-contained single files only:
HEADER_SUFFIXES = (".nhdr", ".mhd", ".hdr", ".img", ".raw")  # header+raw pairs need conversion first

import gradio as gr


def _to_nifti(path: str, workdir: str) -> str:
    """Normalize one single-file volume to .nii.gz via SimpleITK."""
    import SimpleITK as sitk

    low = path.lower()
    if low.endswith((".nii", ".nii.gz")):
        import nibabel as nib

        try:
            nib.load(path)
        except Exception:
            raise gr.Error("upload a .nii or .nii.gz CT volume")
        return path
    if not low.endswith((".nrrd", ".mha")):
        raise gr.Error("upload a .nii.gz, .nrrd, .mha volume, or a DICOM series")
    try:
        img = sitk.ReadImage(path)
    except Exception:
        raise gr.Error(f"could not read {os.path.basename(path)} as an image volume")
    if img.GetDimension() != 3 or min(img.GetSize()) < 1:
        raise gr.Error(f"{os.path.basename(path)} is not a 3D volume")
    out = os.path.join(workdir, "case.nii.gz")
    sitk.WriteImage(img, out)
    return out


def _series_to_nifti(staged_dir: str, workdir: str) -> str:
    """Assemble a DICOM series directory to .nii.gz (largest series wins).

    Walks subdirectories too: zipped PACS studies nest one folder per series.
    """
    import SimpleITK as sitk

    pooled: dict = {}
    for root, _dirs, files in os.walk(staged_dir):
        if not files:
            continue
        try:
            series_ids = sitk.ImageSeriesReader.GetGDCMSeriesIDs(root)
        except Exception:
            continue
        for sid in series_ids:
            fnames = sitk.ImageSeriesReader.GetGDCMSeriesFileNames(root, sid)
            if len(fnames) > len(pooled.get(sid, [])):
                pooled[sid] = fnames
    if not pooled:
        raise gr.Error("no DICOM series found: upload the full series (.dcm files or a .zip)")
    fnames = max(pooled.values(), key=len)
    try:
        reader = sitk.ImageSeriesReader()
        reader.SetFileNames(fnames)
        img = reader.Execute()
    except Exception:
        raise gr.Error("could not assemble the DICOM series (mixed series or corrupt slices?)")
    if img.GetDimension() != 3:
        raise gr.Error("DICOM series is not a single 3D volume")
    if img.GetSize()[2] < 8:
        raise gr.Error(
            f"DICOM series too thin ({img.GetSize()[2]} slices): upload the full abdominal series"
        )
    out = os.path.join(workdir, "case_dcm.nii.gz")
    sitk.WriteImage(img, out)
    return out


def _stage_uploads(paths: list, staged_dir: str) -> None:
    """Copy uploads into staging; extract .zips preserving inner folders.

    Member paths are sanitized (no absolute paths, no ``..``); same-named
    files (e.g. nested PACS series with identical slice names) get a
    counter suffix instead of overwriting each other.
    """
    import shutil
    import zipfile

    def _place(data: bytes | None, src: str, rel: str) -> None:
        parts = [p for p in rel.replace("\\", "/").split("/") if p not in ("", ".")]
        if not parts or ".." in parts or os.path.isabs(rel):
            return
        dest = os.path.join(staged_dir, *parts)
        os.makedirs(os.path.dirname(dest), exist_ok=True)
        stem, ext = os.path.splitext(dest)
        n = 0
        while os.path.exists(dest):
            n += 1
            dest = f"{stem}_{n}{ext}"
        if data is not None:
            with open(dest, "wb") as dst:
                dst.write(data)
        else:
            shutil.copy(src, dest)

    for p in paths:
        if p.lower().endswith(".zip"):
            try:
                with zipfile.ZipFile(p) as zf:
                    members = [m for m in zf.namelist() if not m.endswith("/")]
            except zipfile.BadZipFile:
                raise gr.Error(f"{os.path.basename(p)} is not a readable .zip")
            with zipfile.ZipFile(p) as zf:
                for member in members:
                    with zf.open(member) as src:
                        _place(src.read(), p, member)
        else:
            _place(None, p, os.path.basename(p))