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,
    edge_detection_size: int = 200,
) -> list[matplotlib.axes.Axes]:
    """Overlay moving image edges on fixed images for registration assessment.

    Edges are detected using Canny at a fixed resolution for threshold consistency,
    then upscaled with bilinear interpolation for anti-aliased rendering.

    Args:
        moving: Moving images, shape (B, C, H, W)
        fixed: Fixed images, shape (B, C, H, W)
        title: Optional titles for each image in batch
        ticks: Whether to show pixel coordinate ticks
        axs: Optional pre-existing axes to plot on
        blur_kernel: Gaussian blur kernel size (must be odd)
        canny_low: Canny lower threshold
        canny_high: Canny upper threshold
        edge_color: RGB color tuple for edges, values in [0, 1]
        edge_alpha: Edge opacity in [0, 1]
        edge_detection_size: Resolution for Canny detection

    Returns:
        List of matplotlib Axes objects

    Raises:
        ValueError: If input shapes don't match or parameters are invalid
    """
    if moving.ndim != 4:
        raise ValueError(f"Expected 4D tensors (B, C, H, W), got {moving.ndim}D")
    if title is not None and len(title) != moving.shape[0]:
        raise ValueError(f"Title length {len(title)} != batch size {moving.shape[0]}")
    if blur_kernel % 2 == 0 or blur_kernel < 1:
        raise ValueError(f"blur_kernel must be positive odd integer, got {blur_kernel}")
    if not 0 <= edge_alpha <= 1:
        raise ValueError(f"edge_alpha must be in [0, 1], got {edge_alpha}")

    fixed_gray = fixed.sum(dim=1, keepdim=True)
    moving_gray = moving.sum(dim=1)

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

    H, W = fixed_gray.shape[-2:]
    for moving_img, ax in zip(moving_gray, axs):
        img_np = moving_img.cpu().detach().numpy()
        img_uint8 = cv2.normalize(
            img_np, dst=np.empty_like(img_np), alpha=0, beta=255, norm_type=cv2.NORM_MINMAX
        ).astype(np.uint8)

        img_small = cv2.resize(img_uint8, (edge_detection_size, edge_detection_size), interpolation=cv2.INTER_AREA)
        img_blurred = cv2.GaussianBlur(img_small, (blur_kernel, blur_kernel), 0)
        edges = cv2.Canny(img_blurred, canny_low, canny_high)

        edge_weights = cv2.resize(edges.astype(np.float32) / 255.0, (W, H), interpolation=cv2.INTER_LINEAR)

        rgba = np.zeros((H, W, 4), dtype=np.float32)
        rgba[..., :3] = edge_color
        rgba[..., 3] = edge_weights * edge_alpha

        ax.imshow(rgba)

    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 | 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,
    ticks: bool = True,
    fps: int = 20,
    pause: float = 1.0,
    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,
    edge_detection_size: int = 200,
    verbose: bool = True,
    **kwargs,
) -> 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`.
        ticks: Whether to show pixel coordinate ticks.
        fps: Frames per second for playback.
        pause: Pause duration in seconds at the end of the loop.
        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]`.
        verbose: Whether to display rendering progress.
        **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}

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

    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 fixed_img is not None:
            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], ticks=ticks, **plot_kwargs)
            overlay(
                moving_img[i : i + 1],
                fixed_img,
                title=[frame_titles[2]],
                ticks=ticks,
                axs=axs[2],
                blur_kernel=blur_kernel,
                canny_low=canny_low,
                canny_high=canny_high,
                edge_color=edge_color,
                edge_alpha=edge_alpha,
                edge_detection_size=edge_detection_size,
            )
        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, ticks=ticks, axs=axs, **plot_kwargs)

        fig.canvas.draw()
        frames.append(np.asarray(cast(FigureCanvasAgg, 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 = imwrite("<bytes>", frames_array, extension=".gif", **iio_kwargs)
        display(HTML(f"<img src='data:image/gif;base64,{b64encode(gif_bytes).decode()}'>"))
        return None
    else:
        out_path = Path(out)
        imwrite(out_path, frames_array, **iio_kwargs)
        return out_path