_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 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
# 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("." + e for e in tifffile.TIFF.FILE_EXTENSIONS if "." not in e)
    if tifffile is not None else [".tif", ".tiff"]
)


def _read_tiff(path):
    """Read a TIFF-family file into (array, axes). Raises if not readable."""
    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
                      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 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, 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)


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.
    """
    if 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.
    """
    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