Plot

def plot_drr(
    img: Float[torch.Tensor, "B C H W"],
    mask: Bool[torch.Tensor, "B C H W"] | None = None,
    title: list[str] | None = None,
    ticks: bool = True,
    axs: list[matplotlib.axes.Axes] | None = None,
    cmap: str = "gray",
    mask_cmap: str | matplotlib.colors.Colormap = "Set2",
    mask_n_colors: int = 7,
    interior_alpha: float = 0.3,
    edge_alpha: float = 1.0,
    edge_width: int = 1,
    **imshow_kwargs,
) -> list[matplotlib.axes.Axes]:
    """Plot a batch of DRR images, optionally with a segmentation mask overlay.

    Renders each image by summing across channels, simulating X-ray intensity
    accumulation along a ray. A segmentation mask can be overlaid in two ways:
    passed explicitly via `mask`, or derived automatically when `img` has
    more than one channel (where channel 0 is background and channels 1+ are
    labeled structures). These two modes are mutually exclusive.

    When a mask is rendered, channel 0 is always dropped. It is assumed to
    represent background. Each remaining channel is drawn with a distinct
    color, a translucent interior fill, and an opaque boundary edge detected
    via morphological erosion.

    Args:
        img: Batch of DRR images with shape `(B, C, H, W)`. If `C > 1`,
            channels 1+ are treated as binary segmentation labels and a mask
            is derived as `img > 0`. Channel intensities are summed across
            `C` for display.
        mask: Explicit segmentation mask with shape `(B, C, H, W)`, where
            channel 0 is background and channels 1+ are labeled structures.
            Mutually exclusive with a multi-channel `img`.
        title: Per-image labels of length `B`, rendered as x-axis titles.
            If `None`, no labels are shown.
        ticks: Whether to display 1-indexed pixel coordinate ticks. If
            `False`, all tick marks are hidden. Defaults to `True`.
        axs: Pre-existing axes to plot into. Must have length `B`. If
            `None`, a new figure with `B` subplots is created.
        cmap: Colormap for the DRR image. Defaults to `"gray"`.
        mask_cmap: Colormap used to assign colors to segmentation channels.
            Colors are sampled evenly and cycled if the number of channels
            exceeds `mask_n_colors`. Defaults to `"Set2"`.
        mask_n_colors: Number of evenly spaced colors to sample from
            `mask_cmap` before cycling. Defaults to `7`.
        interior_alpha: Opacity of the filled mask interior, in `[0, 1]`.
            Defaults to `0.3`.
        edge_alpha: Opacity of the mask boundary, in `[0, 1]`.
            Defaults to `1.0`.
        edge_width: Boundary thickness in pixels. Controls the erosion kernel
            size as `2 * edge_width + 1`. Defaults to `1`.
        **imshow_kwargs: Additional keyword arguments forwarded to
            `ax.imshow` for the DRR image only, not the mask.

    Returns:
        List of `Axes` of length `B`, one per image in the batch.

    Raises:
        ValueError: If `img` has more than one channel and `mask` is
            also provided.
    """
    if img.shape[1] > 1 and mask is not None:
        raise ValueError("Pass either a multi-channel img or an explicit mask, not both.")

    axs = _plot_img(img.sum(dim=1, keepdim=True), title, ticks, axs, cmap, **imshow_kwargs)

    if img.shape[1] > 1:
        mask = img > 0
    elif mask is None:
        return axs

    _plot_mask(
        mask[:, 1:].float(),
        axs=axs,
        mask_cmap=mask_cmap,
        mask_n_colors=mask_n_colors,
        interior_alpha=interior_alpha,
        edge_alpha=edge_alpha,
        edge_width=edge_width,
    )
    return axs

def overlay(
    moving: Float[torch.Tensor, "B C H W"],
    fixed: Float[torch.Tensor, "B C H W"],
    title: list[str] | None = None,
    ticks: bool = True,
    axs: list[matplotlib.axes.Axes] | None = None,
    blur_kernel: int = 3,
    canny_low: int = 0,
    canny_high: int = 100,
    edge_color: tuple[float, float, float] = (1.0, 0.0, 0.0),
    edge_alpha: float = 1.0,
) -> list[matplotlib.axes.Axes]:
    """Plot a batch of fixed images with moving image edges overlaid.

    Edges are extracted from the moving image via Canny detection and drawn
    as a transparent layer on top of the fixed image. Useful for visually
    assessing registration quality.

    Args:
        moving: Moving images with shape `(B, C, H, W)`. Channels are summed before edge detection.
        fixed: Fixed images with shape `(B, C, H, W)`. Channels are summed for display.
        title: Per-image labels of length `B`. If `None`, no labels are shown.
        ticks: Whether to display 1-indexed pixel coordinate ticks. Defaults to `True`.
        axs: Pre-existing axes of length `B`. If `None`, a new figure is created.
        blur_kernel: Gaussian blur kernel size applied before Canny. Defaults to `3`.
        canny_low: Lower hysteresis threshold for Canny edge detection. Defaults to `0`.
        canny_high: Upper hysteresis threshold for Canny edge detection. Defaults to `100`.
        edge_color: RGB color of the overlaid edges. Defaults to red `(1.0, 0.0, 0.0)`.
        edge_alpha: Opacity of the overlaid edges, in `[0, 1]`. Defaults to `1.0`.

    Returns:
        List of `Axes` of length `B`, one per image in the batch.
    """
    fixed_gray = fixed.sum(dim=1, keepdim=True)
    moving_gray = moving.sum(dim=1)

    def to_uint8(t: torch.Tensor) -> np.ndarray:
        arr = t.cpu().detach().numpy()
        return cv2.normalize(arr, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)

    axs = _plot_img(fixed_gray, title, ticks, axs, cmap="gray")

    H, W = fixed_gray.shape[-2:]
    rgba = np.zeros((H, W, 4), dtype=np.float32)

    for m, ax in zip(moving_gray, axs):
        m_blurred = cv2.GaussianBlur(to_uint8(m), (blur_kernel, blur_kernel), 0)
        edge_mask = cv2.Canny(m_blurred, canny_low, canny_high).astype(bool)
        rgba[edge_mask] = (*edge_color, edge_alpha)
        ax.imshow(rgba)
        rgba[edge_mask] = 0

    return axs

