Add transfunc and grayscale to render_images (#567)

Author: timtreis

src/spatialdata_plot/pl/basic.py

@@ -3,6 +3,7 @@
 import contextlib
 import sys
 from collections import OrderedDict
+from collections.abc import Callable
 from copy import deepcopy
 from pathlib import Path
 from typing import Any, Literal, cast
@@ -517,6 +518,8 @@ def render_images(
         palette: list[str] | str | None = None,
         alpha: float | int = 1.0,
         scale: str | None = None,
+        grayscale: bool = False,
+        transfunc: Callable[[np.ndarray], np.ndarray] | list[Callable[[np.ndarray], np.ndarray]] | None = None,
         colorbar: bool | str | None = "auto",
         colorbar_params: dict[str, object] | None = None,
         **kwargs: Any,
@@ -561,6 +564,34 @@ def render_images(
                 3) "full": Renders the full image without rasterization. In the case of
                 multiscale images, the highest resolution scale is selected. Note that
                 this may result in long computing times for large images.
+        grayscale : bool, default False
+            Convert the image to grayscale before rendering using luminance
+            weights (Rec. 601: 0.2989 R + 0.5870 G + 0.1140 B). Requires
+            exactly 3 channels at the point of conversion — if ``transfunc``
+            is also provided, it runs first, and the result must have 3
+            channels. The grayscale image is rendered as a single-channel
+            image with ``cmap="gray"`` unless an explicit ``cmap`` is given.
+            Useful for de-emphasising H&E tissue when overlaying colored
+            annotations. Cannot be combined with ``palette``.
+        transfunc : callable or list of callables, optional
+            Transform(s) applied to the raw image array before normalization
+            and rendering.
+
+            **Single callable**: receives a numpy array of shape ``(c, y, x)``
+            (channels first) and must return an array of the same layout.
+            The number of channels may change (e.g., stain deconvolution).
+            Elementwise functions like ``np.log1p`` broadcast naturally.
+            Note that reductions like ``np.percentile`` will compute a
+            *single* value across all channels.
+
+            **List of callables**: one per channel (length must match the
+            number of selected channels). Each receives a ``(y, x)`` array
+            for its channel and must return a ``(y, x)`` array. Use this
+            when each channel needs independent treatment (e.g., different
+            gamma corrections for different fluorescence markers).
+
+            When combined with ``grayscale=True``, ``transfunc`` runs first
+            and ``grayscale`` is applied to the result.
         colorbar :
             Whether to request a colorbar for continuous colors. Use "auto" (default) for automatic selection.
         colorbar_params :
@@ -575,6 +606,8 @@ def render_images(
             The SpatialData object with the rendered images.
         """
         # TODO add Normalize object in tutorial notebook and point to that notebook here
+        if grayscale and palette is not None:
+            raise ValueError("Cannot combine grayscale=True with palette.")
         if "vmin" in kwargs or "vmax" in kwargs:
             logger.warning("`vmin` and `vmax` are deprecated. Pass matplotlib `Normalize` object to norm instead.")
         params_dict = _validate_image_render_params(
@@ -624,6 +657,8 @@ def render_images(
                 zorder=n_steps,
                 colorbar=param_values["colorbar"],
                 colorbar_params=param_values["colorbar_params"],
+                transfunc=transfunc,
+                grayscale=grayscale,
             )
             n_steps += 1
 

src/spatialdata_plot/pl/render.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import dataclasses
 from collections import abc
 from copy import copy
 from typing import Any
@@ -15,6 +16,7 @@
 import pandas as pd
 import scanpy as sc
 import spatialdata as sd
+import xarray as xr
 from anndata import AnnData
 from matplotlib.cm import ScalarMappable
 from matplotlib.colors import ListedColormap, Normalize
@@ -64,6 +66,7 @@
     _maybe_set_colors,
     _mpl_ax_contains_elements,
     _multiscale_to_spatial_image,
+    _prepare_cmap_norm,
     _prepare_transformation,
     _rasterize_if_necessary,
     _set_color_source_vec,
@@ -1019,6 +1022,14 @@ def _render_points(
     )
 
 
+_LUMINANCE_WEIGHTS = np.array([0.2989, 0.5870, 0.1140])
+
+
+def _grayscale_transform(img_cyx: np.ndarray) -> np.ndarray:
+    """Convert a (3, y, x) RGB image to (1, y, x) luminance."""
+    return np.tensordot(_LUMINANCE_WEIGHTS, img_cyx, axes=([0], [0]))[np.newaxis]
+
+
 def _normalize_dtype_to_float(arr: np.ndarray) -> np.ndarray:
     """Normalize an array to float64 in [0, 1] for display with matplotlib.
 
@@ -1122,8 +1133,65 @@ def _render_images(
 
     # the channel parameter has been previously validated, so when not None, render_params.channel is a list
     assert isinstance(channels, list)
+
+    _, trans_data = _prepare_transformation(img, coordinate_system, ax)
+
+    # --- Apply image transforms ---
+    transfunc = render_params.transfunc
+    needs_transform = transfunc is not None or render_params.grayscale
+
+    if needs_transform:
+        raw = np.stack([img.sel(c=ch).values for ch in channels], axis=0)
+
+        # 1) Apply transfunc (before grayscale)
+        if isinstance(transfunc, list):
+            if len(transfunc) != raw.shape[0]:
+                raise ValueError(
+                    f"Length of transfunc list ({len(transfunc)}) must match the number of channels ({raw.shape[0]})."
+                )
+            raw = np.stack([fn(raw[i]) for i, fn in enumerate(transfunc)], axis=0)
+        elif transfunc is not None:
+            raw = transfunc(raw)
+
+        # 2) Apply grayscale (after transfunc)
+        if render_params.grayscale:
+            if raw.shape[0] != 3:
+                raise ValueError(
+                    f"grayscale=True requires exactly 3 channels"
+                    f"{' after transfunc' if transfunc is not None else ''}, "
+                    f"got {raw.shape[0]}. Select 3 channels via the 'channel' parameter."
+                )
+            raw = _grayscale_transform(raw)
+
+        # Rebuild image with new channel coords
+        new_channels = list(range(raw.shape[0]))
+        img = xr.DataArray(
+            data=raw,
+            dims=("c", "y", "x"),
+            coords={"c": new_channels, "y": img.coords["y"], "x": img.coords["x"]},
+        )
+        channels = new_channels
+
     n_channels = len(channels)
 
+    # When grayscale was applied and user didn't provide an explicit cmap,
+    # default to "gray" for intuitive single-channel rendering.
+    got_multiple_cmaps = isinstance(render_params.cmap_params, list)
+    if (
+        render_params.grayscale
+        and not got_multiple_cmaps
+        and isinstance(render_params.cmap_params, CmapParams)
+        and render_params.cmap_params.cmap_is_default
+    ):
+        render_params = dataclasses.replace(
+            render_params,
+            cmap_params=_prepare_cmap_norm(
+                cmap="gray",
+                norm=render_params.cmap_params.norm,
+                na_color=render_params.cmap_params.na_color,
+            ),
+        )
+
     # True if user gave n cmaps for n channels
     got_multiple_cmaps = isinstance(render_params.cmap_params, list)
     if got_multiple_cmaps:
@@ -1139,8 +1207,6 @@ def _render_images(
     if isinstance(render_params.cmap_params, list) and len(render_params.cmap_params) != n_channels:
         raise ValueError("If 'cmap' is provided, its length must match the number of channels.")
 
-    _, trans_data = _prepare_transformation(img, coordinate_system, ax)
-
     # Detect RGB(A) images by channel names — skip when user overrides with palette/cmap
     is_rgb, has_alpha = _is_rgb_image(channels)
     has_explicit_cmap = (

src/spatialdata_plot/pl/render_params.py

@@ -269,6 +269,8 @@ class ImageRenderParams:
     zorder: int = 0
     colorbar: bool | str | None = "auto"
     colorbar_params: dict[str, object] | None = None
+    transfunc: Callable[[np.ndarray], np.ndarray] | list[Callable[[np.ndarray], np.ndarray]] | None = None
+    grayscale: bool = False
 
 
 @dataclass