core/processing.py

"""format-hub — a general image file-format converter + a routing hub.

Two jobs:
  1. CONVERT: read just about any image/array a user throws at us (PNG/JPG/TIFF
     stacks/BMP/WEBP/.npy) and write it back out in any requested format, with
     optional grayscale and bit-depth control. This is the format on-ramp for
     every other app in the repo (e.g. turn a stray .npy into the TIFF stack
     caiman-app expects).
  2. HUB: know the rest of the catalog. `catalog()` returns the routing table of
     every imaging-plaza app; `recommend()` inspects a file and suggests which
     downstream tool fits — so this app is a front door to the others.

Only this file (+ app metadata) is app-specific; core/io.py etc. are shared.
"""
from __future__ import annotations

import os
from typing import Any

import numpy as np
import tifffile
from PIL import Image

from core.io import download_to_tmp, new_tmp_path

# ------------------------------------------------------------------ formats ---

READ_EXTS = {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp", ".webp", ".npy",
             ".dcm", ".zip",          # .dcm = single slice/multiframe; .zip = DICOM series
             ".nii", ".nii.gz"}       # NIfTI volumes
OUT_FORMATS = ["png", "jpg", "tiff", "bmp", "webp", "npy"]
BIT_DEPTHS = ["auto", "8", "16"]

_EXT = {"png": ".png", "jpg": ".jpg", "jpeg": ".jpg", "tiff": ".tif",
        "tif": ".tif", "bmp": ".bmp", "webp": ".webp", "npy": ".npy"}
_MIME = {"png": "image/png", "jpg": "image/jpeg", "tiff": "image/tiff",
         "bmp": "image/bmp", "webp": "image/webp", "npy": "application/octet-stream"}


# ------------------------------------------------------------------ reading ---

def _resolve(file_url: Any) -> str:
    """Normalize any input form to a local path.

    The /process file param is typed `str`, but callers send it three ways: a plain
    URL/path string (MCP, curl), or a Gradio FileData dict (gradio_client upload).
    Handle all of them and download URLs.
    """
    obj = file_url
    if isinstance(obj, dict):  # Gradio FileData from an upload
        obj = obj.get("path") or obj.get("url") or ""
    s = str(obj)
    if s.startswith(("http://", "https://")):
        suffix = ".nii.gz" if s.lower().endswith(".nii.gz") else (os.path.splitext(s)[1] or ".bin")
        return download_to_tmp(s, suffix=suffix)
    return s


def _read_any(path: str) -> np.ndarray:
    """Read npy / TIFF stack / DICOM / NIfTI / any PIL image into an array."""
    from core.nifti import is_nifti, read_nifti
    if is_nifti(path):
        return read_nifti(path)
    ext = os.path.splitext(path)[1].lower()
    if ext == ".npy":
        return np.load(path, allow_pickle=False)
    if ext in (".tif", ".tiff"):
        return np.asarray(tifffile.imread(path))
    if ext in (".dcm", ".zip"):
        from core.dicom import read_dicom_any
        return read_dicom_any(path)
    return np.asarray(Image.open(path))


# --------------------------------------------------------------- conversion ---

def _to_gray(arr: np.ndarray) -> np.ndarray:
    if arr.ndim == 3 and arr.shape[-1] in (3, 4):
        from skimage.color import rgb2gray
        return rgb2gray(arr[..., :3])
    return arr


def _apply_bit_depth(arr: np.ndarray, bit_depth: str) -> np.ndarray:
    if bit_depth == "auto":
        return arr
    a = arr.astype(np.float32)
    mn, mx = float(a.min()), float(a.max())
    rng = max(mx - mn, 1e-8)
    if bit_depth == "8":
        return np.clip((a - mn) / rng * 255.0, 0, 255).astype(np.uint8)
    if bit_depth == "16":
        return np.clip((a - mn) / rng * 65535.0, 0, 65535).astype(np.uint16)
    return arr