def animate(
    moving_img: Float[torch.Tensor, "B C H W"],
    moving_mask: Bool[torch.Tensor, "B C H W"] | None = None,
    out: str | pathlib.Path | None = None,
    fixed_img: Float[torch.Tensor, "1 C H W"] | None = None,
    fixed_mask: Bool[torch.Tensor, "1 C H W"] | None = None,
    titles: list[str] | None = None,
    fps: int = 20,
    pause: float = 1.0,
    verbose: bool = True,
    blur_kernel: int = 3,
    canny_low: int = 0,
    canny_high: int = 100,
    edge_color: tuple[float, float, float] = (1.0, 0.0, 0.0),
    edge_alpha: float = 1.0,
    **kwargs,
) -> pathlib.Path | None:
    """Create an animated GIF from a batch of DRR images.

    Renders a sequence of DRR images as an animated GIF, with optional
    side-by-side comparison against a fixed reference image. When `out` is
    `None`, displays the animation inline in Jupyter notebooks.

    Multi-channel images are automatically converted to single-channel with
    segmentation masks extracted from channels 1+ (channel 0 is background).

    When `fixed_img` is provided, a third column is rendered showing the
    moving image edges overlaid on the fixed image via `overlay`.

    Args:
        moving_img: Batch of moving DRR images.
        moving_mask: Optional segmentation mask for moving images.
        out: Output file path, or `None` for inline display.
        fixed_img: Optional fixed reference image for comparison.
        fixed_mask: Optional segmentation mask for fixed image.
        titles: Optional per-frame titles of length `B`.
        fps: Frames per second for playback.
        pause: Pause duration in seconds at the end of the loop.
        verbose: Whether to display rendering progress.
        blur_kernel: Gaussian blur kernel size applied before Canny edge detection.
        canny_low: Lower hysteresis threshold for Canny edge detection.
        canny_high: Upper hysteresis threshold for Canny edge detection.
        edge_color: RGB color of the overlaid edges.
        edge_alpha: Opacity of the overlaid edges, in `[0, 1]`.
        **kwargs: Additional arguments forwarded to `imageio.v3.imwrite` or `plot_drr`.

    Returns:
        Path to saved file if `out` is provided, otherwise `None`.

    Raises:
        ValueError: If `titles` length does not match batch size.
    """
    B = len(moving_img)
    if titles is not None and len(titles) != B:
        raise ValueError(f"titles length ({len(titles)}) must match batch size ({B})")

    moving_img, moving_mask = _normalize(moving_img, moving_mask)
    if fixed_img is not None:
        fixed_img, fixed_mask = _normalize(fixed_img, fixed_mask)

    iio_keys = {"duration", "loop", "quality", "quantizer", "palettesize"}
    iio_kwargs = {k: v for k, v in kwargs.items() if k in iio_keys}
    iio_kwargs.setdefault("fps", fps)
    iio_kwargs.setdefault("loop", 0)
    plot_kwargs = {k: v for k, v in kwargs.items() if k not in iio_keys}

    has_fixed = fixed_img is not None
    n_cols = 3 if has_fixed else 1
    figsize = (3 * n_cols, 3)

    overlay_kwargs = dict(
        blur_kernel=blur_kernel,
        canny_low=canny_low,
        canny_high=canny_high,
        edge_color=edge_color,
        edge_alpha=edge_alpha,
    )

    iterator = tqdm(range(B), desc="Rendering frames", ncols=75) if verbose else range(B)
    frames = []

    for i in iterator:
        fig, axs = plt.subplots(ncols=n_cols, figsize=figsize, constrained_layout=True)
        axs = [axs] if n_cols == 1 else list(axs)

        if has_fixed:
            frame_img = torch.cat([fixed_img, moving_img[i : i + 1]])
            frame_mask = _concat_masks(fixed_mask, moving_mask[i : i + 1] if moving_mask is not None else None)
            frame_titles = ["Fixed", titles[i] if titles else "Moving", "Overlay"]
            plot_drr(frame_img, frame_mask, title=frame_titles[:2], axs=axs[:2], **plot_kwargs)
            overlay(moving_img[i : i + 1], fixed_img, title=[frame_titles[2]], axs=axs[2], **overlay_kwargs)
        else:
            frame_img = moving_img[i : i + 1]
            frame_mask = moving_mask[i : i + 1] if moving_mask is not None else None
            frame_titles = [titles[i]] if titles else None
            plot_drr(frame_img, frame_mask, title=frame_titles, axs=axs, **plot_kwargs)

        fig.canvas.draw()
        frames.append(np.asarray(fig.canvas.buffer_rgba())[..., :3])
        plt.close(fig)

    if pause > 0:
        frames.extend([frames[-1]] * int(pause * fps))

    frames_array = np.stack(frames)
    if out is None:
        gif_bytes = iio.imwrite("<bytes>", frames_array, extension=".gif", **iio_kwargs)
        ipython_display(HTML(f"<img src='data:image/gif;base64,{b64encode(gif_bytes).decode()}'>"))
        return None
    else:
        out_path = pathlib.Path(out)
        iio.imwrite(out_path, frames_array, **iio_kwargs)
        return out_path