_utils/image_io.py

"""Image standardization.

``standardize_image`` collapses any microscopy TIFF variant (16-bit/float,
multi-channel, multi-page Z/time stacks) to one canonical 8-bit RGB PNG, used
for both the preview and the model.
"""

import os
import tempfile
from contextlib import contextmanager
from typing import NamedTuple

import numpy as np
from PIL import Image

try:
    import tifffile
except ImportError:  # pragma: no cover - tifffile is a project dependency
    tifffile = None

# The model resizes any input to 512x512, so a larger image loses detail rather
# than adding any.
RECOMMENDED_SIZE = 512
WARN_SIZE = 3072
MAX_SIZE = 4096
# Warn below this: a shorter side under 32px means >16x upscaling to 512, so
# there is almost no real detail for the model to work with.
MIN_SIZE = 32
# An integer image is a low-dynamic-range candidate when its full span is tiny
# relative to the dtype's max. Float has no fixed range, so it is skipped.
NARROW_RANGE_FRAC = 0.02
# Sparse fluorescence can have a narrow span but still be valid; only flag it
# when the max is not meaningfully above the 99th percentile.
BRIGHT_TAIL_FRAC = 0.1
# Measured on the uncompressed array, not the file size on disk: compression
# makes disk size a poor proxy for what opening the file actually allocates.
MAX_READ_BYTES = 300 * 1024 ** 2

# Axis roles, per tifffile's `series.axes` naming. Anything else (Z, T, ...) is a
# stack axis, indexed by frame.
_CHANNEL = ("C", "S")      # C = separate channel planes, S = interleaved RGB samples
_UNKNOWN = ("Q", "I")      # file named no axis; only these may be *guessed* as channels


TIFF_EXTENSIONS = (
    sorted(
        {".ome.tif", ".ome.tiff"}
        | {"." + e for e in tifffile.TIFF.FILE_EXTENSIONS if "." not in e}
    )
    if tifffile is not None else [".ome.tif", ".ome.tiff", ".tif", ".tiff"]
)


def _series_planes(series):
    """Number of 2D planes in a series: the product of every axis that is not
    spatial (Y/X) or RGB samples (S). Axes-aware on purpose - a positional
    shape[:-2] would miscount RGB, which tifffile stores as a trailing S axis."""
    planes = int(np.prod([
        dim for axis, dim in zip(str(series.axes), series.shape)
        if axis not in ("Y", "X", "S")
    ]))
    return planes or 1


def _ome_needs_single_file(tif):
    """True for a multi-file OME-TIFF whose logical series spans more planes than
    this file holds: its data lives in sibling files that are not part of a
    single upload, so the assembled series would zero-fill the out-of-file
    planes. A self-contained OME has all its planes in-file (planes == pages)."""
    try:
        return bool(tif.is_ome) and _series_planes(tif.series[0]) > len(tif.pages)
    except Exception:  # noqa: BLE001 - detection must never break the read
        return False


@contextmanager
def _open_tiff(path):
    """Open a TIFF; for a multi-file OME set whose series spans beyond this file,
    reopen just this file's own pages (is_ome=False) so out-of-file planes are
    not zero-filled. Detection is lazy (shape/pages only, no pixel decode)."""
    tif = tifffile.TiffFile(path)
    try:
        if _ome_needs_single_file(tif):
            tif.close()
            tif = tifffile.TiffFile(path, is_ome=False)
        yield tif
    finally:
        tif.close()


def _read_tiff(path):
    """Read a TIFF-family file into (array, axes). Raises if not readable."""
    with _open_tiff(path) as tif:
        series = tif.series[0]
        page = tif.pages[0]
        arr = np.asarray(series.asarray())
        axes = str(series.axes)

        is_palette = (getattr(page, "photometric", None) == tifffile.PHOTOMETRIC.PALETTE
                      and not any(a in _CHANNEL for a in axes))
        if is_palette and page.colormap is not None:
            colormap = np.asarray(page.colormap)  # (3, 2**bits), uint16
            rgb = np.moveaxis(colormap[:, arr], 0, -1)  # (..., H, W, 3)
            # TIFF colormaps are 16-bit; scale down to 8-bit (65535/255=257).
            arr = np.round(rgb / 257.0).astype(np.uint8)
            axes = axes + "S"  # the LUT added an RGB sample axis
        return arr, axes


