src/ai_agent/utils/image_io.py

# utils/image_io.py
from __future__ import annotations
from pathlib import Path
import shutil
import time
import tempfile
import zipfile
import numpy as np
import imageio.v3 as iio
import pydicom
import nibabel as nib
from typing import Tuple, Dict, Any
from pydicom.pixels import apply_modality_lut, apply_voi_lut


def is_dicom_path(path: str | Path) -> bool:
    """Improved DICOM detection"""
    p = Path(path)
    if p.is_dir():
        # Check if directory contains any .dcm files
        return any(f.suffix.lower() == ".dcm" for f in p.rglob("*"))

    # For single files, do proper DICOM validation
    try:
        pydicom.dcmread(str(p), stop_before_pixels=True)
        return True
    except Exception:
        return False


def _safe_rmtree(p: Path) -> None:
    """Remove temp directory if it was created by us (safety guard)."""
    try:
        p = Path(p)
        troot = Path(tempfile.gettempdir())
        if (
            p.is_dir()
            and p.parent == troot
            and (p.name.startswith("dicom_zip_") or p.name.startswith("preview_"))
        ):
            shutil.rmtree(p, ignore_errors=True)
    except Exception:
        pass


def _cleanup_old_dicom_zips(hours: int = 6) -> None:
    """Cleanup stale dicom_zip_* temp folders older than `hours`. Best-effort."""
    troot = Path(tempfile.gettempdir())
    cutoff = time.time() - hours * 3600
    try:
        for d in troot.glob("dicom_zip_*"):
            try:
                if d.is_dir() and d.stat().st_mtime < cutoff:
                    shutil.rmtree(d, ignore_errors=True)
            except Exception:
                pass
    except Exception:
        pass


def maybe_unzip(path: str | Path) -> Path:
    """Safely extract zip file to temp directory, with better error handling."""
    p = Path(path)
    if p.is_dir() or p.suffix.lower() != ".zip":
        return p

    try:
        _cleanup_old_dicom_zips(hours=6)

        tmp = Path(tempfile.mkdtemp(prefix="dicom_zip_"))
        with zipfile.ZipFile(p) as z:
            # Check if zip contains DICOM files
            has_dicom = any(name.lower().endswith(".dcm") for name in z.namelist())
            if not has_dicom:
                raise ValueError("ZIP file contains no DICOM files")

            # Extract with path sanitization
            for item in z.namelist():
                if ".." not in item:  # Basic path traversal protection
                    z.extract(item, tmp)
        return tmp
    except Exception as e:
        raise ValueError(f"Failed to process ZIP file: {str(e)}")


def load_nifti(path: str | Path) -> Tuple[np.ndarray, Dict[str, Any]]:
    img = nib.load(str(path))
    data = img.get_fdata(dtype=np.float32)
    hdr = img.header
    zooms = tuple(float(z) for z in hdr.get_zooms())
    return data, {
        "format": "NIfTI",
        "shape": data.shape,
        "zooms": zooms,
        "datatype": str(hdr.get_data_dtype()),
    }


