utils

sleap_nn.export.utils

Utilities for export workflows.

Functions:

Name	Description
build_bottomup_candidate_template	Build candidate template matching ONNX wrapper's line_scores ordering.
load_training_config	Load training configuration from a model directory.
resolve_anchor_part	Resolve anchor_part from config for the model types that center a crop.
resolve_backbone_source	Resolve a human-readable backbone weight source for metadata.
resolve_backbone_type	Return backbone type from config.
resolve_background_fill	Resolve the masked-out background fill the trained embedder used (burn-in).
resolve_burn_in	Resolve whether the trained embedder masked the crop (mask burn-in).
resolve_centroid_method	Resolve the trained centroid method for the model types that center a crop.
resolve_class_maps_output_stride	Resolve class maps output stride for multiclass bottom-up models.
resolve_class_names	Resolve class names for multiclass models.
resolve_crop_size	Resolve crop size from preprocessing config.
resolve_edge_inds	Resolve edge indices for metadata.
resolve_embedding_dim	Resolve the embedding (output vector) dimensionality for an embedding model.
resolve_embedding_input_channels	Resolve the DATA channels an embedding crop is fed with (not the backbone).
resolve_input_channels	Resolve input channels from backbone config.
resolve_input_scale	Resolve preprocessing scale from config.
resolve_input_shape	Resolve a dummy input shape for export.
resolve_model_type	Return model type from config.
resolve_n_classes	Resolve number of classes for multiclass models.
resolve_node_names	Resolve node names for metadata.
resolve_normalize	Resolve whether the embedding head L2-normalizes its output.
resolve_output_stride	Resolve output stride from head config.
resolve_pafs_output_stride	Resolve PAFs output stride for bottom-up models.
warn_on_tiled_export	Warn that a tiled model is exported without its tiling wrapper (deferred).

build_bottomup_candidate_template(n_nodes, max_peaks_per_node, edge_inds)

Build candidate template matching ONNX wrapper's line_scores ordering.

The ONNX BottomUpONNXWrapper produces line_scores with shape (n_edges, k*k) where for each edge connecting (src_node, dst_node), position i*k + j corresponds to: - src peak flat index: src_node * k + i - dst peak flat index: dst_node * k + j

This function builds edge_inds and edge_peak_inds tensors that match this exact ordering, so that line_scores_flat[idx] corresponds to edge_peak_inds[idx].

Parameters:

Name	Type	Description	Default
n_nodes	int	Number of nodes in the skeleton.	required
max_peaks_per_node	int	Maximum peaks per node (k) used during export.	required
edge_inds	List[Tuple[int, int]]	List of (src_node, dst_node) tuples defining skeleton edges.	required

Returns:

Type	Description
Tuple['torch.Tensor', 'torch.Tensor', 'torch.Tensor']	Tuple of (peak_channel_inds, edge_inds_tensor, edge_peak_inds_tensor): - peak_channel_inds: (n_nodes * k,) tensor mapping flat peak index to node - edge_inds_tensor: (n_edges * k * k,) tensor of edge indices for each candidate - edge_peak_inds_tensor: (n_edges * k * k, 2) tensor of (src, dst) peak indices

Example

from sleap_nn.export.utils import build_bottomup_candidate_template peak_ch, edge_inds, edge_peaks = build_bottomup_candidate_template( ... n_nodes=15, max_peaks_per_node=20, edge_inds=[(1, 2), (1, 5)] ... )

Use with ONNX output:

line_scores_flat = line_scores.reshape(-1) valid_scores = line_scores_flat[valid_mask] valid_edge_peaks = edge_peaks[valid_mask]

Note

This function is necessary because get_connection_candidates() in sleap_nn.inference.paf_grouping uses unstable argsort, which shuffles peak indices within each node and breaks alignment with ONNX output ordering.

Source code in sleap_nn/export/utils.py

