utils

sleap_nn.training.utils

Miscellaneous utility functions for training.

Classes:

Name	Description
MatplotlibRenderer	Renders VisualizationData to matplotlib figures.
VisualizationData	Container for visualization data from a single sample.
WandBRenderer	Renders VisualizationData to wandb.Image objects.

Functions:

Name	Description
get_dist_rank	Return the rank of the current process if torch.distributed is initialized.
get_gpu_memory	Get the available memory on each GPU.
imgfig	Create a tight figure for image plotting.
is_distributed_initialized	Check if distributed processes are initialized.
plot_confmaps	Plot confidence maps reduced over channels.
plot_img	Plot an image in a tight figure.
plot_peaks	Plot ground truth and detected peaks.
plot_tile_grid	Visualize the sliding-window tile grid (and blend coverage) over a frame.
xavier_init_weights	Function to initilaise the model weights with Xavier initialization method.

MatplotlibRenderer

Renders VisualizationData to matplotlib figures.

Methods:

Name	Description
render	Render visualization data to a matplotlib figure.
render_gt_mask	Render the predicted foreground vs the ground-truth mask overlay.
render_instance_masks	Render grouped predicted per-instance masks as a colored overlay.
render_offsets	Render the center-offset field DIRECTION for segmentation models.
render_pafs	Render PAF magnitude visualization.

Source code in sleap_nn/training/utils.py

class MatplotlibRenderer:
    """Renders VisualizationData to matplotlib figures."""

    def render(self, data: VisualizationData) -> matplotlib.figure.Figure:
        """Render visualization data to a matplotlib figure.

        Args:
            data: VisualizationData containing image, confmaps, peaks, etc.

        Returns:
            A matplotlib Figure object.
        """
        img = data.image
        scale = 1.0
        if img.shape[0] < 512:
            scale = 2.0
        if img.shape[0] < 256:
            scale = 4.0

        fig = plot_img(img, dpi=72 * scale, scale=scale)
        plot_confmaps(data.pred_confmaps, output_scale=data.output_scale)
        plot_peaks(data.gt_instances, data.pred_peaks, paired=data.is_paired)
        return fig

    def render_gt_mask(self, data: VisualizationData) -> matplotlib.figure.Figure:
        """Render the predicted foreground vs the ground-truth mask overlay.

        Draws the image, overlays the predicted foreground probability (as for the
        standard confmap render), then outlines the ground-truth foreground mask so
        the prediction can be compared against truth each epoch (the segmentation
        analog of pose's GT-keypoint overlay).

        Args:
            data: VisualizationData with ``gt_mask`` populated.

        Returns:
            A matplotlib Figure object.
        """
        if data.gt_mask is None:
            raise ValueError("gt_mask is None, cannot render GT mask overlay")

        img = data.image
        scale = 1.0
        if img.shape[0] < 512:
            scale = 2.0
        if img.shape[0] < 256:
            scale = 4.0

        gt = np.asarray(data.gt_mask)
        if gt.ndim == 3:
            gt = gt[..., 0]
        gt = gt.astype(np.float32)

        fig = plot_img(img, dpi=72 * scale, scale=scale)
        # Predicted foreground probability as a translucent heatmap (same overlay
        # the confmap render uses), so pred and GT are visible together.
        plot_confmaps(data.pred_confmaps, output_scale=data.output_scale)
        ax = plt.gca()
        gt_output_scale = gt.shape[0] / img.shape[0]
        extent = [
            -0.5,
            gt.shape[1] / gt_output_scale - 0.5,
            gt.shape[0] / gt_output_scale - 0.5,
            -0.5,
        ]
        # GT foreground outline (single contour at the 0.5 level).
        ax.contour(
            np.linspace(extent[0], extent[1], gt.shape[1]),
            np.linspace(extent[3], extent[2], gt.shape[0]),
            gt,
            levels=[0.5],
            colors="lime",
            linewidths=1.5,
        )
        return fig

    def render_pafs(self, data: VisualizationData) -> matplotlib.figure.Figure:
        """Render PAF magnitude visualization.

        Args:
            data: VisualizationData with pred_pafs populated.

        Returns:
            A matplotlib Figure object showing PAF magnitudes.
        """
        if data.pred_pafs is None:
            raise ValueError("pred_pafs is None, cannot render PAFs")

        img = data.image
        scale = 1.0
        if img.shape[0] < 512:
            scale = 2.0
        if img.shape[0] < 256:
            scale = 4.0

        # Compute PAF magnitude
        pafs = data.pred_pafs  # (H, W, 2*edges) or (H, W, edges, 2)
        if pafs.ndim == 3:
            n_edges = pafs.shape[-1] // 2
            pafs = pafs.reshape(pafs.shape[0], pafs.shape[1], n_edges, 2)
        magnitude = np.sqrt(pafs[..., 0] ** 2 + pafs[..., 1] ** 2)
        magnitude = magnitude.max(axis=-1)  # Max over edges

        fig = plot_img(img, dpi=72 * scale, scale=scale)
        ax = plt.gca()

        # Calculate PAF output scale from actual PAF dimensions, not confmap output_scale
        # PAFs may have a different output_stride than confmaps
        paf_output_scale = magnitude.shape[0] / img.shape[0]

        ax.imshow(
            magnitude,
            alpha=0.5,
            origin="upper",
            cmap="viridis",
            extent=[
                -0.5,
                magnitude.shape[1] / paf_output_scale - 0.5,
                magnitude.shape[0] / paf_output_scale - 0.5,
                -0.5,
            ],
        )
        return fig

    def render_offsets(self, data: VisualizationData) -> matplotlib.figure.Figure:
        """Render the center-offset field DIRECTION for segmentation models.

        Colors each pixel by the ANGLE of its center-offset vector (hue =
        ``atan2(dy, dx)``, the standard optical-flow color wheel), so the field
        reads as "which way does each foreground pixel point to its instance
        center." Opacity is modulated by the offset magnitude so near-zero
        (background) offsets stay transparent and only meaningful vectors show
        their direction.

        Args:
            data: VisualizationData with ``pred_offsets`` (H, W, 2) populated,
                laid out as ``[..., 0] = dx`` and ``[..., 1] = dy``.

        Returns:
            A matplotlib Figure object showing the per-pixel offset direction.
        """
        if data.pred_offsets is None:
            raise ValueError("pred_offsets is None, cannot render offsets")

        import matplotlib.colors as mcolors

        img = data.image
        scale = 1.0
        if img.shape[0] < 512:
            scale = 2.0
        if img.shape[0] < 256:
            scale = 4.0

        offsets = data.pred_offsets  # (H, W, 2) -> (dx, dy)
        dx, dy = offsets[..., 0], offsets[..., 1]
        angle = np.arctan2(dy, dx)  # [-pi, pi]
        magnitude = np.hypot(dx, dy)

        # Hue = direction (wrapped to [0, 1]); full saturation/value for vivid color.
        hue = (angle + np.pi) / (2.0 * np.pi)
        hsv = np.stack([hue, np.ones_like(hue), np.ones_like(hue)], axis=-1)
        rgb = mcolors.hsv_to_rgb(hsv)  # (H, W, 3)
        # Opacity ramps with magnitude so background (~0 offset) is transparent.
        mmax = float(magnitude.max())
        alpha = (magnitude / mmax if mmax > 0 else magnitude) * 0.7
        rgba = np.dstack([rgb, alpha])  # (H, W, 4)

        fig = plot_img(img, dpi=72 * scale, scale=scale)
        ax = plt.gca()
        offset_output_scale = rgba.shape[0] / img.shape[0]
        ax.imshow(
            rgba,
            origin="upper",
            extent=[
                -0.5,
                rgba.shape[1] / offset_output_scale - 0.5,
                rgba.shape[0] / offset_output_scale - 0.5,
                -0.5,
            ],
        )
        return fig

    #: Distinct per-instance overlay colors (RGB, 0-1), cycled by instance index.
    _INSTANCE_COLORS = [
        (1.0, 0.31, 0.31),
        (0.31, 1.0, 0.31),
        (0.31, 0.31, 1.0),
        (1.0, 1.0, 0.31),
        (1.0, 0.31, 1.0),
        (0.31, 1.0, 1.0),
        (1.0, 0.63, 0.31),
        (0.63, 0.31, 1.0),
    ]

    def render_instance_masks(
        self, data: VisualizationData
    ) -> matplotlib.figure.Figure:
        """Render grouped predicted per-instance masks as a colored overlay.

        Each predicted instance's mask (from the offset-grouping that already runs
        in ``get_visualization_data``) is drawn in a distinct color on the input
        image, so instance separation — the model's actual output, and the single
        most informative seg diagnostic — is visible each epoch. When a GT
        foreground mask is present (``gt_mask``), its silhouette is outlined for
        reference. Empty instance lists render the plain image (no instances
        detected this epoch).

        Args:
            data: VisualizationData with ``instance_masks`` populated (a list of
                ``(H, W)`` boolean masks at output-stride resolution).

        Returns:
            A matplotlib Figure object.
        """
        if data.instance_masks is None:
            raise ValueError(
                "instance_masks is None, cannot render instance-mask overlay"
            )

        img = data.image
        scale = 1.0
        if img.shape[0] < 512:
            scale = 2.0
        if img.shape[0] < 256:
            scale = 4.0

        fig = plot_img(img, dpi=72 * scale, scale=scale)
        ax = plt.gca()

        masks = [np.asarray(m, dtype=bool) for m in data.instance_masks]
        masks = [m[..., 0] if m.ndim == 3 else m for m in masks]
        if masks:
            h, w = masks[0].shape[:2]
            mask_output_scale = h / img.shape[0]
            extent = [
                -0.5,
                w / mask_output_scale - 0.5,
                h / mask_output_scale - 0.5,
                -0.5,
            ]
            # Build a translucent RGBA overlay: one color per instance, alpha only
            # where that instance's mask is set (background stays transparent).
            overlay = np.zeros((h, w, 4), dtype=np.float32)
            for i, m in enumerate(masks):
                if m.shape[:2] != (h, w):
                    continue
                r, g, b = self._INSTANCE_COLORS[i % len(self._INSTANCE_COLORS)]
                overlay[m, 0] = r
                overlay[m, 1] = g
                overlay[m, 2] = b
                overlay[m, 3] = 0.5
            ax.imshow(overlay, origin="upper", extent=extent, interpolation="nearest")

        # Optional GT foreground silhouette outline for reference.
        if data.gt_mask is not None:
            gt = np.asarray(data.gt_mask)
            if gt.ndim == 3:
                gt = gt[..., 0]
            gt = gt.astype(np.float32)
            gt_output_scale = gt.shape[0] / img.shape[0]
            gextent = [
                -0.5,
                gt.shape[1] / gt_output_scale - 0.5,
                gt.shape[0] / gt_output_scale - 0.5,
                -0.5,
            ]
            ax.contour(
                np.linspace(gextent[0], gextent[1], gt.shape[1]),
                np.linspace(gextent[3], gextent[2], gt.shape[0]),
                gt,
                levels=[0.5],
                colors="white",
                linewidths=1.0,
            )
        return fig