def _coerce_for_pil(arr: np.ndarray, fmt: str) -> np.ndarray:
    """PIL-writable formats need uint8; jpg/bmp have no alpha channel."""
    a = arr
    if a.dtype != np.uint8:
        a = _apply_bit_depth(a, "8")
    if fmt in ("jpg", "jpeg", "bmp") and a.ndim == 3 and a.shape[-1] == 4:
        a = a[..., :3]
    return a


def convert(
    file_url: str,
    target_format: str = "png",
    to_gray: bool = False,
    bit_depth: str = "auto",
) -> tuple[str, dict]:
    """Convert an input file to `target_format`. Returns (output_path, report)."""
    fmt = target_format.lower().strip()
    if fmt not in _EXT:
        raise ValueError(f"unsupported target_format '{target_format}'; choose one of {OUT_FORMATS}")

    src = _resolve(file_url)
    arr = _read_any(src)
    in_shape, in_dtype = list(arr.shape), str(arr.dtype)
    is_stack = arr.ndim == 3 and arr.shape[-1] not in (3, 4)

    if to_gray:
        arr = _to_gray(arr)
    if bit_depth != "auto":
        arr = _apply_bit_depth(arr, bit_depth)

    out_path = new_tmp_path(f"converted{_EXT[fmt]}")

    if fmt == "npy":
        np.save(out_path, arr)
    elif fmt in ("tiff", "tif"):
        tifffile.imwrite(out_path, arr, compression="zlib")  # keeps stacks + 16-bit
    else:
        if is_stack:
            raise ValueError(
                f"input is a {in_shape} stack; '{fmt}' is single-frame — convert to "
                f"'tiff' or 'npy' to keep all frames"
            )
        Image.fromarray(_coerce_for_pil(arr, fmt)).save(out_path)

    out_arr = np.asarray(arr)
    report = {
        "input_shape": in_shape,
        "input_dtype": in_dtype,
        "is_stack": bool(is_stack),
        "target_format": fmt,
        "to_gray": bool(to_gray),
        "bit_depth": bit_depth,
        "output_shape": list(out_arr.shape),
        "output_dtype": str(out_arr.dtype),
        "output_bytes": os.path.getsize(out_path),
    }
    return out_path, report


def process(file_url: str, target_format: str = "png", to_gray: bool = False,
            bit_depth: str = "auto") -> str:
    """Contract entry point: convert and return the output file path."""
    out_path, _ = convert(file_url, target_format=target_format, to_gray=to_gray,
                           bit_depth=bit_depth)
    return out_path


# --------------------------------------------------------------------- hub ---