def build_bottomup_candidate_template(
    n_nodes: int, max_peaks_per_node: int, edge_inds: List[Tuple[int, int]]
) -> Tuple["torch.Tensor", "torch.Tensor", "torch.Tensor"]:
    """Build candidate template matching ONNX wrapper's line_scores ordering.

    The ONNX BottomUpONNXWrapper produces line_scores with shape (n_edges, k*k) where
    for each edge connecting (src_node, dst_node), position i*k + j corresponds to:
    - src peak flat index: src_node * k + i
    - dst peak flat index: dst_node * k + j

    This function builds edge_inds and edge_peak_inds tensors that match this exact
    ordering, so that line_scores_flat[idx] corresponds to edge_peak_inds[idx].

    Args:
        n_nodes: Number of nodes in the skeleton.
        max_peaks_per_node: Maximum peaks per node (k) used during export.
        edge_inds: List of (src_node, dst_node) tuples defining skeleton edges.

    Returns:
        Tuple of (peak_channel_inds, edge_inds_tensor, edge_peak_inds_tensor):
        - peak_channel_inds: (n_nodes * k,) tensor mapping flat peak index to node
        - edge_inds_tensor: (n_edges * k * k,) tensor of edge indices for each candidate
        - edge_peak_inds_tensor: (n_edges * k * k, 2) tensor of (src, dst) peak indices

    Example:
        >>> from sleap_nn.export.utils import build_bottomup_candidate_template
        >>> peak_ch, edge_inds, edge_peaks = build_bottomup_candidate_template(
        ...     n_nodes=15, max_peaks_per_node=20, edge_inds=[(1, 2), (1, 5)]
        ... )
        >>> # Use with ONNX output:
        >>> line_scores_flat = line_scores.reshape(-1)
        >>> valid_scores = line_scores_flat[valid_mask]
        >>> valid_edge_peaks = edge_peaks[valid_mask]

    Note:
        This function is necessary because `get_connection_candidates()` in
        `sleap_nn.inference.paf_grouping` uses unstable argsort, which shuffles
        peak indices within each node and breaks alignment with ONNX output ordering.
    """
    import torch

    k = max_peaks_per_node
    n_edges = len(edge_inds)

    # peak_channel_inds: [0,0,...0, 1,1,...1, ...] (k times each)
    peak_channel_inds = torch.arange(n_nodes, dtype=torch.int32).repeat_interleave(k)

    edge_inds_list = []
    edge_peak_inds_list = []

    for edge_idx, (src_node, dst_node) in enumerate(edge_inds):
        # Build k*k candidate pairs in row-major order (i*k + j)
        # src indices: [src_node*k + 0, src_node*k + 0, ..., src_node*k + 1, ...]
        # dst indices: [dst_node*k + 0, dst_node*k + 1, ..., dst_node*k + 0, ...]
        src_base = src_node * k
        dst_base = dst_node * k

        src_indices = torch.arange(k, dtype=torch.int32).repeat_interleave(k) + src_base
        dst_indices = torch.arange(k, dtype=torch.int32).repeat(k) + dst_base

        edge_inds_list.append(torch.full((k * k,), edge_idx, dtype=torch.int32))
        edge_peak_inds_list.append(torch.stack([src_indices, dst_indices], dim=1))

    if edge_inds_list:
        edge_inds_tensor = torch.cat(edge_inds_list)
        edge_peak_inds_tensor = torch.cat(edge_peak_inds_list)
    else:
        edge_inds_tensor = torch.empty((0,), dtype=torch.int32)
        edge_peak_inds_tensor = torch.empty((0, 2), dtype=torch.int32)

    return peak_channel_inds, edge_inds_tensor, edge_peak_inds_tensor

load_training_config(model_dir)

Load training configuration from a model directory.

Source code in sleap_nn/export/utils.py

def load_training_config(model_dir: str | Path) -> DictConfig:
    """Load training configuration from a model directory."""
    model_dir = Path(model_dir)
    yaml_path = model_dir / "training_config.yaml"
    json_path = model_dir / "training_config.json"

    if yaml_path.exists():
        return OmegaConf.load(yaml_path.as_posix())
    if json_path.exists():
        return TrainingJobConfig.load_sleap_config(json_path.as_posix())

    raise FileNotFoundError(
        f"No training_config.yaml or training_config.json found in {model_dir}"
    )

resolve_anchor_part(cfg, model_type)

Resolve anchor_part from config for the model types that center a crop.

Parameters:

Name	Type	Description	Default
cfg	DictConfig	The training job configuration.	required
model_type	str	The model type (e.g., "centroid", "centered_instance", "embedding").	required

Returns:

