Render 2D crops

biahub/analysis/AnalysisSettings.py

@@ -1,17 +1,21 @@
 import warnings
 
-from typing import Literal, Optional, Union
+from pathlib import Path
+from typing import Annotated, Literal, NamedTuple, Optional, Union
 
 import numpy as np
 import torch
 
+from cmap import Colormap
 from pydantic import (
     BaseModel,
     ConfigDict,
+    Field,
     NonNegativeInt,
     PositiveFloat,
     PositiveInt,
     field_validator,
+    model_validator,
 )
 
 
@@ -206,3 +210,65 @@ def __init__(self, **data):
                     DeprecationWarning,
                 )
         super().__init__(**data)
+
+
+class IndexRange(NamedTuple):
+    """Index range for a single axis."""
+
+    start: NonNegativeInt
+    stop: NonNegativeInt
+
+
+class BaseChannelRender2DSettings(MyBaseModel):
+    """Settings for rendering a single channel in 2D.
+    Base model for a discriminated union."""
+
+    path: Path
+    name: str
+    multiscale_level: str = "0"
+
+
+PositiveFloatLE1 = Annotated[float, Field(gt=0.0, le=1.0)]
+
+
+class ImageChannelRender2DSettings(BaseChannelRender2DSettings):
+    """Settings for rendering an image as a bitmap in 2D."""
+
+    lut: Colormap
+    channel_type: Literal["image"]
+    clim: tuple[float, float] | None
+    clim_mode: Literal["absolute", "percentile"] | None
+    alpha: PositiveFloatLE1 = 1.0
+    gamma: PositiveFloatLE1 = 1.0
+
+    @model_validator(mode="after")
+    @classmethod
+    def check_clim_mode(cls, data):
+        if data.clim_mode is not None and data.clim is None:
+            raise ValueError("clim_mode requires clim to be set")
+        return data
+
+
+class ContourChannelRender2DSettings(BaseChannelRender2DSettings):
+    """Settings for rendering contours as vectors in 2D."""
+
+    channel_type: Literal["contour"]
+    linewidth: PositiveFloat = 2.0
+
+
+ChannelRender2DSettings = Annotated[
+    ImageChannelRender2DSettings | ContourChannelRender2DSettings,
+    Field(discriminator="channel_type"),
+]
+
+
+class Render2DSettings(MyBaseModel):
+    """Settings for rendering 2D images."""
+
+    time_index: NonNegativeInt
+    channels: list[ChannelRender2DSettings]
+    z_range: IndexRange
+    y_range: IndexRange
+    x_range: IndexRange
+    figsize: tuple[PositiveFloat, PositiveFloat] = (4, 4)
+    dpi: PositiveInt = 300

biahub/analysis/settings/example_render_2d_settings.yml

@@ -0,0 +1,27 @@
+# Render 2D settings
+time_index: 0
+channels:
+  - path: /path/to/plate.zarr/row/col/fov
+    name: GFP
+    # optional
+    multiscale_level: 0
+    channel_type: image
+    lut: green
+    clim: [100, 500]
+    clim_mode: absolute
+    alpha: 1.0
+    gamma: 1.0
+  - path: /path/to/plate.zarr/row/col/fov
+    name: nuclei_masks
+    channel_type: contour
+    linewidth: 2.0
+# Z index range
+z_range: [0, 1]
+# Y index range
+y_range: [0, 512]
+# X index range
+x_range: [0, 512]
+# export figsize in inches
+figsize: [4, 4]
+# export DPI
+dpi: 300

biahub/cli/main.py

@@ -12,6 +12,7 @@
 from biahub.cli.estimate_stitch import estimate_stitch
 from biahub.cli.optimize_registration import optimize_registration
 from biahub.cli.register import register
+from biahub.cli.render import render_2d
 from biahub.cli.stabilize import stabilize
 from biahub.cli.stitch import stitch
 from biahub.cli.update_scale_metadata import update_scale_metadata
@@ -45,3 +46,4 @@ def cli():
 cli.add_command(estimate_psf)
 cli.add_command(deconvolve)
 cli.add_command(characterize_psf)
