_utils/image_io.py

"""Image standardization.

Microscopy TIFFs come in many flavors: 16-bit or float pixels, single or
multi channel, and multi-page Z / time stacks. ``standardize_image`` collapses
any of those variants to a single canonical 8-bit RGB PNG used for BOTH the
on-screen preview and the model, so what you see is what gets segmented.

Which axis means what is read from the file's own metadata (``series.axes``
from tifffile: 'C'=channel, 'S'=RGB samples, 'Z'/'T'=stack, 'Y'/'X'=spatial)
rather than guessed from the array shape -- guessing cannot tell a 3-channel
(C, Y, X) image apart from a 3-slice (Z, Y, X) stack. Only when a file names
no axes at all (tifffile reports 'Q' = unknown) do we fall back to the dominant
convention: the first unknown axis of size 2-4 is the channel axis.

Reduction policy:

  * multi-page / Z / time stacks -> one frame (default: the first)
  * more than 3 channels         -> first 3 channels (as R, G, B)
  * intensity -> 0-255           -> 1st-99th percentile auto-contrast for
                                    16-bit/float input (outlier-robust: a hot /
                                    saturated pixel would make plain min/max
                                    scaling collapse the real signal to black)

Standard 8-bit inputs (PNG/JPG/8-bit TIFF) are passed through unchanged in
value; only channel layout is normalized.
"""

import os
import tempfile
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

# Size policy. The model resizes any input to 512x512 (see segmentation.py), so a
# larger image costs memory/time and *loses* detail rather than adding any.
RECOMMENDED_SIZE = 512   # cropping to about this preserves the most detail
WARN_SIZE = 3072         # above this, advise the user to crop
MAX_SIZE = 4096          # above this, refuse: reading the pixels risks an OOM
# Refuse files whose full array would not comfortably fit in memory. Measured on
# the uncompressed array (shape x dtype) rather than the file size on disk, since
# compression makes disk size a poor proxy for what we actually allocate.
MAX_READ_BYTES = 300 * 1024 ** 2   # 300 MiB

# Axis roles, per tifffile's `series.axes` naming.
_SPATIAL = ("Y", "X")
_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
# Anything else (Z, T, ...) is a stack axis and is indexed by frame.


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.
    """
    ext = os.path.splitext(path)[1].lower()
    if ext in (".tif", ".tiff") and tifffile is not None:
        with tifffile.TiffFile(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
            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
    # Everything else (and if tifffile is unavailable): PIL. Expand palette
    # images to RGB so the LUT is applied instead of returning bare indices.
    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."""
    ext = os.path.splitext(path)[1].lower()
    if ext in (".tif", ".tiff") and tifffile is not None:
        try:
            with tifffile.TiffFile(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"


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):
    """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. The remaining stack axes
    (Z / time / page) are indexed: the first by ``frame``, any deeper ones by 0.
    """
    axes = list(axes)

    # Drop size-1 axes, keeping arr and axes in sync.
    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)

    # Apply the plan to the array: infer the channel axis, then move it last.
    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))

    # Index the leading stack axes: first by `frame`, deeper by 0.
    target = 3 if ci is not None else 2
    first = True
    while len(axes) > target:
        idx = int(frame) if first else 0
        idx = max(0, min(idx, arr.shape[0] - 1))
        arr = arr[idx]
        axes.pop(0)
        first = False

    return arr


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)
    """
    frames: int
    channels: int
    axes: str
    guessed: bool
    shape: tuple
    width: int
    height: int


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
        target = 3 if has_c else 2
        frames = int(planned_shape[0]) if len(planned_axes) > target else 1

        # 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)
    except Exception:  # noqa: BLE001
        return ImageInfo(1, 1, "", False, (), 0, 0)


def count_frames(path):
    """Return how many frames a file's stack (Z/T) axis has (1 if not a stack)."""
    return inspect_image(path).frames


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.
    """
    ext = os.path.splitext(path)[1].lower()
    if ext in (".tif", ".tiff") and tifffile is not None:
        try:
            with tifffile.TiffFile(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.
    """
    ext = os.path.splitext(path)[1].lower()
    if ext in (".tif", ".tiff") and 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):
    """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)
    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, 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 a stack to use (0-based; ignored for non-stacks).
        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)
        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