Type	Description
Optional[str]	The anchor part name if configured, None otherwise. Only returns a value for the "centroid", "centered_instance" and "embedding" model types -- the heads that carry the crop/centroid-center knobs (#586). For an embedding model the consumer of the export has to produce the crops itself, so this is the one piece of geometry it cannot do without.

Source code in sleap_nn/export/utils.py

def resolve_anchor_part(cfg: DictConfig, model_type: str) -> Optional[str]:
    """Resolve anchor_part from config for the model types that center a crop.

    Args:
        cfg: The training job configuration.
        model_type: The model type (e.g., "centroid", "centered_instance",
            "embedding").

    Returns:
        The anchor part name if configured, None otherwise. Only returns a value
        for the "centroid", "centered_instance" and "embedding" model types -- the
        heads that carry the crop/centroid-center knobs (#586). For an
        ``embedding`` model the consumer of the export has to produce the crops
        itself, so this is the one piece of geometry it cannot do without.
    """
    leaf = _centroid_head_leaf(cfg, model_type)
    if leaf is _NO_HEAD:
        return None
    return getattr(leaf, "anchor_part", None)

resolve_backbone_source(cfg)

Resolve a human-readable backbone weight source for metadata.

Returns "imagenet" if the backbone was initialized from pretrained ImageNet weights (convnext/swint pre_trained_weights or unet pretrained_backbone_weights), else "scratch".

Source code in sleap_nn/export/utils.py

def resolve_backbone_source(cfg: DictConfig) -> Optional[str]:
    """Resolve a human-readable backbone weight source for metadata.

    Returns ``"imagenet"`` if the backbone was initialized from pretrained ImageNet
    weights (convnext/swint ``pre_trained_weights`` or unet
    ``pretrained_backbone_weights``), else ``"scratch"``.
    """
    backbone_type = get_backbone_type_from_cfg(cfg)
    backbone_cfg = cfg.model_config.backbone_config.get(backbone_type)
    if backbone_cfg is not None:
        pre_trained = OmegaConf.select(
            backbone_cfg, "pre_trained_weights", default=None
        )
        if pre_trained:
            return "imagenet"
    if OmegaConf.select(cfg, "model_config.pretrained_backbone_weights", default=None):
        return "imagenet"
    return "scratch"

resolve_backbone_type(cfg)

Return backbone type from config.

Source code in sleap_nn/export/utils.py

def resolve_backbone_type(cfg: DictConfig) -> str:
    """Return backbone type from config."""
    return get_backbone_type_from_cfg(cfg)

resolve_background_fill(cfg)

Resolve the masked-out background fill the trained embedder used (burn-in).

Source code in sleap_nn/export/utils.py

def resolve_background_fill(cfg: DictConfig) -> str:
    """Resolve the masked-out background fill the trained embedder used (burn-in)."""
    return str(
        OmegaConf.select(
            cfg, "data_config.preprocessing.background_fill", default="black"
        )
    )

resolve_burn_in(cfg)

Resolve whether the trained embedder masked the crop (mask burn-in).

A burn_in=True model standardizes over the foreground only and replaces the background, which the single-input ONNX wrapper (maskless whole-crop standardize) cannot reproduce — so this drives the export-time divergence warning + metadata. Mirrors the canonical default in PreprocessingConfig.burn_in (False).

Source code in sleap_nn/export/utils.py

def resolve_burn_in(cfg: DictConfig) -> bool:
    """Resolve whether the trained embedder masked the crop (mask burn-in).

    A ``burn_in=True`` model standardizes over the foreground only and replaces the
    background, which the single-input ONNX wrapper (maskless whole-crop standardize)
    cannot reproduce — so this drives the export-time divergence warning + metadata.
    Mirrors the canonical default in ``PreprocessingConfig.burn_in`` (``False``).
    """
    return bool(
        OmegaConf.select(cfg, "data_config.preprocessing.burn_in", default=False)
    )

resolve_centroid_method(cfg, model_type)

Resolve the trained centroid method for the model types that center a crop.

Companion to :func:resolve_anchor_part. Recorded in the export metadata so a consumer of the exported model can tag predicted centroids with the method the model was actually trained on (#586) — without it, a bbox_center model's predictions would be recorded as center_of_mass.

The same applies to an embedding (re-ID) model, more strongly: its consumer must produce the crops itself, so the crop-center recipe the embedder was trained with is what makes its vectors comparable.

Parameters:

Name	Type	Description	Default
cfg	DictConfig	The training job configuration.	required
model_type	str	The model type (e.g., "centroid", "centered_instance", "embedding").	required

Returns:

Type	Description
Optional[str]	The resolved method (one of sleap_nn.data.instance_centroids.CENTROID_METHODS), or None for model types that have no centroid.

Source code in sleap_nn/export/utils.py

def resolve_centroid_method(cfg: DictConfig, model_type: str) -> Optional[str]:
    """Resolve the trained centroid method for the model types that center a crop.

    Companion to :func:`resolve_anchor_part`. Recorded in the export metadata so a
    consumer of the exported model can tag predicted centroids with the method the
    model was actually trained on (#586) — without it, a ``bbox_center`` model's
    predictions would be recorded as ``center_of_mass``.

    The same applies to an ``embedding`` (re-ID) model, more strongly: its
    consumer must produce the crops itself, so the crop-center recipe the
    embedder was trained with is what makes its vectors comparable.

    Args:
        cfg: The training job configuration.
        model_type: The model type (e.g., "centroid", "centered_instance",
            "embedding").

    Returns:
        The resolved method (one of
        ``sleap_nn.data.instance_centroids.CENTROID_METHODS``), or ``None`` for
        model types that have no centroid.
    """
    from sleap_nn.data.instance_centroids import centroid_method_from_config

    leaf = _centroid_head_leaf(cfg, model_type)
    if leaf is _NO_HEAD:
        return None
    return centroid_method_from_config(leaf)[0]

resolve_class_maps_output_stride(cfg)

Resolve class maps output stride for multiclass bottom-up models.

Source code in sleap_nn/export/utils.py

def resolve_class_maps_output_stride(cfg: DictConfig) -> int:
    """Resolve class maps output stride for multiclass bottom-up models."""
    mc_bottomup_cfg = getattr(
        cfg.model_config.head_configs, "multi_class_bottomup", None
    )
    if mc_bottomup_cfg is not None and mc_bottomup_cfg.class_maps is not None:
        return int(mc_bottomup_cfg.class_maps.output_stride)
    return 8

resolve_class_names(cfg, model_type)

Resolve class names for multiclass models.

Source code in sleap_nn/export/utils.py

def resolve_class_names(cfg: DictConfig, model_type: str) -> List[str]:
    """Resolve class names for multiclass models."""
    head_cfg = cfg.model_config.head_configs.get(model_type)
    if head_cfg is None:
        return []

    # Top-down multiclass: class_vectors.classes
    if hasattr(head_cfg, "class_vectors") and head_cfg.class_vectors is not None:
        classes = getattr(head_cfg.class_vectors, "classes", None)
        if classes:
            return list(classes)

    # Bottom-up multiclass: class_maps.classes
    if hasattr(head_cfg, "class_maps") and head_cfg.class_maps is not None:
        classes = getattr(head_cfg.class_maps, "classes", None)
        if classes:
            return list(classes)

    return []

resolve_crop_size(cfg)

Resolve crop size from preprocessing config.

Source code in sleap_nn/export/utils.py

def resolve_crop_size(cfg: DictConfig) -> Optional[Tuple[int, int]]:
    """Resolve crop size from preprocessing config."""
    crop_size = cfg.data_config.preprocessing.crop_size
    if crop_size is None:
        return None
    # Check for list/tuple or OmegaConf ListConfig
    if isinstance(crop_size, (list, tuple)) or (
        hasattr(crop_size, "__iter__")
        and hasattr(crop_size, "__len__")
        and not isinstance(crop_size, (str, int))
    ):
        if len(crop_size) == 2:
            return int(crop_size[0]), int(crop_size[1])
        if len(crop_size) == 1:
            return int(crop_size[0]), int(crop_size[0])
    return int(crop_size), int(crop_size)

resolve_edge_inds(cfg, node_names)

Resolve edge indices for metadata.

Source code in sleap_nn/export/utils.py

def resolve_edge_inds(cfg: DictConfig, node_names: List[str]) -> List[Tuple[int, int]]:
    """Resolve edge indices for metadata."""
    edges = _edge_inds_from_skeletons(cfg.data_config.skeletons)
    if edges:
        return _normalize_edges(edges, node_names)

    bottomup_cfg = getattr(cfg.model_config.head_configs, "bottomup", None)
    if bottomup_cfg is not None and bottomup_cfg.pafs is not None:
        edges = bottomup_cfg.pafs.edges
        if edges:
            return _normalize_edges(edges, node_names)

    return []

resolve_embedding_dim(cfg)

Resolve the embedding (output vector) dimensionality for an embedding model.

Source code in sleap_nn/export/utils.py

def resolve_embedding_dim(cfg: DictConfig) -> int:
    """Resolve the embedding (output vector) dimensionality for an embedding model."""
    leaf = OmegaConf.select(
        cfg, "model_config.head_configs.embedding.embedding", default=None
    )
    if leaf is None:
        return 128
    return int(OmegaConf.select(leaf, "embedding_dim", default=128))

resolve_embedding_input_channels(cfg)

Resolve the DATA channels an embedding crop is fed with (not the backbone).

The embedder forces grayscale by default (1 channel); a 3-channel ImageNet backbone repeats the gray channel internally (Model.forward). Only when the user explicitly opts into RGB (ensure_rgb) are the crops 3-channel. This is the channel count the exported graph should accept, independent of the backbone in_channels.

Source code in sleap_nn/export/utils.py

def resolve_embedding_input_channels(cfg: DictConfig) -> int:
    """Resolve the DATA channels an embedding crop is fed with (not the backbone).

    The embedder forces grayscale by default (1 channel); a 3-channel ImageNet
    backbone repeats the gray channel internally (``Model.forward``). Only when the
    user explicitly opts into RGB (``ensure_rgb``) are the crops 3-channel. This is
    the channel count the exported graph should accept, independent of the backbone
    ``in_channels``.
    """
    preprocessing = cfg.data_config.preprocessing
    if bool(getattr(preprocessing, "ensure_rgb", False)):
        return 3
    if bool(getattr(preprocessing, "ensure_grayscale", False)):
        return 1
    # Embedding default is grayscale even when neither flag is set.
    return 1

resolve_input_channels(cfg)

Resolve input channels from backbone config.

Source code in sleap_nn/export/utils.py

def resolve_input_channels(cfg: DictConfig) -> int:
    """Resolve input channels from backbone config."""
    backbone_type = get_backbone_type_from_cfg(cfg)
    return int(cfg.model_config.backbone_config[backbone_type].in_channels)

resolve_input_scale(cfg)

Resolve preprocessing scale from config.

Source code in sleap_nn/export/utils.py

def resolve_input_scale(cfg: DictConfig) -> float:
    """Resolve preprocessing scale from config."""
    scale = cfg.data_config.preprocessing.scale
    # Check for list/tuple or OmegaConf ListConfig
    if isinstance(scale, (list, tuple)) or (
        hasattr(scale, "__iter__")
        and hasattr(scale, "__len__")
        and not isinstance(scale, str)
    ):
        return float(scale[0]) if len(scale) > 0 else 1.0
    return float(scale)

resolve_input_shape(cfg, input_height=None, input_width=None)

Resolve a dummy input shape for export.

Source code in sleap_nn/export/utils.py

def resolve_input_shape(
    cfg: DictConfig,
    input_height: Optional[int] = None,
    input_width: Optional[int] = None,
) -> Tuple[int, int, int, int]:
    """Resolve a dummy input shape for export."""
    channels = resolve_input_channels(cfg)
    height = input_height or cfg.data_config.preprocessing.max_height or 512
    width = input_width or cfg.data_config.preprocessing.max_width or 512
    return 1, channels, int(height), int(width)

resolve_model_type(cfg)

Return model type from config.

Source code in sleap_nn/export/utils.py

def resolve_model_type(cfg: DictConfig) -> str:
    """Return model type from config."""
    return get_model_type_from_cfg(cfg)

resolve_n_classes(cfg, model_type)

Resolve number of classes for multiclass models.

Source code in sleap_nn/export/utils.py

def resolve_n_classes(cfg: DictConfig, model_type: str) -> int:
    """Resolve number of classes for multiclass models."""
    class_names = resolve_class_names(cfg, model_type)
    return len(class_names) if class_names else 0

resolve_node_names(cfg, model_type)

Resolve node names for metadata.

Source code in sleap_nn/export/utils.py

def resolve_node_names(cfg: DictConfig, model_type: str) -> List[str]:
    """Resolve node names for metadata."""
    skeleton_nodes = _node_names_from_skeletons(cfg.data_config.skeletons)
    if skeleton_nodes:
        return skeleton_nodes

    head_cfg = cfg.model_config.head_configs.get(model_type)
    if head_cfg is None:
        return []

    if hasattr(head_cfg, "confmaps") and head_cfg.confmaps is not None:
        part_names = getattr(head_cfg.confmaps, "part_names", None)
        if part_names:
            return list(part_names)

    if model_type == "centroid":
        anchor = getattr(head_cfg.confmaps, "anchor_part", None) if head_cfg else None
        return [anchor] if anchor else ["centroid"]

    return []

resolve_normalize(cfg)

Resolve whether the embedding head L2-normalizes its output.

Source code in sleap_nn/export/utils.py

def resolve_normalize(cfg: DictConfig) -> bool:
    """Resolve whether the embedding head L2-normalizes its output."""
    leaf = OmegaConf.select(
        cfg, "model_config.head_configs.embedding.embedding", default=None
    )
    if leaf is None:
        return True
    return bool(OmegaConf.select(leaf, "normalize", default=True))

resolve_output_stride(cfg, model_type)

Resolve output stride from head config.

Source code in sleap_nn/export/utils.py

def resolve_output_stride(cfg: DictConfig, model_type: str) -> int:
    """Resolve output stride from head config."""
    head_cfg = cfg.model_config.head_configs[model_type]
    if head_cfg is None:
        return 1
    if model_type == "embedding":
        # Pooled re-ID head: no confmaps/pafs. The head taps the backbone's
        # ``middle_output`` so its output_stride equals the backbone max_stride.
        leaf = getattr(head_cfg, "embedding", None)
        if leaf is not None and getattr(leaf, "output_stride", None) is not None:
            return int(leaf.output_stride)
        return 1
    if hasattr(head_cfg, "confmaps") and head_cfg.confmaps is not None:
        return int(head_cfg.confmaps.output_stride)
    if hasattr(head_cfg, "pafs") and head_cfg.pafs is not None:
        return int(head_cfg.pafs.output_stride)
    return 1

resolve_pafs_output_stride(cfg)

Resolve PAFs output stride for bottom-up models.

Source code in sleap_nn/export/utils.py

def resolve_pafs_output_stride(cfg: DictConfig) -> int:
    """Resolve PAFs output stride for bottom-up models."""
    bottomup_cfg = getattr(cfg.model_config.head_configs, "bottomup", None)
    if bottomup_cfg is not None and bottomup_cfg.pafs is not None:
        return int(bottomup_cfg.pafs.output_stride)
    return 1

warn_on_tiled_export(cfg)

Warn that a tiled model is exported without its tiling wrapper (deferred).

Tiled ONNX/TensorRT export is not yet implemented (design DQ15): the exported graph is the plain per-frame network, so an exported tiled model runs whole-frame (no sliding-window tiling / stitching). PyTorch inference still tiles correctly.

Parameters:

Name	Type	Description	Default
cfg	DictConfig	A loaded training config.	required

Source code in sleap_nn/export/utils.py

def warn_on_tiled_export(cfg: DictConfig) -> None:
    """Warn that a tiled model is exported without its tiling wrapper (deferred).

    Tiled ONNX/TensorRT export is not yet implemented (design DQ15): the exported
    graph is the plain per-frame network, so an exported tiled model runs whole-frame
    (no sliding-window tiling / stitching). PyTorch inference still tiles correctly.

    Args:
        cfg: A loaded training config.
    """
    tiling = OmegaConf.select(cfg, "data_config.preprocessing.tiling")
    if tiling is not None and getattr(tiling, "enabled", False):
        logger.warning(
            "This model was trained with tiling enabled, but tiled ONNX/TensorRT "
            "export is not yet supported. The exported model runs on WHOLE frames "
            "without tiling/stitching and will not reproduce tiled-inference results. "
            "Use PyTorch inference for tiled prediction."
        )