+cli.add_command(render_2d)

biahub/cli/render.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import click
+import matplotlib.pyplot as plt
+import numpy as np
+
+from iohub.ngff import open_ome_zarr
+from skimage.exposure import rescale_intensity
+from skimage.measure import find_contours
+
+from biahub.analysis.AnalysisSettings import Render2DSettings
+from biahub.cli.parsing import config_filepath, output_filepath
+from biahub.cli.utils import yaml_to_model
+
+if TYPE_CHECKING:
+    from iohub.ngff import ImageArray
+    from numpy.typing import NDArray
+
+    from biahub.analysis.AnalysisSettings import (
+        ContourChannelRender2DSettings,
+        ImageChannelRender2DSettings,
+    )
+
+
+def render_image(image: NDArray, settings: ImageChannelRender2DSettings) -> NDArray:
+    """Render an image in RGBA and scale to RGB by A."""
+    if image.ndim != 2:
+        raise ValueError(f"Image must be 2D, got shape: {image.shape}")
+    image = image.astype(np.float32)
+    if settings.clim is not None:
+        if settings.clim_mode == "absolute":
+            low, high = settings.clim
+        elif settings.clim_mode == "percentile":
+            low, high = np.percentile(image, settings.clim)
+    else:
+        low, high = np.min(image), np.max(image)
+    if low >= high:
+        raise ValueError("Could not determine valid contrast limits.")
+    image = rescale_intensity(image, in_range=(low, high), out_range=(0, 1))
+    image = settings.lut(image, gamma=settings.gamma)
+    image[..., 3] *= settings.alpha
+    return image[..., :3] * image[..., 3:]
+
+
+@click.command()
+@output_filepath()
+@config_filepath()
+def render_2d(config_filepath: Path, output_filepath: Path) -> None:
+    """Render a 2D slice in the given LUT."""
+    settings = yaml_to_model(config_filepath, Render2DSettings)
+
+    def _slice_tczyx(array: ImageArray, channel_index: int) -> NDArray:
+        crop = array[
+            settings.time_index,
+            channel_index,
+            slice(*settings.z_range),
+            slice(*settings.y_range),
+            slice(*settings.x_range),
+        ]
+        return np.squeeze(crop)
+
+    rendered_images: list[NDArray] = []
+    labels_and_settings: list[tuple[NDArray, ContourChannelRender2DSettings]] = []
+    for ch_setting in settings.channels:
+        with open_ome_zarr(ch_setting.path, layout="fov") as dataset:
+            ch_idx = dataset.get_channel_index(ch_setting.name)
+            crop = _slice_tczyx(dataset[ch_setting.multiscale_level], ch_idx)
+        crop_and_settings = (crop, ch_setting)
+        if ch_setting.channel_type == "image":
+            rendered_images.append(render_image(crop, ch_setting))
+        elif ch_setting.channel_type == "contour":
+            labels_and_settings.append(crop_and_settings)
+
+    # render images
+    image = np.sum(rendered_images, axis=0).clip(0, 1)
+
+    # render contours
+    figure, ax = plt.subplots(figsize=settings.figsize, dpi=300)
+    ax.imshow(image)
+    for idx, (labels, ct_settings) in enumerate(labels_and_settings):
+        unique_labels = np.unique(labels)
+        for label in unique_labels:
+            if label == 0:
+                continue
+            mask = labels == label
+            contours = find_contours(mask, level=0.5)
+            for contour in contours:
+                ax.plot(
+                    contour[:, 1],
+                    contour[:, 0],
+                    linewidth=ct_settings.linewidth,
+                    color=ct_settings.lut(idx),
+                )
+    # Use image coordinates, not data coordinates
+    ax.axis("image")
+    ax.set_xticks([])
+    ax.set_yticks([])
+    figure.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
+    figure.savefig(output_filepath, bbox_inches="tight", pad_inches=0)

pyproject.toml

@@ -41,6 +41,8 @@ dependencies = [
   "antspyx",
   "pystackreg",
   "napari-psf-analysis",
+  "cmap",
+  "scikit-image",
 ]