def load_dicom_series(dir_or_file: str | Path) -> Tuple[np.ndarray, Dict[str, Any]]:
    root = Path(dir_or_file)
    if root.is_file():
        # Could be a single multi-frame file; keep it as-is
        files = [root]
    else:
        files = sorted([p for p in root.rglob("*") if p.is_file()])

    dsets = []
    for p in files:
        try:
            ds = pydicom.dcmread(str(p), stop_before_pixels=False, force=True)
            if hasattr(ds, "PixelData"):
                dsets.append(ds)
        except Exception:
            continue
    if not dsets:
        raise ValueError("No DICOM slices with pixel data found.")

    # If the first file is multi-frame, load frames from it
    first = dsets[0]

    def _prep_pixels(ds):
        arr = ds.pixel_array  # may be (frames, rows, cols) or (rows, cols)
        # Apply LUTs/windowing for proper display range
        try:
            arr = apply_modality_lut(arr, ds)
        except Exception:
            pass
        try:
            arr = apply_voi_lut(arr, ds)
        except Exception:
            pass
        arr = arr.astype(np.float32)

        # Handle MONOCHROME1 (invert)
        if getattr(ds, "PhotometricInterpretation", "").upper() == "MONOCHROME1":
            # invert per frame
            arr = arr.max() - arr

        return arr

    if getattr(first, "NumberOfFrames", None):
        arr = _prep_pixels(first)
        # ensure (H, W, Z)
        if arr.ndim == 3:  # (frames, rows, cols)
            arr = np.transpose(arr, (1, 2, 0))
        elif arr.ndim == 2:  # (rows, cols)
            arr = arr[..., None]
        vol = arr
        dsets = [first]  # meta from first
    else:
        # classic per-slice series
        def sort_key(ds):
            inst = getattr(ds, "InstanceNumber", 0)
            try:
                inst = int(inst)
            except Exception:
                inst = 0
            ipp = getattr(ds, "ImagePositionPatient", None)
            z = (
                float(ipp[2])
                if (isinstance(ipp, (list, tuple)) and len(ipp) == 3)
                else 0.0
            )
            return (inst, z)

        dsets.sort(key=sort_key)
        frames = [_prep_pixels(ds) for ds in dsets]

        # each frame is (H,W); make (H,W,Z)
        vol = np.stack([f if f.ndim == 2 else f[0] for f in frames], axis=-1).astype(
            np.float32
        )

    # Normalize to [0,1] for consistent downstream PNG saving
    vmin = np.nanmin(vol)
    vmax = np.nanmax(vol)
    if np.isfinite(vmax) and vmax > vmin:
        vol = (vol - vmin) / (vmax - vmin)
    else:
        vol = np.zeros_like(vol, dtype=np.float32)

    # Spacing
    pxsp = getattr(dsets[0], "PixelSpacing", None)
    if not (isinstance(pxsp, (list, tuple)) and len(pxsp) == 2):
        # XA often uses ImagerPixelSpacing instead
        pxsp = getattr(dsets[0], "ImagerPixelSpacing", None)
    sy, sx = (float(pxsp[0]), float(pxsp[1])) if pxsp else (1.0, 1.0)

    # Slice spacing
    if vol.shape[-1] > 1:
        try:
            # use IPP difference if available
            z0 = float(getattr(dsets[0], "ImagePositionPatient", [0, 0, 0])[2])
            z1 = float(
                getattr(
                    dsets[min(1, len(dsets) - 1)], "ImagePositionPatient", [0, 0, 0]
                )[2]
            )
            sz = (
                abs(z1 - z0)
                if z1 != z0
                else float(getattr(dsets[0], "SliceThickness", 1.0))
            )
        except Exception:
            sz = float(getattr(dsets[0], "SliceThickness", 1.0))
    else:
        sz = float(getattr(dsets[0], "SliceThickness", 1.0))

    meta = {
        "format": "DICOM",
        "shape": vol.shape,  # (H, W, Z)
        "spacing": (sy, sx, sz),  # mm
        "Modality": getattr(dsets[0], "Modality", None),
        "BodyPartExamined": getattr(dsets[0], "BodyPartExamined", None),
        "SeriesDescription": getattr(dsets[0], "SeriesDescription", None),
        "SeriesInstanceUID": getattr(dsets[0], "SeriesInstanceUID", None),
        "StudyDescription": getattr(dsets[0], "StudyDescription", None),
        "PatientSex": getattr(dsets[0], "PatientSex", None),
        "PatientAge": getattr(dsets[0], "PatientAge", None),
        "PhotometricInterpretation": getattr(
            dsets[0], "PhotometricInterpretation", None
        ),
        "NumberOfFrames": getattr(dsets[0], "NumberOfFrames", None),
        "BitsStored": getattr(dsets[0], "BitsStored", None),
    }
    return vol.astype(np.float32), meta


def load_any(path: str | Path) -> Tuple[np.ndarray, Dict[str, Any]]:
    """Load any supported image format with better error handling and temp cleanup.

    If the input is a ZIP containing DICOMs, we extract to a temp dir and ensure
    it is deleted after loading.
    """
    temp_dir: Path | None = None
    try:
        p = maybe_unzip(path)
        # Track whether we created a temp extraction dir
        pt = Path(p)
        if (
            pt.is_dir()
            and pt.parent == Path(tempfile.gettempdir())
            and pt.name.startswith("dicom_zip_")
        ):
            temp_dir = pt

        # Handle DICOM
        if Path(p).is_dir() or is_dicom_path(p):
            data, meta = load_dicom_series(p)
            return data, meta

        # Handle NIfTI
        s = str(p).lower()
        if s.endswith(".nii") or s.endswith(".nii.gz"):
            return load_nifti(p)

        # Handle regular images
        arr = iio.imread(str(p))
        meta = {"format": Path(p).suffix.upper().lstrip("."), "shape": arr.shape}
        return arr.astype(np.float32), meta
    except Exception as e:
        raise ValueError(f"Failed to load image {path}: {str(e)}")
    finally:
        if temp_dir is not None:
            _safe_rmtree(temp_dir)