render(data)

Render visualization data to a matplotlib figure.

Parameters:

Name	Type	Description	Default
data	VisualizationData	VisualizationData containing image, confmaps, peaks, etc.	required

Returns:

Type	Description
Figure	A matplotlib Figure object.

Source code in sleap_nn/training/utils.py

def render(self, data: VisualizationData) -> matplotlib.figure.Figure:
    """Render visualization data to a matplotlib figure.

    Args:
        data: VisualizationData containing image, confmaps, peaks, etc.

    Returns:
        A matplotlib Figure object.
    """
    img = data.image
    scale = 1.0
    if img.shape[0] < 512:
        scale = 2.0
    if img.shape[0] < 256:
        scale = 4.0

    fig = plot_img(img, dpi=72 * scale, scale=scale)
    plot_confmaps(data.pred_confmaps, output_scale=data.output_scale)
    plot_peaks(data.gt_instances, data.pred_peaks, paired=data.is_paired)
    return fig

render_gt_mask(data)

Render the predicted foreground vs the ground-truth mask overlay.

Draws the image, overlays the predicted foreground probability (as for the standard confmap render), then outlines the ground-truth foreground mask so the prediction can be compared against truth each epoch (the segmentation analog of pose's GT-keypoint overlay).

Parameters:

Name	Type	Description	Default
data	VisualizationData	VisualizationData with gt_mask populated.	required

Returns:

Type	Description
Figure	A matplotlib Figure object.

Source code in sleap_nn/training/utils.py

def render_gt_mask(self, data: VisualizationData) -> matplotlib.figure.Figure:
    """Render the predicted foreground vs the ground-truth mask overlay.

    Draws the image, overlays the predicted foreground probability (as for the
    standard confmap render), then outlines the ground-truth foreground mask so
    the prediction can be compared against truth each epoch (the segmentation
    analog of pose's GT-keypoint overlay).

    Args:
        data: VisualizationData with ``gt_mask`` populated.

    Returns:
        A matplotlib Figure object.
    """
    if data.gt_mask is None:
        raise ValueError("gt_mask is None, cannot render GT mask overlay")

    img = data.image
    scale = 1.0
    if img.shape[0] < 512:
        scale = 2.0
    if img.shape[0] < 256:
        scale = 4.0

    gt = np.asarray(data.gt_mask)
    if gt.ndim == 3:
        gt = gt[..., 0]
    gt = gt.astype(np.float32)

    fig = plot_img(img, dpi=72 * scale, scale=scale)
    # Predicted foreground probability as a translucent heatmap (same overlay
    # the confmap render uses), so pred and GT are visible together.
    plot_confmaps(data.pred_confmaps, output_scale=data.output_scale)
    ax = plt.gca()
    gt_output_scale = gt.shape[0] / img.shape[0]
    extent = [
        -0.5,
        gt.shape[1] / gt_output_scale - 0.5,
        gt.shape[0] / gt_output_scale - 0.5,
        -0.5,
    ]
    # GT foreground outline (single contour at the 0.5 level).
    ax.contour(
        np.linspace(extent[0], extent[1], gt.shape[1]),
        np.linspace(extent[3], extent[2], gt.shape[0]),
        gt,
        levels=[0.5],
        colors="lime",
        linewidths=1.5,
    )
    return fig

render_instance_masks(data)

Render grouped predicted per-instance masks as a colored overlay.

Each predicted instance's mask (from the offset-grouping that already runs in get_visualization_data) is drawn in a distinct color on the input image, so instance separation — the model's actual output, and the single most informative seg diagnostic — is visible each epoch. When a GT foreground mask is present (gt_mask), its silhouette is outlined for reference. Empty instance lists render the plain image (no instances detected this epoch).

Parameters:

Name	Type	Description	Default
data	VisualizationData	VisualizationData with instance_masks populated (a list of (H, W) boolean masks at output-stride resolution).	required

Returns:

Type	Description
Figure	A matplotlib Figure object.

Source code in sleap_nn/training/utils.py

def render_instance_masks(
    self, data: VisualizationData
) -> matplotlib.figure.Figure:
    """Render grouped predicted per-instance masks as a colored overlay.

    Each predicted instance's mask (from the offset-grouping that already runs
    in ``get_visualization_data``) is drawn in a distinct color on the input
    image, so instance separation — the model's actual output, and the single
    most informative seg diagnostic — is visible each epoch. When a GT
    foreground mask is present (``gt_mask``), its silhouette is outlined for
    reference. Empty instance lists render the plain image (no instances
    detected this epoch).

    Args:
        data: VisualizationData with ``instance_masks`` populated (a list of
            ``(H, W)`` boolean masks at output-stride resolution).

    Returns:
        A matplotlib Figure object.
    """
    if data.instance_masks is None:
        raise ValueError(
            "instance_masks is None, cannot render instance-mask overlay"
        )

    img = data.image
    scale = 1.0
    if img.shape[0] < 512:
        scale = 2.0
    if img.shape[0] < 256:
        scale = 4.0

    fig = plot_img(img, dpi=72 * scale, scale=scale)
    ax = plt.gca()

    masks = [np.asarray(m, dtype=bool) for m in data.instance_masks]
    masks = [m[..., 0] if m.ndim == 3 else m for m in masks]
    if masks:
        h, w = masks[0].shape[:2]
        mask_output_scale = h / img.shape[0]
        extent = [
            -0.5,
            w / mask_output_scale - 0.5,
            h / mask_output_scale - 0.5,
            -0.5,
        ]
        # Build a translucent RGBA overlay: one color per instance, alpha only
        # where that instance's mask is set (background stays transparent).
        overlay = np.zeros((h, w, 4), dtype=np.float32)
        for i, m in enumerate(masks):
            if m.shape[:2] != (h, w):
                continue
            r, g, b = self._INSTANCE_COLORS[i % len(self._INSTANCE_COLORS)]
            overlay[m, 0] = r
            overlay[m, 1] = g
            overlay[m, 2] = b
            overlay[m, 3] = 0.5
        ax.imshow(overlay, origin="upper", extent=extent, interpolation="nearest")

    # Optional GT foreground silhouette outline for reference.
    if data.gt_mask is not None:
        gt = np.asarray(data.gt_mask)
        if gt.ndim == 3:
            gt = gt[..., 0]
        gt = gt.astype(np.float32)
        gt_output_scale = gt.shape[0] / img.shape[0]
        gextent = [
            -0.5,
            gt.shape[1] / gt_output_scale - 0.5,
            gt.shape[0] / gt_output_scale - 0.5,
            -0.5,
        ]
        ax.contour(
            np.linspace(gextent[0], gextent[1], gt.shape[1]),
            np.linspace(gextent[3], gextent[2], gt.shape[0]),
            gt,
            levels=[0.5],
            colors="white",
            linewidths=1.0,
        )
    return fig

render_offsets(data)

Render the center-offset field DIRECTION for segmentation models.

Colors each pixel by the ANGLE of its center-offset vector (hue = atan2(dy, dx), the standard optical-flow color wheel), so the field reads as "which way does each foreground pixel point to its instance center." Opacity is modulated by the offset magnitude so near-zero (background) offsets stay transparent and only meaningful vectors show their direction.

Parameters:

Name	Type	Description	Default
data	VisualizationData	VisualizationData with pred_offsets (H, W, 2) populated, laid out as [..., 0] = dx and [..., 1] = dy.	required

Returns:

Type	Description
Figure	A matplotlib Figure object showing the per-pixel offset direction.

Source code in sleap_nn/training/utils.py

def render_offsets(self, data: VisualizationData) -> matplotlib.figure.Figure:
    """Render the center-offset field DIRECTION for segmentation models.

    Colors each pixel by the ANGLE of its center-offset vector (hue =
    ``atan2(dy, dx)``, the standard optical-flow color wheel), so the field
    reads as "which way does each foreground pixel point to its instance
    center." Opacity is modulated by the offset magnitude so near-zero
    (background) offsets stay transparent and only meaningful vectors show
    their direction.

    Args:
        data: VisualizationData with ``pred_offsets`` (H, W, 2) populated,
            laid out as ``[..., 0] = dx`` and ``[..., 1] = dy``.

    Returns:
        A matplotlib Figure object showing the per-pixel offset direction.
    """
    if data.pred_offsets is None:
        raise ValueError("pred_offsets is None, cannot render offsets")

    import matplotlib.colors as mcolors

    img = data.image
    scale = 1.0
    if img.shape[0] < 512:
        scale = 2.0
    if img.shape[0] < 256:
        scale = 4.0

    offsets = data.pred_offsets  # (H, W, 2) -> (dx, dy)
    dx, dy = offsets[..., 0], offsets[..., 1]
    angle = np.arctan2(dy, dx)  # [-pi, pi]
    magnitude = np.hypot(dx, dy)

    # Hue = direction (wrapped to [0, 1]); full saturation/value for vivid color.
    hue = (angle + np.pi) / (2.0 * np.pi)
    hsv = np.stack([hue, np.ones_like(hue), np.ones_like(hue)], axis=-1)
    rgb = mcolors.hsv_to_rgb(hsv)  # (H, W, 3)
    # Opacity ramps with magnitude so background (~0 offset) is transparent.
    mmax = float(magnitude.max())
    alpha = (magnitude / mmax if mmax > 0 else magnitude) * 0.7
    rgba = np.dstack([rgb, alpha])  # (H, W, 4)

    fig = plot_img(img, dpi=72 * scale, scale=scale)
    ax = plt.gca()
    offset_output_scale = rgba.shape[0] / img.shape[0]
    ax.imshow(
        rgba,
        origin="upper",
        extent=[
            -0.5,
            rgba.shape[1] / offset_output_scale - 0.5,
            rgba.shape[0] / offset_output_scale - 0.5,
            -0.5,
        ],
    )
    return fig

render_pafs(data)

Render PAF magnitude visualization.

Parameters:

Name	Type	Description	Default
data	VisualizationData	VisualizationData with pred_pafs populated.	required

Returns:

Type	Description
Figure	A matplotlib Figure object showing PAF magnitudes.

Source code in sleap_nn/training/utils.py

def render_pafs(self, data: VisualizationData) -> matplotlib.figure.Figure:
    """Render PAF magnitude visualization.

    Args:
        data: VisualizationData with pred_pafs populated.

    Returns:
        A matplotlib Figure object showing PAF magnitudes.
    """
    if data.pred_pafs is None:
        raise ValueError("pred_pafs is None, cannot render PAFs")

    img = data.image
    scale = 1.0
    if img.shape[0] < 512:
        scale = 2.0
    if img.shape[0] < 256:
        scale = 4.0

    # Compute PAF magnitude
    pafs = data.pred_pafs  # (H, W, 2*edges) or (H, W, edges, 2)
    if pafs.ndim == 3:
        n_edges = pafs.shape[-1] // 2
        pafs = pafs.reshape(pafs.shape[0], pafs.shape[1], n_edges, 2)
    magnitude = np.sqrt(pafs[..., 0] ** 2 + pafs[..., 1] ** 2)
    magnitude = magnitude.max(axis=-1)  # Max over edges

    fig = plot_img(img, dpi=72 * scale, scale=scale)
    ax = plt.gca()

    # Calculate PAF output scale from actual PAF dimensions, not confmap output_scale
    # PAFs may have a different output_stride than confmaps
    paf_output_scale = magnitude.shape[0] / img.shape[0]

    ax.imshow(
        magnitude,
        alpha=0.5,
        origin="upper",
        cmap="viridis",
        extent=[
            -0.5,
            magnitude.shape[1] / paf_output_scale - 0.5,
            magnitude.shape[0] / paf_output_scale - 0.5,
            -0.5,
        ],
    )
    return fig

VisualizationData dataclass

Container for visualization data from a single sample.

This dataclass decouples data extraction from rendering, allowing the same data to be rendered to different output targets (matplotlib, wandb, etc.).

Attributes:

Name	Type	Description
image	ndarray	Input image as (H, W, C) numpy array, normalized to [0, 1].
pred_confmaps	ndarray	Predicted confidence maps as (H, W, nodes) array, values in [0, 1].
pred_peaks	ndarray	Predicted keypoints as (instances, nodes, 2) or (nodes, 2) array.
pred_peak_values	ndarray	Confidence values as (instances, nodes) or (nodes,) array.
gt_instances	ndarray	Ground truth keypoints, same shape as pred_peaks.
node_names	List[str]	List of node/keypoint names, e.g., ["head", "thorax", ...].
output_scale	float	Ratio of confmap size to image size (confmap_h / image_h).
is_paired	bool	Whether GT and predictions can be paired for error visualization.
pred_pafs	Optional[ndarray]	Part affinity fields for bottom-up models, optional.
pred_class_maps	Optional[ndarray]	Class maps for multi-class models, optional.
pred_center_heatmap	Optional[ndarray]	Center heatmap for segmentation models, optional.
pred_offsets	Optional[ndarray]	Center-offset field (H, W, 2) for segmentation models, optional. Rendered as a per-pixel offset-magnitude map.
gt_mask	Optional[ndarray]	Ground-truth foreground mask (H, W) or (H, W, 1) for segmentation models, optional. Rendered as a translucent overlay so the predicted foreground probability can be compared against truth each epoch.
instance_masks	Optional[List[ndarray]]	Grouped predicted per-instance masks for bottom-up segmentation, as a list of (H, W) boolean arrays at output-stride resolution (one per detected instance), optional. Rendered as a colored overlay (one color per instance) so instance separation — the model's actual output — is visible each epoch. These are the masks the offset-grouping in get_visualization_data already computes.

Source code in sleap_nn/training/utils.py

@dataclass
class VisualizationData:
    """Container for visualization data from a single sample.

    This dataclass decouples data extraction from rendering, allowing the same
    data to be rendered to different output targets (matplotlib, wandb, etc.).

    Attributes:
        image: Input image as (H, W, C) numpy array, normalized to [0, 1].
        pred_confmaps: Predicted confidence maps as (H, W, nodes) array, values in [0, 1].
        pred_peaks: Predicted keypoints as (instances, nodes, 2) or (nodes, 2) array.
        pred_peak_values: Confidence values as (instances, nodes) or (nodes,) array.
        gt_instances: Ground truth keypoints, same shape as pred_peaks.
        node_names: List of node/keypoint names, e.g., ["head", "thorax", ...].
        output_scale: Ratio of confmap size to image size (confmap_h / image_h).
        is_paired: Whether GT and predictions can be paired for error visualization.
        pred_pafs: Part affinity fields for bottom-up models, optional.
        pred_class_maps: Class maps for multi-class models, optional.
        pred_center_heatmap: Center heatmap for segmentation models, optional.
        pred_offsets: Center-offset field (H, W, 2) for segmentation models,
            optional. Rendered as a per-pixel offset-magnitude map.
        gt_mask: Ground-truth foreground mask (H, W) or (H, W, 1) for segmentation
            models, optional. Rendered as a translucent overlay so the predicted
            foreground probability can be compared against truth each epoch.
        instance_masks: Grouped predicted per-instance masks for bottom-up
            segmentation, as a list of ``(H, W)`` boolean arrays at output-stride
            resolution (one per detected instance), optional. Rendered as a
            colored overlay (one color per instance) so instance separation — the
            model's actual output — is visible each epoch. These are the masks the
            offset-grouping in ``get_visualization_data`` already computes.
    """

    image: np.ndarray
    pred_confmaps: np.ndarray
    pred_peaks: np.ndarray
    pred_peak_values: np.ndarray
    gt_instances: np.ndarray
    node_names: List[str] = field(default_factory=list)
    output_scale: float = 1.0
    is_paired: bool = True
    pred_pafs: Optional[np.ndarray] = None
    pred_class_maps: Optional[np.ndarray] = None
    pred_center_heatmap: Optional[np.ndarray] = None
    pred_offsets: Optional[np.ndarray] = None
    gt_mask: Optional[np.ndarray] = None
    instance_masks: Optional[List[np.ndarray]] = None

WandBRenderer

Renders VisualizationData to wandb.Image objects.

Supports multiple rendering modes: - "direct": Pre-render with matplotlib, convert to wandb.Image - "boxes": Use wandb boxes for interactive keypoint visualization - "masks": Use wandb masks for confidence map overlay

Methods:

Name	Description
__init__	Initialize the renderer.
render	Render visualization data to a wandb.Image.

Source code in sleap_nn/training/utils.py

class WandBRenderer:
    """Renders VisualizationData to wandb.Image objects.

    Supports multiple rendering modes:
    - "direct": Pre-render with matplotlib, convert to wandb.Image
    - "boxes": Use wandb boxes for interactive keypoint visualization
    - "masks": Use wandb masks for confidence map overlay
    """

    def __init__(
        self,
        mode: str = "direct",
        box_size: float = 5.0,
        confmap_threshold: float = 0.1,
        min_size: int = 512,
    ):
        """Initialize the renderer.

        Args:
            mode: Rendering mode - "direct", "boxes", or "masks".
            box_size: Size of keypoint boxes in pixels (for "boxes" mode).
            confmap_threshold: Threshold for confmap mask (for "masks" mode).
            min_size: Minimum image dimension. Smaller images will be upscaled.
        """
        self.mode = mode
        self.box_size = box_size
        self.confmap_threshold = confmap_threshold
        self.min_size = min_size
        self._mpl_renderer = MatplotlibRenderer()

    def render(
        self, data: VisualizationData, caption: Optional[str] = None
    ) -> "wandb.Image":
        """Render visualization data to a wandb.Image.

        Args:
            data: VisualizationData containing image, confmaps, peaks, etc.
            caption: Optional caption for the image.

        Returns:
            A wandb.Image object.
        """
        import wandb

        if self.mode == "boxes":
            return self._render_with_boxes(data, caption)
        elif self.mode == "masks":
            return self._render_with_masks(data, caption)
        else:  # "direct"
            return self._render_direct(data, caption)

    def _get_scale_factor(self, img_h: int, img_w: int) -> int:
        """Calculate scale factor to ensure minimum image size."""
        min_dim = min(img_h, img_w)
        if min_dim >= self.min_size:
            return 1
        return int(np.ceil(self.min_size / min_dim))

    def _render_direct(
        self, data: VisualizationData, caption: Optional[str] = None
    ) -> "wandb.Image":
        """Pre-render with matplotlib, return as wandb.Image."""
        import wandb
        from PIL import Image

        fig = self._mpl_renderer.render(data)

        # Convert figure to PIL Image
        buf = BytesIO()
        fig.savefig(buf, format="png", bbox_inches="tight", pad_inches=0)
        buf.seek(0)
        plt.close(fig)

        pil_image = Image.open(buf)
        return wandb.Image(pil_image, caption=caption)

    def _render_with_boxes(
        self, data: VisualizationData, caption: Optional[str] = None
    ) -> "wandb.Image":
        """Use wandb boxes for interactive keypoint visualization."""
        import wandb
        from PIL import Image

        # Prepare class labels from node names
        class_labels = {i: name for i, name in enumerate(data.node_names)}
        if not class_labels:
            class_labels = {i: f"node_{i}" for i in range(data.pred_peaks.shape[-2])}

        # Convert image to uint8
        img_uint8 = (np.clip(data.image, 0, 1) * 255).astype(np.uint8)
        # Handle single-channel images: squeeze (H, W, 1) -> (H, W)
        if img_uint8.ndim == 3 and img_uint8.shape[2] == 1:
            img_uint8 = img_uint8.squeeze(axis=2)
        img_h, img_w = img_uint8.shape[:2]

        # Scale up small images for better visibility in wandb
        scale = self._get_scale_factor(img_h, img_w)
        if scale > 1:
            pil_img = Image.fromarray(img_uint8)
            pil_img = pil_img.resize(
                (img_w * scale, img_h * scale), resample=Image.BILINEAR
            )
            img_uint8 = np.array(pil_img)

        # Build ground truth boxes (use percent domain for proper thumbnail scaling)
        gt_box_data = self._peaks_to_boxes(
            data.gt_instances, data.node_names, img_w, img_h, is_gt=True
        )

        # Build prediction boxes
        pred_box_data = self._peaks_to_boxes(
            data.pred_peaks,
            data.node_names,
            img_w,
            img_h,
            peak_values=data.pred_peak_values,
            is_gt=False,
        )

        return wandb.Image(
            img_uint8,
            boxes={
                "ground_truth": {"box_data": gt_box_data, "class_labels": class_labels},
                "predictions": {
                    "box_data": pred_box_data,
                    "class_labels": class_labels,
                },
            },
            caption=caption,
        )

    def _peaks_to_boxes(
        self,
        peaks: np.ndarray,
        node_names: List[str],
        img_w: int,
        img_h: int,
        peak_values: Optional[np.ndarray] = None,
        is_gt: bool = False,
    ) -> List[dict]:
        """Convert peaks array to wandb box_data format.

        Args:
            peaks: Keypoints as (instances, nodes, 2) or (nodes, 2).
            node_names: List of node names.
            img_w: Image width in pixels.
            img_h: Image height in pixels.
            peak_values: Optional confidence values.
            is_gt: Whether these are ground truth points.

        Returns:
            List of box dictionaries for wandb.
        """
        box_data = []

        # Normalize shape to (instances, nodes, 2)
        if peak_values is not None:
            peak_values = np.asarray(peak_values)
        if peaks.ndim == 2:
            peaks = peaks[np.newaxis, ...]

        if peak_values is not None and peak_values.ndim == 1:
            if peaks.shape[0] == 1 and peak_values.shape[0] == peaks.shape[1]:
                peak_values = peak_values[np.newaxis, ...]
            elif peaks.shape[1] == 1 and peak_values.shape[0] == peaks.shape[0]:
                peak_values = peak_values[:, np.newaxis]

        # Convert box_size from pixels to percent
        box_w_pct = self.box_size / img_w
        box_h_pct = self.box_size / img_h

        for inst_idx, instance in enumerate(peaks):
            for node_idx, (x, y) in enumerate(instance):
                if np.isnan(x) or np.isnan(y):
                    continue

                node_name = (
                    node_names[node_idx]
                    if node_idx < len(node_names)
                    else f"node_{node_idx}"
                )

                # Convert pixel coordinates to percent (0-1 range)
                x_pct = float(x) / img_w
                y_pct = float(y) / img_h

                box = {
                    "position": {
                        "middle": [x_pct, y_pct],
                        "width": box_w_pct,
                        "height": box_h_pct,
                    },
                    "domain": "percent",
                    "class_id": node_idx,
                }

                if is_gt:
                    box["box_caption"] = f"GT: {node_name}"
                else:
                    if peak_values is not None:
                        conf = float(peak_values[inst_idx, node_idx])
                        box["box_caption"] = f"{node_name} ({conf:.2f})"
                        box["scores"] = {"confidence": conf}
                    else:
                        box["box_caption"] = node_name

                box_data.append(box)

        return box_data

    def _render_with_masks(
        self, data: VisualizationData, caption: Optional[str] = None
    ) -> "wandb.Image":
        """Use wandb masks for confidence map overlay.

        Uses argmax approach: each pixel shows the dominant node.
        """
        import wandb

        # Prepare class labels (0 = background, 1+ = nodes)
        class_labels = {0: "background"}
        for i, name in enumerate(data.node_names):
            class_labels[i + 1] = name
        if not data.node_names:
            n_nodes = data.pred_confmaps.shape[-1]
            for i in range(n_nodes):
                class_labels[i + 1] = f"node_{i}"

        # Create argmax mask from confmaps
        confmaps = data.pred_confmaps  # (H/stride, W/stride, nodes)
        max_vals = confmaps.max(axis=-1)
        argmax_map = confmaps.argmax(axis=-1) + 1  # +1 for background offset
        argmax_map[max_vals < self.confmap_threshold] = 0  # Background

        # Convert image to uint8
        img_uint8 = (np.clip(data.image, 0, 1) * 255).astype(np.uint8)
        # Handle single-channel images: (H, W, 1) -> (H, W)
        if img_uint8.ndim == 3 and img_uint8.shape[2] == 1:
            img_uint8 = img_uint8.squeeze(axis=2)
        img_h, img_w = img_uint8.shape[:2]

        # Resize mask to match image dimensions (confmaps are H/stride, W/stride)
        from PIL import Image

        mask_pil = Image.fromarray(argmax_map.astype(np.uint8))
        mask_pil = mask_pil.resize((img_w, img_h), resample=Image.NEAREST)
        argmax_map = np.array(mask_pil)

        return wandb.Image(
            img_uint8,
            masks={
                "confidence_maps": {
                    "mask_data": argmax_map.astype(np.uint8),
                    "class_labels": class_labels,
                }
            },
            caption=caption,
        )

__init__(mode='direct', box_size=5.0, confmap_threshold=0.1, min_size=512)

Initialize the renderer.

Parameters:

Name	Type	Description	Default
mode	str	Rendering mode - "direct", "boxes", or "masks".	'direct'
box_size	float	Size of keypoint boxes in pixels (for "boxes" mode).	5.0
confmap_threshold	float	Threshold for confmap mask (for "masks" mode).	0.1
min_size	int	Minimum image dimension. Smaller images will be upscaled.	512

Source code in sleap_nn/training/utils.py

def __init__(
    self,
    mode: str = "direct",
    box_size: float = 5.0,
    confmap_threshold: float = 0.1,
    min_size: int = 512,
):
    """Initialize the renderer.

    Args:
        mode: Rendering mode - "direct", "boxes", or "masks".
        box_size: Size of keypoint boxes in pixels (for "boxes" mode).
        confmap_threshold: Threshold for confmap mask (for "masks" mode).
        min_size: Minimum image dimension. Smaller images will be upscaled.
    """
    self.mode = mode
    self.box_size = box_size
    self.confmap_threshold = confmap_threshold
    self.min_size = min_size
    self._mpl_renderer = MatplotlibRenderer()

render(data, caption=None)

Render visualization data to a wandb.Image.

Parameters:

Name	Type	Description	Default
data	VisualizationData	VisualizationData containing image, confmaps, peaks, etc.	required
caption	Optional[str]	Optional caption for the image.	None

Returns:

Type	Description
Image	A wandb.Image object.

Source code in sleap_nn/training/utils.py

def render(
    self, data: VisualizationData, caption: Optional[str] = None
) -> "wandb.Image":
    """Render visualization data to a wandb.Image.

    Args:
        data: VisualizationData containing image, confmaps, peaks, etc.
        caption: Optional caption for the image.

    Returns:
        A wandb.Image object.
    """
    import wandb

    if self.mode == "boxes":
        return self._render_with_boxes(data, caption)
    elif self.mode == "masks":
        return self._render_with_masks(data, caption)
    else:  # "direct"
        return self._render_direct(data, caption)

get_dist_rank()

Return the rank of the current process if torch.distributed is initialized.

Source code in sleap_nn/training/utils.py

def get_dist_rank():
    """Return the rank of the current process if torch.distributed is initialized."""
    return dist.get_rank() if is_distributed_initialized() else None

get_gpu_memory()

Get the available memory on each GPU.

Returns:

Type	Description
List[int]	A list of the available memory on each GPU in MiB.

Source code in sleap_nn/training/utils.py

def get_gpu_memory() -> List[int]:
    """Get the available memory on each GPU.

    Returns:
        A list of the available memory on each GPU in MiB.
    """
    if shutil.which("nvidia-smi") is None:
        return []

    command = [
        "nvidia-smi",
        "--query-gpu=index,memory.free",
        "--format=csv",
    ]

    try:
        memory_poll = subprocess.run(command, capture_output=True)
    except (subprocess.SubprocessError, FileNotFoundError):
        return []

    subprocess_result = memory_poll.stdout
    memory_string = subprocess_result.decode("ascii").split("\n")[1:-1]

    if "CUDA_VISIBLE_DEVICES" in os.environ.keys():
        cuda_visible_devices = os.environ["CUDA_VISIBLE_DEVICES"].split(",")
    else:
        cuda_visible_devices = None

    memory_list = []
    for row in memory_string:
        gpu_index, available_memory = row.split(", ")
        available_memory = available_memory.split(" MiB")[0]

        if cuda_visible_devices is None or gpu_index in cuda_visible_devices:
            memory_list.append(int(available_memory))

    return memory_list

imgfig(size=6, dpi=72, scale=1.0)

Create a tight figure for image plotting.

Parameters:

Name	Type	Description	Default
size	float | tuple	Scalar or 2-tuple specifying the (width, height) of the figure in inches. If scalar, will assume equal width and height.	6
dpi	int | float	Dots per inch, controlling the resolution of the image.	72
scale	float	Factor to scale the size of the figure by. This is a convenience for increasing the size of the plot at the same DPI.	1.0

Returns:

Type	Description
Figure	A matplotlib.figure.Figure to use for plotting.

Source code in sleap_nn/training/utils.py

def imgfig(
    size: float | tuple = 6, dpi: int | float = 72, scale: float = 1.0
) -> matplotlib.figure.Figure:
    """Create a tight figure for image plotting.

    Args:
        size: Scalar or 2-tuple specifying the (width, height) of the figure in inches.
            If scalar, will assume equal width and height.
        dpi: Dots per inch, controlling the resolution of the image.
        scale: Factor to scale the size of the figure by. This is a convenience for
            increasing the size of the plot at the same DPI.

    Returns:
        A matplotlib.figure.Figure to use for plotting.
    """
    if not isinstance(size, (tuple, list)):
        size = (size, size)
    fig = plt.figure(figsize=(scale * size[0], scale * size[1]), dpi=dpi)
    ax = fig.add_axes([0, 0, 1, 1], frameon=False)
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)
    plt.autoscale(tight=True)
    ax.set_xticks([])
    ax.set_yticks([])
    ax.grid(False)
    return fig

