Refactor render channels to render images (#55)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: sagar87

src/spatialdata_plot/pl/basic.py

@@ -290,7 +290,7 @@ def render_points(
     def render_images(
         self,
         element: str | None = None,
-        channel: str | None = None,
+        channel: list[str] | list[int] | int | str | None = None,
         cmap: Colormap | str | None = None,
         norm: Optional[Normalize] = None,
         na_color: str | tuple[float, ...] | None = (0.0, 0.0, 0.0, 0.0),
@@ -326,6 +326,7 @@ def render_images(
         sdata = self._copy()
         sdata = _verify_plotting_tree(sdata)
         n_steps = len(sdata.plotting_tree.keys())
+
         cmap_params = _prepare_cmap_norm(
             cmap=cmap,
             norm=norm,

src/spatialdata_plot/pl/render.py

@@ -12,6 +12,7 @@
 import pandas as pd
 import scanpy as sc
 import spatialdata as sd
+import xarray as xr
 from anndata import AnnData
 from geopandas import GeoDataFrame
 from matplotlib import colors
@@ -28,6 +29,8 @@
     OutlineParams,
     ScalebarParams,
     _decorate_axs,
+    _get_colors_for_categorical_obs,
+    _get_linear_colormap,
     _map_color_seg,
     _maybe_set_colors,
     _normalize,
@@ -292,7 +295,7 @@ class ImageRenderParams:
 
     cmap_params: CmapParams
     element: str | None = None
-    channel: Sequence[str] | None = None
+    channel: list[str] | list[int] | int | str | None = None
     palette: Palette_t = None
     alpha: float = 1.0
 
@@ -319,10 +322,32 @@ def _render_images(
     if (len(img.c) > 3 or len(img.c) == 2) and render_params.channel is None:
         raise NotImplementedError("Only 1 or 3 channels are supported at the moment.")
 
-    img = _normalize(img, clip=True)
-
     if render_params.channel is not None:
-        img = img.sel(c=[render_params.channel])
+        channels = [render_params.channel] if isinstance(render_params.channel, (str, int)) else render_params.channel
+        img = img.sel(c=channels)
+        num_channels = img.sizes["c"]
+
+        if render_params.palette is not None:
+            if num_channels > len(render_params.palette):
+                raise ValueError("If palette is provided, it must match the number of channels.")
+
+            color = render_params.palette
+
+        else:
+            color = _get_colors_for_categorical_obs(img.coords["c"].values.tolist())
+
+        cmaps = _get_linear_colormap([str(c) for c in color[:num_channels]], "k")
+        img = _normalize(img, clip=True)
+        colored = np.stack([cmaps[i](img.values[i]) for i in range(num_channels)], 0).sum(0)
+        img = xr.DataArray(
+            data=colored,
+            coords=[
+                img.coords["y"],
+                img.coords["x"],
+                ["R", "G", "B", "A"],
+            ],
+            dims=["y", "x", "c"],
+        )
 
     img = img.transpose("y", "x", "c")  # for plotting
 

src/spatialdata_plot/pl/utils.py

@@ -21,7 +21,7 @@
 from matplotlib.axes import Axes
 from matplotlib.cm import get_cmap
 from matplotlib.collections import PatchCollection
-from matplotlib.colors import Colormap, ListedColormap, Normalize, TwoSlopeNorm, to_rgba
+from matplotlib.colors import Colormap, LinearSegmentedColormap, ListedColormap, Normalize, TwoSlopeNorm, to_rgba
 from matplotlib.figure import Figure
 from matplotlib.gridspec import GridSpec
 from matplotlib_scalebar.scalebar import ScaleBar
@@ -918,3 +918,14 @@ def _multiscale_to_image(sdata: sd.SpatialData) -> sd.SpatialData:
             sdata.images[k] = Image2DModel.parse(v["scale0"].ds.to_array().squeeze(axis=0))
 
     return sdata
+
+
+def _get_linear_colormap(colors: list[str], background: str) -> list[LinearSegmentedColormap]:
+    return [LinearSegmentedColormap.from_list(c, [background, c], N=256) for c in colors]
+
+
+def _get_listed_colormap(color_dict: dict[str, str]) -> ListedColormap:
+    sorted_labels = sorted(color_dict.keys())
+    colors = [color_dict[k] for k in sorted_labels]
+
+    return ListedColormap(["black"] + colors, N=len(colors) + 1)

src/spatialdata_plot/pp/utils.py

@@ -1,23 +1,10 @@
 from collections import OrderedDict
 
-import matplotlib
 import spatialdata as sd
-from matplotlib.colors import LinearSegmentedColormap, ListedColormap
 from spatialdata.models import TableModel
 from spatialdata.transformations import get_transformation
 
 
-def _get_linear_colormap(colors: list[str], background: str) -> list[matplotlib.colors.LinearSegmentedColormap]:
-    return [LinearSegmentedColormap.from_list(c, [background, c], N=256) for c in colors]
-
-
-def _get_listed_colormap(color_dict: dict[str, str]) -> matplotlib.colors.ListedColormap:
-    sorted_labels = sorted(color_dict.keys())
-    colors = [color_dict[k] for k in sorted_labels]
-
-    return ListedColormap(["black"] + colors, N=len(colors) + 1)
-
-
 def _get_region_key(sdata: sd.SpatialData) -> str:
     """Quick access to the data's region key."""
     return str(sdata.table.uns[TableModel.ATTRS_KEY][TableModel.REGION_KEY_KEY])