def _read_array(path):
    """Load an image file into (array, axes) where axes names each dimension.

    axes uses tifffile's convention ('C' channel, 'S' RGB samples, 'Z'/'T'
    stack, 'Y'/'X' spatial, 'Q' unknown). Indexed / palette images (ImageJ
    "8-bit Color", palette PNG/GIF) are expanded through their color lookup
    table so we return true RGB, not bare indices.
    """
    if tifffile is not None:
        try:
            return _read_tiff(path)
        except Exception:  # noqa: BLE001 - not a TIFF, or tifffile cannot parse it
            pass

    img = Image.open(path)
    if img.mode in ("P", "PA"):
        img = img.convert("RGB")
    arr = np.asarray(img)
    return arr, ("YXS" if arr.ndim == 3 else "YX")


def _shape_axes(path):
    """Return (shape, axes) from metadata only - no pixel decode."""
    if tifffile is not None:
        try:
            with _open_tiff(path) as tif:
                series = tif.series[0]
                return tuple(series.shape), str(series.axes)
        except Exception:  # noqa: BLE001 - fall through to PIL
            pass
    with Image.open(path) as img:
        w, h = img.size
        if img.mode in ("P", "PA", "RGB", "RGBA", "CMYK", "YCbCr", "LAB", "HSV"):
            return (h, w, len(img.getbands())), "YXS"
        return (h, w), "YX"


_PIL_BITS = {"1": 1, "L": 8, "P": 8, "LA": 8, "PA": 8, "RGB": 8, "RGBA": 8,
             "CMYK": 8, "YCbCr": 8, "LAB": 8, "HSV": 8,
             "I;16": 16, "I;16B": 16, "I;16L": 16, "I": 32, "F": 32}


def bit_depth(path):
    """Bits per sample, read from metadata only (0 if it cannot be determined)."""
    if tifffile is not None:
        try:
            with _open_tiff(path) as tif:
                return int(np.dtype(tif.series[0].dtype).itemsize) * 8
        except Exception:  # noqa: BLE001 - fall through to PIL
            pass
    try:
        with Image.open(path) as img:
            return int(_PIL_BITS.get(img.mode, 0))
    except Exception:  # noqa: BLE001
        return 0


def _plan_axes(shape, axes):
    """Drop size-1 axes, infer a channel axis when the file names none, and move
    it last.

    Returns (shape, axes, guessed) as lists + a flag saying whether the channel
    axis had to be inferred rather than read from the file.
    """
    axes = list(axes)
    shape = list(shape)
    if len(axes) != len(shape):  # be defensive; keep the trailing spatial axes
        axes = ["Q"] * (len(shape) - 2) + ["Y", "X"]

    pairs = [(a, d) for a, d in zip(axes, shape) if d != 1]
    axes = [a for a, _ in pairs]
    shape = [d for _, d in pairs]

    guessed = False
    if not any(a in _CHANNEL for a in axes):
        # Only guess on axes the file left unnamed. An axis the file explicitly
        # calls Z/T is a stack even when its length happens to be 3.
        for i, (a, d) in enumerate(zip(axes, shape)):
            if a in _UNKNOWN and d in (2, 3, 4):
                axes[i] = "C"
                guessed = True
                break

    ci = next((i for i, a in enumerate(axes) if a in _CHANNEL), None)
    if ci is not None:
        axes.append(axes.pop(ci))
        shape.append(shape.pop(ci))
    return shape, axes, guessed