is_distributed_initialized()

Check if distributed processes are initialized.

Source code in sleap_nn/training/utils.py

def is_distributed_initialized():
    """Check if distributed processes are initialized."""
    return dist.is_available() and dist.is_initialized()

plot_confmaps(confmaps, output_scale=1.0)

Plot confidence maps reduced over channels.

Parameters:

Name	Type	Description	Default
confmaps	ndarray	Confidence maps to plot with shape (height, width, channel).	required
output_scale	float	Factor to scale the size of the figure by.	1.0

Returns:

Type	Description
	A matplotlib.figure.Figure to use for plotting.

Source code in sleap_nn/training/utils.py

def plot_confmaps(confmaps: np.ndarray, output_scale: float = 1.0):
    """Plot confidence maps reduced over channels.

    Args:
        confmaps: Confidence maps to plot with shape (height, width, channel).
        output_scale: Factor to scale the size of the figure by.

    Returns:
        A matplotlib.figure.Figure to use for plotting.
    """
    ax = plt.gca()
    return ax.imshow(
        np.squeeze(confmaps.max(axis=-1)),
        alpha=0.5,
        origin="upper",
        vmin=0,
        vmax=1,
        extent=[
            -0.5,
            confmaps.shape[1] / output_scale - 0.5,
            confmaps.shape[0] / output_scale - 0.5,
            -0.5,
        ],
    )