# input type -> downstream imaging-plaza app (mirrors use-imaging-plaza/catalog.md)
CATALOG: list[dict[str, Any]] = [
    {"input": "2D image — general scikit-image ops", "app": "skimage-classic", "port": 7860},
    {"input": "2D image — fibre/structure orientation", "app": "orientationpy-app", "port": 7861},
    {"input": "2D photo — remove background", "app": "rembg-app", "port": 7862},
    {"input": "2D slice — CT scan + reconstruction", "app": "tomo-recon", "port": 7863},
    {"input": "particle-image pair (2-frame TIFF) — flow", "app": "piv-app", "port": 7864},
    {"input": "2D noisy image — denoise", "app": "denoise-app", "port": 7865},
    {"input": "2D microscopy image — segment cells/nuclei", "app": "cellpose-app", "port": 7866},
    {"input": "camera-trap image (colour/IR) — wildlife detection", "app": "pyrvision-app", "port": 7875},
    {"input": "3D CT volume (Z,H,W) — baggage analysis", "app": "ct-baggage", "port": 7867},
    {"input": "top-down multi-fly video — behavior", "app": "fly-behavior", "port": 7868},
    {"input": "top-down animal video — markerless pose", "app": "deeplabcut-app", "port": 7869},
    {"input": "multi-camera montage video — 3D fly pose", "app": "deepfly3d-app", "port": 7870},
    {"input": "single-fly video — microbehavior + sleep", "app": "flyvista-app", "port": 7871},
    {"input": "2-photon calcium movie (T,H,W) — source extraction", "app": "caiman-app", "port": 7872},
    {"input": "fluorescence movie (T,H,W) — spatiotemporal event detection (AQuA)", "app": "aqua-app", "port": 7881},
    {"input": "fluorescence movie (T,H,W) — events + Consensus Functional Units (AQuA2)", "app": "aqua2-app", "port": 7882},
    {"input": "fluorescence movie (T,H,W) — self-supervised denoising (DeepInterpolation)", "app": "deepinterpolation-app", "port": 7883},
    {"input": "noisy 2D image — self-supervised denoising (Noise2Void)", "app": "noise2void-app", "port": 7884},
    {"input": "degraded 2D image — content-aware restoration (CSBDeep/CARE)", "app": "csbdeep-app", "port": 7885},
    {"input": "video or audio file — convert/analyze/extract-audio (ffmpeg)", "app": "media-hub", "port": 7886},
    {"input": "3D mesh (STL/OBJ/PLY) — wind/CFD", "app": "openfoam-windtunnel", "port": 7873},
    {"input": "georeferenced raster (GeoTIFF/GDAL) — reproject/resample/crop", "app": "rasterio-geo", "port": 7876},
    {"input": "vector dataset (GeoJSON/Shapefile/GPKG) — reproject/buffer/clip", "app": "vector-geo", "port": 7877},
    {"input": "LiDAR point cloud (LAS/LAZ) — crop/downsample/reproject", "app": "pointcloud-geo", "port": 7878},
    {"input": "DEM / Cloud-Optimized GeoTIFF — fetch/window by bbox + hillshade", "app": "stac-dem", "port": 7879},
    {"input": "3D mesh (OBJ/PLY/STL/GLB/OFF) — convert/decimate/scale", "app": "mesh-geo", "port": 7880},
]


def catalog() -> dict:
    """Return the routing table of every imaging-plaza app (the hub directory)."""
    return {
        "hub": "format-hub",
        "convert_formats": OUT_FORMATS,
        "apps": CATALOG,
        "note": "Convert your file here, then send it to the app for your input type. "
                "Each app's /process accepts a URL directly (no upload).",
    }


def dicom_info(file_url: str) -> dict:
    """Report the study/series/segmentation structure of a DICOM .dcm or .zip
    (a series, a multi-series study, a DICOMDIR, or a SEG)."""
    from core.dicom import dicom_overview
    return dicom_overview(_resolve(file_url))


def recommend(file_url: str) -> dict:
    """Inspect a file and suggest which downstream app(s) fit its shape."""
    src = _resolve(file_url)
    arr = _read_any(src)
    nd = arr.ndim
    rgb = nd == 3 and arr.shape[-1] in (3, 4)
    stack = nd == 3 and not rgb
    if nd == 2 or rgb:
        kind = "a single 2D image"
        apps = ["skimage-classic", "orientationpy-app", "rembg-app", "denoise-app",
                "cellpose-app", "tomo-recon"]
    elif stack and arr.shape[0] == 2:
        kind = "a 2-frame pair"
        apps = ["piv-app"]
    elif stack:
        kind = f"a {arr.shape[0]}-frame stack / volume"
        apps = ["caiman-app", "ct-baggage", "fly-behavior", "deeplabcut-app",
                "flyvista-app", "deepfly3d-app"]
    else:
        kind = f"an array of shape {list(arr.shape)}"
        apps = []
    by = {c["app"]: c for c in CATALOG}
    return {
        "detected": kind,
        "shape": list(arr.shape),
        "dtype": str(arr.dtype),
        "suggested_apps": [by[a] for a in apps if a in by],
    }


# ---- one-shot helper used by the UI / smoke test ----

def simulate_full(file_url: str, target_format: str = "png", to_gray: bool = False,
                  bit_depth: str = "auto") -> dict:
    out_path, report = convert(file_url, target_format=target_format, to_gray=to_gray,
                               bit_depth=bit_depth)
    return {"out_path": out_path, "report": report, "mime": _MIME[report["target_format"]]}