def _reduce_to_hwc(arr, axes, frame=0, sub_frame=None):
    """Collapse a labelled array to 2D (H, W) or 3D (H, W, C).

    The channel axis (named by the file, or inferred by ``_plan_axes`` only when
    the file names none) is moved last and kept whole. Stack axes (Z / time /
    page) are handled in order:
      * the first is indexed by ``frame`` (the frame slider);
      * the second is indexed by ``sub_frame`` if given, else collapsed by a
        maximum-intensity projection (the natural default for a focal stack -
        keeps the brightest signal across planes rather than an arbitrary one);
      * any deeper ones are projected too.
    """
    axes = list(axes)

    for i in range(len(axes) - 1, -1, -1):
        if arr.shape[i] == 1:
            arr = arr.reshape(arr.shape[:i] + arr.shape[i + 1:])
            axes.pop(i)

    _, planned, _ = _plan_axes(arr.shape, axes)

    if "C" in planned and not any(a in _CHANNEL for a in axes):
        for i, (a, d) in enumerate(zip(axes, arr.shape)):
            if a in _UNKNOWN and d in (2, 3, 4):
                axes[i] = "C"
                break
    ci = next((i for i, a in enumerate(axes) if a in _CHANNEL), None)
    if ci is not None:
        arr = np.moveaxis(arr, ci, -1)
        axes.append(axes.pop(ci))

    target = 3 if ci is not None else 2
    stack = 0
    while len(axes) > target:
        if stack == 0:
            idx = max(0, min(int(frame), arr.shape[0] - 1))
            arr = arr[idx]
        elif stack == 1 and sub_frame is not None:
            idx = max(0, min(int(sub_frame), arr.shape[0] - 1))
            arr = arr[idx]
        else:
            arr = arr.max(axis=0)  # maximum-intensity projection over this axis
        axes.pop(0)
        stack += 1

    return arr


# tifffile axis letters -> words a microscopist uses, for messages and slider
# labels. Anything unmapped (an unnamed 'Q'/'I' axis, a raw page axis) is just a
# "frame".
_AXIS_LABELS = {"T": "timepoint", "Z": "z-plane"}


def _axis_label(a):
    return _AXIS_LABELS.get(a, "frame")


class ImageInfo(NamedTuple):
    """How a file's dimensions were interpreted (read from metadata only).

    frames    - length of the first stack (Z/T) axis, 1 if not a stack
    channels  - length of the channel axis, 1 if single-channel
    axes      - the file's own axes string, e.g. 'CZYX' ('Q' = unnamed)
    guessed   - True if the channel axis was inferred rather than read
    shape     - the file's raw shape as stored, e.g. (1, 4, 1, 1024, 1024)
    width     - pixels along the X axis (0 if unknown)
    height    - pixels along the Y axis (0 if unknown)
    sub_frames- length of a *second* stack axis (e.g. Z in a time+Z file), 1 if
                there is only one stack axis
    frame_label / sub_label - the words for those two axes ('timepoint', ...)
    """
    frames: int
    channels: int
    axes: str
    guessed: bool
    shape: tuple
    width: int
    height: int
    sub_frames: int = 1
    frame_label: str = "frame"
    sub_label: str = "frame"


def inspect_image(path):
    """Describe a file's structure from metadata only (no pixel decode)."""
    try:
        shape, axes = _shape_axes(path)
        planned_shape, planned_axes, guessed = _plan_axes(shape, axes)
        has_c = bool(planned_axes) and planned_axes[-1] in _CHANNEL
        channels = int(planned_shape[-1]) if has_c else 1

        stack = [(a, d) for a, d in zip(planned_axes, planned_shape)
                 if a not in ("Y", "X") and a not in _CHANNEL]
        frames = int(stack[0][1]) if stack else 1
        frame_label = _axis_label(stack[0][0]) if stack else "frame"
        sub_frames = int(stack[1][1]) if len(stack) > 1 else 1
        sub_label = _axis_label(stack[1][0]) if len(stack) > 1 else "frame"

        # Spatial size comes from the named Y/X axes, so it stays correct
        # whatever order the other axes are in.
        width = height = 0
        for a, d in zip(axes, shape):
            if a == "Y":
                height = int(d)
            elif a == "X":
                width = int(d)
        if not (width and height):  # no Y/X named; fall back to the header probe
            width, height = image_size(path)

        return ImageInfo(frames, channels, str(axes), guessed, tuple(shape),
                         width, height, sub_frames, frame_label, sub_label)
    except Exception:  # noqa: BLE001
        return ImageInfo(1, 1, "", False, (), 0, 0)