plot_img(img, dpi=72, scale=1.0)

Plot an image in a tight figure.

Parameters:

Name	Type	Description	Default
img	ndarray	Image to plot of shape (height, width, channel).	required
dpi	int | float	Dots per inch, controlling the resolution of the image.	72
scale	float	Factor to scale the size of the figure by. This is a convenience for increasing the size of the plot at the same DPI.	1.0

Returns:

Type	Description
Figure	A matplotlib.figure.Figure to use for plotting.

Source code in sleap_nn/training/utils.py

def plot_img(
    img: np.ndarray, dpi: int | float = 72, scale: float = 1.0
) -> matplotlib.figure.Figure:
    """Plot an image in a tight figure.

    Args:
        img: Image to plot of shape (height, width, channel).
        dpi: Dots per inch, controlling the resolution of the image.
        scale: Factor to scale the size of the figure by. This is a convenience for
            increasing the size of the plot at the same DPI.

    Returns:
        A matplotlib.figure.Figure to use for plotting.
    """
    if hasattr(img, "numpy"):
        img = img.numpy()

    if img.shape[0] == 1:
        # Squeeze out batch singleton dimension.
        img = img.squeeze(axis=0)

    # Check if image is grayscale (single channel).
    grayscale = img.shape[-1] == 1
    if grayscale:
        # Squeeze out singleton channel.
        img = img.squeeze(axis=-1)

    # Normalize the range of pixel values.
    img_min = img.min()
    img_max = img.max()
    if img_min < 0.0 or img_max > 1.0:
        img = (img - img_min) / (img_max - img_min)

    fig = imgfig(
        size=(float(img.shape[1]) / dpi, float(img.shape[0]) / dpi),
        dpi=dpi,
        scale=scale,
    )

    ax = fig.gca()
    ax.imshow(
        img,
        cmap="gray" if grayscale else None,
        origin="upper",
        extent=[-0.5, img.shape[1] - 0.5, img.shape[0] - 0.5, -0.5],
    )
    return fig

plot_peaks(pts_gt, pts_pr=None, paired=False)

Plot ground truth and detected peaks.

Parameters:

Name	Type	Description	Default
pts_gt	ndarray	Ground-truth keypoints of shape (num_instances, nodes, 2). To plot centroids, shape: (1, num_instances, 2).	required
pts_pr	ndarray | None	Predicted keypoints of shape (num_instances, nodes, 2). To plot centroids, shape: (1, num_instances, 2)	None
paired	bool	True if error lines should be plotted else False.	False

Returns:

Type	Description
	A matplotlib.figure.Figure to use for plotting.

Source code in sleap_nn/training/utils.py

def plot_peaks(
    pts_gt: np.ndarray, pts_pr: np.ndarray | None = None, paired: bool = False
):
    """Plot ground truth and detected peaks.

    Args:
        pts_gt: Ground-truth keypoints of shape (num_instances, nodes, 2). To plot centroids, shape: (1, num_instances, 2).
        pts_pr: Predicted keypoints of shape (num_instances, nodes, 2). To plot centroids, shape: (1, num_instances, 2)
        paired: True if error lines should be plotted else False.

    Returns:
        A matplotlib.figure.Figure to use for plotting.
    """
    handles = []
    ax = plt.gca()
    if paired and pts_pr is not None:
        for pt_gt, pt_pr in zip(pts_gt, pts_pr):
            for p_gt, p_pr in zip(pt_gt, pt_pr):
                handles.append(
                    ax.plot(
                        [p_gt[0], p_pr[0]], [p_gt[1], p_pr[1]], "r-", alpha=0.5, lw=2
                    )
                )
    if pts_pr is not None:
        handles.append(
            ax.plot(
                pts_gt[..., 0].ravel(),
                pts_gt[..., 1].ravel(),
                "g.",
                alpha=0.7,
                ms=10,
                mew=1,
                mec="w",
            )
        )
        handles.append(
            ax.plot(
                pts_pr[..., 0].ravel(),
                pts_pr[..., 1].ravel(),
                "r.",
                alpha=0.7,
                ms=10,
                mew=1,
                mec="w",
            )
        )
    else:
        cmap = sns.color_palette("tab20")
        for i, pts in enumerate(pts_gt):
            handles.append(
                ax.plot(
                    pts[:, 0],
                    pts[:, 1],
                    ".",
                    alpha=0.7,
                    ms=15,
                    mew=1,
                    mfc=cmap[i % len(cmap)],
                    mec="w",
                )
            )
    return handles