class PixelReport(NamedTuple):
    """Pixel-level sanity checks on decoded image data (see ``pixel_stats``).

    decoded   - False if the pixels could not be read (corrupt / truncated file)
    finite    - False if the frame contains NaN/inf (only checked for float data)
    vmin/vmax - min and max pixel value of the frame
    dtype_max - np.iinfo(dtype).max for integer data, else 0.0
    low_range - True if an integer frame's robust span is a tiny fraction of the
                dtype range (low contrast / narrow dynamic range)
    """
    decoded: bool
    finite: bool
    vmin: float
    vmax: float
    dtype_max: float
    low_range: bool


def pixel_stats(path, frame=0, sub_frame=None):
    """Decode image pixels and report problems (blank / non-finite / corrupt /
    low dynamic range).

    Works on the raw reduced frame, not the RGB-padded one, so a 2-channel
    image's zero-padded third channel does not skew the min/max/finite stats.
    Only call once the header size guard has passed, so decoding is bounded.
    """
    try:
        raw, axes = _read_array(path)
        arr = _reduce_to_hwc(raw, axes, frame=frame, sub_frame=sub_frame)
    except Exception:  # noqa: BLE001 - unreadable pixels
        return PixelReport(False, True, 0.0, 0.0, 0.0, False)

    if np.issubdtype(arr.dtype, np.floating) and not np.isfinite(arr).all():
        return PixelReport(True, False, 0.0, 0.0, 0.0, False)

    vmin, vmax = float(arr.min()), float(arr.max())
    is_int = np.issubdtype(arr.dtype, np.integer)
    dtype_max = float(np.iinfo(arr.dtype).max) if is_int else 0.0
    low_range = False
    if is_int and vmax > vmin and dtype_max > 0:
        narrow = (vmax - vmin) < NARROW_RANGE_FRAC * dtype_max
        if narrow:
            # Among narrow images, spare fluorescence (dark background + sparse
            # bright objects) by checking for a bright tail: its max sits far
            # above the 99th percentile. A dense narrow band (values packed in a
            # thin range) has none, so only that is flagged as low dynamic range.
            p99 = float(np.percentile(arr, 99))
            bright_tail = (vmax - p99) / (vmax - vmin)
            low_range = bright_tail < BRIGHT_TAIL_FRAC
    return PixelReport(True, True, vmin, vmax, dtype_max, low_range)



def array_nbytes(path):
    """Bytes that opening this file's full array would allocate.

    Computed from shape + dtype in the header - no pixel decode - so it is safe
    to call on a file that is too large to open. Returns 0 if unknown.
    """
    if tifffile is not None:
        try:
            with _open_tiff(path) as tif:
                series = tif.series[0]
                return int(np.prod(series.shape)) * int(np.dtype(series.dtype).itemsize)
        except Exception:  # noqa: BLE001 - fall through to PIL
            pass
    try:
        with Image.open(path) as img:
            w, h = img.size
            return int(w) * int(h) * len(img.getbands())  # 8-bit assumption
    except Exception:  # noqa: BLE001
        return 0


def image_size(path):
    """Return an image's (width, height) by reading only its header.

    Never decodes pixel data, so this is safe to call as a size guard on a file
    that would be too large to load. Returns (0, 0) if the size cannot be
    determined, so callers treat it as "unknown" and proceed.
    """
    if tifffile is not None:
        try:
            with tifffile.TiffFile(path) as tif:
                page = tif.pages[0]
                return int(page.imagewidth), int(page.imagelength)
        except Exception:  # noqa: BLE001 - fall through to PIL
            pass
    try:
        with Image.open(path) as img:  # PIL parses the header lazily
            return int(img.size[0]), int(img.size[1])
    except Exception:  # noqa: BLE001
        return 0, 0