plot_tile_grid(img, tile_size, overlap, output_stride, max_stride=1, min_overlap_fraction=0.25, blend='gaussian', sigma_scale=0.125, show_coverage=True, dpi=72, scale=1.0)

Visualize the sliding-window tile grid (and blend coverage) over a frame.

Draws each tile's extent as a rectangle over the image and, optionally, overlays the summed importance-window coverage (the CNT denominator used by the tiled-inference stitcher) as a heatmap so seams and edge down-weighting are visible. Useful for choosing tile_size / overlap and for documentation figures.

Parameters:

Name	Type	Description	Default
img	ndarray	Frame image; accepts (H, W), (H, W, C), (C, H, W) or a leading singleton batch, as a numpy array or torch tensor.	required
tile_size	int	Square tile side length in pixels.	required
overlap	int	Tile overlap in pixels.	required
output_stride	int	Model output stride (grid snapping + coverage resolution).	required
max_stride	int	Backbone max stride (used to snap the inter-tile step).	1
min_overlap_fraction	float	Minimum overlap as a fraction of tile_size.	0.25
blend	str	Importance-window mode, one of "gaussian"/"pyramid"/"constant".	'gaussian'
sigma_scale	float	Gaussian importance-window sigma scale.	0.125
show_coverage	bool	If True, overlay the summed-window coverage heatmap.	True
dpi	int | float	Dots per inch passed to :func:plot_img.	72
scale	float	Figure size scale passed to :func:plot_img.	1.0