def _to_rgb(arr):
    """Turn a 2D or (H, W, C) array into exactly 3 channels."""
    if arr.ndim == 2:
        return np.stack([arr] * 3, axis=-1)

    channels = arr.shape[2]
    if channels == 1:
        return np.repeat(arr, 3, axis=2)
    if channels == 2:
        # Pad a zero third channel rather than inventing signal.
        return np.concatenate([arr, np.zeros_like(arr[:, :, :1])], axis=2)
    return arr[:, :, :3]


def _load_rgb(path, frame=0, sub_frame=None):
    """Read a file and reduce it to an (H, W, 3) array plus its source dtype."""
    raw, axes = _read_array(path)
    arr = _reduce_to_hwc(raw, axes, frame=frame, sub_frame=sub_frame)
    arr = _to_rgb(arr)
    return arr, raw.dtype


def _to_uint8(arr, stretch):
    """Map pixel values to uint8.

    When ``stretch`` is True (non-8-bit input) a 1st-99th percentile auto-
    contrast stretch is applied - outlier-robust, so a hot/saturated pixel does
    not collapse the visible signal to black. Otherwise values are only clipped,
    so standard 8-bit images are unchanged.
    """
    arr = arr.astype(np.float32)

    if not stretch:
        return np.clip(arr, 0, 255).astype(np.uint8)

    lo, hi = np.percentile(arr, (1, 99))
    if hi <= lo:  # near-flat image; fall back to full min/max
        lo, hi = float(arr.min()), float(arr.max())
    if hi <= lo:  # truly constant image
        return np.zeros(arr.shape, dtype=np.uint8)

    arr = np.clip((arr - lo) / (hi - lo), 0.0, 1.0)
    return (arr * 255.0).astype(np.uint8)


def _save_png(arr, out_path=None, out_dir=None, base=None, suffix="_std.png"):
    """Save an (H, W, 3) uint8 array as a PNG and return the path."""
    if out_path is None:
        if out_dir is not None:
            os.makedirs(out_dir, exist_ok=True)
            out_path = os.path.join(out_dir, (base or "image") + suffix)
        else:
            tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
            out_path = tmp.name
            tmp.close()
    Image.fromarray(arr, mode="RGB").save(out_path)
    return out_path


def standardize_image(path, frame=0, sub_frame=None, out_dir=None, out_path=None):
    """Standardize any image/TIFF to a canonical 8-bit RGB PNG.

    Used for both the preview and the model: one frame, <=3 channels, with a
    1st-99th percentile auto-contrast stretch for 16-bit/float input (8-bit is
    passed through unchanged).

    Args:
        path: path to the input image (TIFF, PNG, JPG, ...).
        frame: which frame of the first stack axis to use (0-based; ignored for
            non-stacks).
        sub_frame: which plane of a second stack axis (e.g. Z in a time+Z file)
            to use (0-based). None means combine those planes by a maximum-
            intensity projection.
        out_dir: optional directory for the output PNG.
        out_path: optional explicit output file path (overrides out_dir).

    Returns:
        Path to the standardized PNG. On any failure the original ``path`` is
        returned unchanged so callers degrade gracefully.
    """
    try:
        arr, dtype = _load_rgb(path, frame=frame, sub_frame=sub_frame)
        arr = _to_uint8(arr, stretch=(dtype != np.uint8))
        base = os.path.splitext(os.path.basename(path))[0]
        return _save_png(arr, out_path=out_path, out_dir=out_dir, base=base, suffix="_std.png")
    except Exception as e:  # noqa: BLE001 - never let preprocessing crash a run
        print(f"⚠️ standardize_image failed for {path}: {e}; using original file")
        return path