Returns:

Type	Description
Figure	A matplotlib.figure.Figure with the image, tile rectangles and (optionally) the coverage heatmap.

Source code in sleap_nn/training/utils.py

def plot_tile_grid(
    img: np.ndarray,
    tile_size: int,
    overlap: int,
    output_stride: int,
    max_stride: int = 1,
    min_overlap_fraction: float = 0.25,
    blend: str = "gaussian",
    sigma_scale: float = 0.125,
    show_coverage: bool = True,
    dpi: int | float = 72,
    scale: float = 1.0,
) -> matplotlib.figure.Figure:
    """Visualize the sliding-window tile grid (and blend coverage) over a frame.

    Draws each tile's extent as a rectangle over the image and, optionally,
    overlays the summed importance-window coverage (the ``CNT`` denominator used
    by the tiled-inference stitcher) as a heatmap so seams and edge
    down-weighting are visible. Useful for choosing ``tile_size`` / ``overlap``
    and for documentation figures.

    Args:
        img: Frame image; accepts ``(H, W)``, ``(H, W, C)``, ``(C, H, W)`` or a
            leading singleton batch, as a numpy array or torch tensor.
        tile_size: Square tile side length in pixels.
        overlap: Tile overlap in pixels.
        output_stride: Model output stride (grid snapping + coverage resolution).
        max_stride: Backbone max stride (used to snap the inter-tile step).
        min_overlap_fraction: Minimum overlap as a fraction of ``tile_size``.
        blend: Importance-window mode, one of ``"gaussian"``/``"pyramid"``/``"constant"``.
        sigma_scale: Gaussian importance-window sigma scale.
        show_coverage: If ``True``, overlay the summed-window coverage heatmap.
        dpi: Dots per inch passed to :func:`plot_img`.
        scale: Figure size scale passed to :func:`plot_img`.

    Returns:
        A ``matplotlib.figure.Figure`` with the image, tile rectangles and
        (optionally) the coverage heatmap.
    """
    import matplotlib.patches as mpatches
    import torch

    from sleap_nn.data.tiling import generate_tile_grid
    from sleap_nn.inference.tile_merger import build_importance_window

    if hasattr(img, "detach"):
        img = img.detach().cpu().numpy()
    img = np.asarray(img)
    img = np.squeeze(img)
    if img.ndim == 2:
        img = img[..., None]
    elif img.ndim == 3 and img.shape[0] in (1, 3) and img.shape[-1] not in (1, 3):
        # (C, H, W) -> (H, W, C)
        img = np.transpose(img, (1, 2, 0))

    height, width = int(img.shape[0]), int(img.shape[1])
    origins = generate_tile_grid(
        (height, width),
        tile_size=tile_size,
        overlap=overlap,
        output_stride=output_stride,
        max_stride=max_stride,
        min_overlap_fraction=min_overlap_fraction,
    )

    fig = plot_img(img, dpi=dpi, scale=scale)
    ax = fig.gca()

    if show_coverage:
        stride = output_stride
        cnt = torch.zeros((max(1, height // stride), max(1, width // stride)))
        win = build_importance_window(
            (tile_size // stride, tile_size // stride),
            mode=blend,
            sigma_scale=sigma_scale,
        )
        th, tw = win.shape
        for y0, x0 in origins:
            yy, xx = y0 // stride, x0 // stride
            ye, xe = min(yy + th, cnt.shape[0]), min(xx + tw, cnt.shape[1])
            if ye > yy and xe > xx:
                cnt[yy:ye, xx:xe] += win[: ye - yy, : xe - xx]
        ax.imshow(
            cnt.numpy(),
            alpha=0.5,
            cmap="viridis",
            origin="upper",
            extent=[-0.5, width - 0.5, height - 0.5, -0.5],
        )

    for y0, x0 in origins:
        ax.add_patch(
            mpatches.Rectangle(
                (x0 - 0.5, y0 - 0.5),
                tile_size,
                tile_size,
                fill=False,
                edgecolor="red",
                linewidth=1.0,
                alpha=0.7,
            )
        )

    return fig

xavier_init_weights(x)

Function to initilaise the model weights with Xavier initialization method.

Source code in sleap_nn/training/utils.py

def xavier_init_weights(x):
    """Function to initilaise the model weights with Xavier initialization method."""
    if isinstance(x, nn.Conv2d) or isinstance(x, nn.Linear):
        if x.weight is not None:
            nn.init.xavier_uniform_(x.weight)
        if x.bias is not None:
            nn.init.constant_(x.bias, 0)