Skip to content

instance_cropping

sleap_nn.data.instance_cropping

Handle cropping of instances.

Classes:

Name Description
BBoxCropSpec

How a box-scaled centered-instance crop is cut (crop_mode="bbox").

Functions:

Name Description
bbox_crop_window_sides

Unfloored box-scaled crop window sides of every training instance.

bbox_crop_windows

Crop windows that resize each box to a crop_size x crop_size crop.

clip_rects

Clip [y0, y1, x0, x1) rectangles to a (height, width) frame.

compute_augmentation_padding

Compute padding needed to accommodate augmentation transforms.

count_clipped_instances

Count labeled instances that a given crop size would clip.

crop_windows

Resample crop_size x crop_size crops of windows (inference path).

find_instance_crop_size

Compute a crop size that contains every labeled instance.

find_mask_crop_size

Compute a crop size that contains every mask (the mask analog of find_instance_crop_size).

find_max_instance_bbox_size

Find the maximum bounding box dimension across all instances in labels.

generate_crops

Generate cropped image for the given centroid.

iter_mask_extents

Yield (required_crop_size, bbox_size) for every mask, in size-matched pixels.

iter_required_crop_sizes

Yield the crop size each labeled instance needs to avoid being clipped.

letterbox_crops

Zero the pixels of (N, C, h, w) crops outside the centered padded box.

make_centered_bboxes

Create centered bounding boxes around centroid.

mask_foreground_rects

A mask's foreground on the image-pixel grid, as disjoint rectangles.

mask_rects_center

The (cx, cy) a mask crop centers on, from the mask's rectangles.

mask_rects_moments

(n, sum_x, sum_y): a mask's pixel count and coordinate sums, exactly.

points_from_crop

Map crop keypoints (N, ..., 2) back to the frame.

points_to_crop

Map frame keypoints (N, ..., 2) into the crops of windows.

pre_crop_size

Side of the training pre-crop that leaves room for rotation augmentation.

rasterize_rects

The [y0:y1, x0:x1] window of the mask rects cover, as a bool array.

scale_windows

Grow (or shrink) (N, 4) edge-coordinate windows about their centers.

window_corners

(N, 4) edge-coordinate windows -> (N, 4, 2) TL, TR, BR, BL corners.

BBoxCropSpec dataclass

How a box-scaled centered-instance crop is cut (crop_mode="bbox").

One object describes the crop for training (the datasets) and inference (the top-down layer), built from the centered-instance model's config by :meth:from_config.

Attributes:

Name Type Description
aspect str

preprocessing.crop_bbox_aspect (:data:BBOX_CROP_ASPECTS).

padding float

preprocessing.crop_bbox_padding (fraction of the longest side).

min_size Optional[float]

preprocessing.crop_bbox_min_size (frame pixels) or None.

bbox_source str

data_config.bbox_source -- where training boxes come from.

bbox_padding float

data_config.bbox_padding -- pixels added to derived boxes.

Methods:

Name Description
__post_init__

Validate the aspect mode.

from_config

The spec of a training config, or None for crop_mode="fixed".

windows

:func:bbox_crop_windows with this spec's aspect, padding and floor.

Source code in sleap_nn/data/instance_cropping.py
@dataclass(frozen=True)
class BBoxCropSpec:
    """How a box-scaled centered-instance crop is cut (``crop_mode="bbox"``).

    One object describes the crop for training (the datasets) and inference (the
    top-down layer), built from the centered-instance model's config by
    :meth:`from_config`.

    Attributes:
        aspect: ``preprocessing.crop_bbox_aspect`` (:data:`BBOX_CROP_ASPECTS`).
        padding: ``preprocessing.crop_bbox_padding`` (fraction of the longest side).
        min_size: ``preprocessing.crop_bbox_min_size`` (frame pixels) or ``None``.
        bbox_source: ``data_config.bbox_source`` -- where training boxes come from.
        bbox_padding: ``data_config.bbox_padding`` -- pixels added to derived boxes.
    """

    aspect: str = "square"
    padding: float = 0.1
    min_size: Optional[float] = None
    bbox_source: str = "auto"
    bbox_padding: float = 0.0

    def __post_init__(self):
        """Validate the aspect mode."""
        if self.aspect not in BBOX_CROP_ASPECTS:
            raise ValueError(
                f"crop_bbox_aspect must be one of {BBOX_CROP_ASPECTS}, got "
                f"{self.aspect!r}."
            )

    @classmethod
    def from_config(cls, config) -> Optional["BBoxCropSpec"]:
        """The spec of a training config, or ``None`` for ``crop_mode="fixed"``."""
        from omegaconf import OmegaConf

        def get(key, default):
            value = OmegaConf.select(config, key, default=default)
            return default if value is None else value

        mode = get("data_config.preprocessing.crop_mode", "fixed")
        if mode not in CROP_MODES:
            raise ValueError(f"crop_mode must be one of {CROP_MODES}, got {mode!r}.")
        if mode == "fixed":
            return None
        min_size = OmegaConf.select(
            config, "data_config.preprocessing.crop_bbox_min_size", default=None
        )
        return cls(
            aspect=str(get("data_config.preprocessing.crop_bbox_aspect", "square")),
            padding=float(get("data_config.preprocessing.crop_bbox_padding", 0.1)),
            min_size=None if min_size is None else float(min_size),
            bbox_source=str(get("data_config.bbox_source", "auto")),
            bbox_padding=float(get("data_config.bbox_padding", 0.0)),
        )

    def windows(
        self, boxes: torch.Tensor, crop_size: int
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """:func:`bbox_crop_windows` with this spec's aspect, padding and floor."""
        return bbox_crop_windows(
            boxes, crop_size, self.aspect, self.padding, self.min_size
        )

__post_init__()

Validate the aspect mode.

Source code in sleap_nn/data/instance_cropping.py
def __post_init__(self):
    """Validate the aspect mode."""
    if self.aspect not in BBOX_CROP_ASPECTS:
        raise ValueError(
            f"crop_bbox_aspect must be one of {BBOX_CROP_ASPECTS}, got "
            f"{self.aspect!r}."
        )

from_config(config) classmethod

The spec of a training config, or None for crop_mode="fixed".

Source code in sleap_nn/data/instance_cropping.py
@classmethod
def from_config(cls, config) -> Optional["BBoxCropSpec"]:
    """The spec of a training config, or ``None`` for ``crop_mode="fixed"``."""
    from omegaconf import OmegaConf

    def get(key, default):
        value = OmegaConf.select(config, key, default=default)
        return default if value is None else value

    mode = get("data_config.preprocessing.crop_mode", "fixed")
    if mode not in CROP_MODES:
        raise ValueError(f"crop_mode must be one of {CROP_MODES}, got {mode!r}.")
    if mode == "fixed":
        return None
    min_size = OmegaConf.select(
        config, "data_config.preprocessing.crop_bbox_min_size", default=None
    )
    return cls(
        aspect=str(get("data_config.preprocessing.crop_bbox_aspect", "square")),
        padding=float(get("data_config.preprocessing.crop_bbox_padding", 0.1)),
        min_size=None if min_size is None else float(min_size),
        bbox_source=str(get("data_config.bbox_source", "auto")),
        bbox_padding=float(get("data_config.bbox_padding", 0.0)),
    )

windows(boxes, crop_size)

:func:bbox_crop_windows with this spec's aspect, padding and floor.

Source code in sleap_nn/data/instance_cropping.py
def windows(
    self, boxes: torch.Tensor, crop_size: int
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """:func:`bbox_crop_windows` with this spec's aspect, padding and floor."""
    return bbox_crop_windows(
        boxes, crop_size, self.aspect, self.padding, self.min_size
    )

bbox_crop_window_sides(labels, spec, max_hw, user_instances_only=True)

Unfloored box-scaled crop window sides of every training instance.

The window side is the padded box's longest side (square/letterbox) or its larger padded extent (stretch), in size-matched pixels, before min_size applies -- what :meth:ModelTrainer sizes crop_size and crop_bbox_min_size from.

Parameters:

Name Type Description Default
labels Labels

Training labels.

required
spec BBoxCropSpec

The crop spec (its min_size is ignored).

required
max_hw Optional[Tuple[Optional[int], Optional[int]]]

The size matcher's (max_height, max_width).

required
user_instances_only bool

Measure user instances only, as the dataset trains.

True

Returns:

Type Description
ndarray

(n_instances,) window sides in pixels.

Source code in sleap_nn/data/instance_cropping.py
def bbox_crop_window_sides(
    labels: sio.Labels,
    spec: "BBoxCropSpec",
    max_hw: Optional[Tuple[Optional[int], Optional[int]]],
    user_instances_only: bool = True,
) -> np.ndarray:
    """Unfloored box-scaled crop window sides of every training instance.

    The window side is the padded box's longest side (``square``/``letterbox``) or
    its larger padded extent (``stretch``), in size-matched pixels, before
    ``min_size`` applies -- what :meth:`ModelTrainer` sizes ``crop_size`` and
    ``crop_bbox_min_size`` from.

    Args:
        labels: Training labels.
        spec: The crop spec (its ``min_size`` is ignored).
        max_hw: The size matcher's ``(max_height, max_width)``.
        user_instances_only: Measure user instances only, as the dataset trains.

    Returns:
        ``(n_instances,)`` window sides in pixels.
    """
    from sleap_nn.data.bboxes import instance_crop_bbox, resolve_bbox_source

    source = resolve_bbox_source(labels, spec.bbox_source)
    boxes = []
    for lf in labels:
        instances = lf.user_instances if user_instances_only else lf.instances
        eff = None
        for inst in instances:
            if inst.is_empty:
                continue
            box = instance_crop_bbox(lf, inst, source, spec.bbox_padding)
            if box is None:
                continue
            if eff is None:
                eff = _frame_eff_scale(lf, max_hw)
            boxes.append(box * eff)
    if not boxes:
        return np.zeros((0,), dtype=np.float64)
    windows, _, _ = bbox_crop_windows(
        torch.as_tensor(np.stack(boxes)), 1, spec.aspect, spec.padding, None
    )
    sides = (windows[:, 2:] - windows[:, :2]).max(dim=-1).values
    return sides.numpy().astype(np.float64)

bbox_crop_windows(boxes, crop_size, aspect='square', padding=0.1, min_size=None)

Crop windows that resize each box to a crop_size x crop_size crop.

Each box is padded on every side by padding times its longest side. For square and letterbox the window is the square on the padded box's longest side (one scale for both axes); for stretch it is the padded box itself (one scale per axis). Window sides are floored at min_size (and 1 px) so a tiny or degenerate box is not magnified without bound. Windows are centered on the box center.

Parameters:

Name Type Description Default
boxes Tensor

(N, 4) boxes (x1, y1, x2, y2) in keypoint coordinates.

required
crop_size int

Output crop side in pixels.

required
aspect str

One of :data:BBOX_CROP_ASPECTS.

'square'
padding float

Margin per side as a fraction of the box's longest side.

0.1
min_size Optional[float]

Smallest window side in frame pixels, or None.

None

Returns:

Type Description
Tuple[Tensor, Tensor, Tensor]

(windows, scales, keep_half): windows (N, 4) as edge-coordinate (X1, Y1, X2, Y2); scales (N, 2) the (sx, sy) frame -> crop scale; keep_half (N, 2) the half-extent, in crop pixels, of the padded box around the crop center (the region :func:letterbox_crops keeps).

Source code in sleap_nn/data/instance_cropping.py
def bbox_crop_windows(
    boxes: torch.Tensor,
    crop_size: int,
    aspect: str = "square",
    padding: float = 0.1,
    min_size: Optional[float] = None,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Crop windows that resize each box to a ``crop_size`` x ``crop_size`` crop.

    Each box is padded on every side by ``padding`` times its longest side. For
    ``square`` and ``letterbox`` the window is the square on the padded box's longest
    side (one scale for both axes); for ``stretch`` it is the padded box itself (one
    scale per axis). Window sides are floored at ``min_size`` (and 1 px) so a tiny
    or degenerate box is not magnified without bound. Windows are centered on the
    box center.

    Args:
        boxes: ``(N, 4)`` boxes ``(x1, y1, x2, y2)`` in keypoint coordinates.
        crop_size: Output crop side in pixels.
        aspect: One of :data:`BBOX_CROP_ASPECTS`.
        padding: Margin per side as a fraction of the box's longest side.
        min_size: Smallest window side in frame pixels, or ``None``.

    Returns:
        ``(windows, scales, keep_half)``: ``windows (N, 4)`` as edge-coordinate
        ``(X1, Y1, X2, Y2)``; ``scales (N, 2)`` the ``(sx, sy)`` frame -> crop scale;
        ``keep_half (N, 2)`` the half-extent, in crop pixels, of the padded box around
        the crop center (the region :func:`letterbox_crops` keeps).
    """
    if aspect not in BBOX_CROP_ASPECTS:
        raise ValueError(
            f"crop_bbox_aspect must be one of {BBOX_CROP_ASPECTS}, got {aspect!r}."
        )
    boxes = torch.as_tensor(boxes, dtype=torch.float32).reshape(-1, 4)
    x1, y1, x2, y2 = boxes.unbind(-1)
    w = (x2 - x1).clamp(min=0)
    h = (y2 - y1).clamp(min=0)
    margin = 2.0 * float(padding) * torch.maximum(w, h)
    padded = torch.stack([w + margin, h + margin], dim=-1)  # (N, 2)
    if aspect == "stretch":
        sides = padded.clone()
    else:
        sides = padded.max(dim=-1, keepdim=True).values.expand(-1, 2).clone()
    floor = max(float(min_size or 0.0), 1.0)
    sides = sides.clamp(min=floor)
    center = torch.stack([(x1 + x2) / 2, (y1 + y2) / 2], dim=-1) + 0.5
    windows = torch.cat([center - sides / 2, center + sides / 2], dim=-1)
    scales = float(crop_size) / sides
    keep_half = torch.minimum(padded * scales, sides * scales) / 2
    return windows, scales, keep_half

clip_rects(rects, frame_hw)

Clip [y0, y1, x0, x1) rectangles to a (height, width) frame.

Source code in sleap_nn/data/instance_cropping.py
def clip_rects(rects: np.ndarray, frame_hw: Tuple[int, int]) -> np.ndarray:
    """Clip ``[y0, y1, x0, x1)`` rectangles to a ``(height, width)`` frame."""
    clipped = rects.copy()
    clipped[:, 0:2] = np.clip(clipped[:, 0:2], 0, int(frame_hw[0]))
    clipped[:, 2:4] = np.clip(clipped[:, 2:4], 0, int(frame_hw[1]))
    return clipped[(clipped[:, 1] > clipped[:, 0]) & (clipped[:, 3] > clipped[:, 2])]

compute_augmentation_padding(bbox_size, rotation_max=0.0, scale_max=1.0)

Compute padding needed to accommodate augmentation transforms.

When rotation and scaling augmentations are applied, the bounding box of an instance can expand beyond its original size. This function calculates the padding needed to ensure the full instance remains visible after augmentation.

Parameters:

Name Type Description Default
bbox_size float

The size of the instance bounding box (max of width/height).

required
rotation_max float

Maximum absolute rotation angle in degrees. For symmetric rotation ranges like [-180, 180], pass 180.

0.0
scale_max float

Maximum scaling factor. For scale range [0.9, 1.1], pass 1.1.

1.0

Returns:

Type Description
int

Padding in pixels to add around the bounding box (total, not per side).

Source code in sleap_nn/data/instance_cropping.py
def compute_augmentation_padding(
    bbox_size: float,
    rotation_max: float = 0.0,
    scale_max: float = 1.0,
) -> int:
    """Compute padding needed to accommodate augmentation transforms.

    When rotation and scaling augmentations are applied, the bounding box of an
    instance can expand beyond its original size. This function calculates the
    padding needed to ensure the full instance remains visible after augmentation.

    Args:
        bbox_size: The size of the instance bounding box (max of width/height).
        rotation_max: Maximum absolute rotation angle in degrees. For symmetric
            rotation ranges like [-180, 180], pass 180.
        scale_max: Maximum scaling factor. For scale range [0.9, 1.1], pass 1.1.

    Returns:
        Padding in pixels to add around the bounding box (total, not per side).
    """
    if rotation_max == 0.0 and scale_max <= 1.0:
        return 0

    # For a square bbox rotated by angle θ, the new bbox has side length:
    # L' = L * (|cos(θ)| + |sin(θ)|)
    # Maximum expansion occurs at 45°: L' = L * sqrt(2)
    # For arbitrary angle: we use the worst case within the rotation range
    rotation_rad = math.radians(min(abs(rotation_max), 90))
    rotation_factor = abs(math.cos(rotation_rad)) + abs(math.sin(rotation_rad))

    # For angles > 45°, the factor increases, max at 45° = sqrt(2)
    # But for angles approaching 90°, it goes back to 1
    # Worst case in any range including 45° is sqrt(2)
    if abs(rotation_max) >= 45:
        rotation_factor = math.sqrt(2)

    # Combined expansion factor
    expansion_factor = rotation_factor * max(scale_max, 1.0)

    # Total padding needed (both sides)
    expanded_size = bbox_size * expansion_factor
    padding = expanded_size - bbox_size

    return int(math.ceil(padding))

count_clipped_instances(labels, crop_size, max_hw=None, anchor_ind=None, centroid_method=None, centroid_fallback=None, user_instances_only=False, superseded_only=False)

Count labeled instances that a given crop size would clip.

Used to warn about an explicitly configured crop_size, which we must not silently override -- some users knowingly accept clipping an extremity.

Parameters:

Name Type Description Default
labels Labels

A sio.Labels containing user-labeled instances.

required
crop_size int

The configured crop size, in size-matched pixels.

required
max_hw Optional[Tuple[Optional[int], Optional[int]]]

The configured (max_height, max_width).

None
anchor_ind Optional[int]

Index of the anchor node, or None.

None
centroid_method Optional[str]

The resolved centroid method, or None to infer.

None
centroid_fallback Optional[str]

The reduce method for a non-visible anchor node.

None
user_instances_only bool

When True, ignore predicted instances.

False
superseded_only bool

See :func:_iter_frame_instances.

False

Returns:

Type Description
Tuple[int, int, float]

(n_clipped, n_total, max_required) -- how many instances have at least one node outside the crop, how many were examined, and the crop size that would contain all of them (0.0 if there are none).

Source code in sleap_nn/data/instance_cropping.py
def count_clipped_instances(
    labels: sio.Labels,
    crop_size: int,
    max_hw: Optional[Tuple[Optional[int], Optional[int]]] = None,
    anchor_ind: Optional[int] = None,
    centroid_method: Optional[str] = None,
    centroid_fallback: Optional[str] = None,
    user_instances_only: bool = False,
    superseded_only: bool = False,
) -> Tuple[int, int, float]:
    """Count labeled instances that a given crop size would clip.

    Used to warn about an explicitly configured ``crop_size``, which we must not
    silently override -- some users knowingly accept clipping an extremity.

    Args:
        labels: A `sio.Labels` containing user-labeled instances.
        crop_size: The configured crop size, in size-matched pixels.
        max_hw: The configured ``(max_height, max_width)``.
        anchor_ind: Index of the anchor node, or ``None``.
        centroid_method: The resolved centroid method, or ``None`` to infer.
        centroid_fallback: The reduce method for a non-visible anchor node.
        user_instances_only: When ``True``, ignore predicted instances.
        superseded_only: See :func:`_iter_frame_instances`.

    Returns:
        ``(n_clipped, n_total, max_required)`` -- how many instances have at
        least one node outside the crop, how many were examined, and the crop
        size that would contain all of them (``0.0`` if there are none).
    """
    n_clipped = 0
    n_total = 0
    max_required = 0.0
    for required in iter_required_crop_sizes(
        labels,
        max_hw=max_hw,
        anchor_ind=anchor_ind,
        centroid_method=centroid_method,
        centroid_fallback=centroid_fallback,
        user_instances_only=user_instances_only,
        superseded_only=superseded_only,
    ):
        n_total += 1
        if required > crop_size:
            n_clipped += 1
        max_required = max(max_required, required)
    return n_clipped, n_total, max_required

crop_windows(images, windows, sample_inds, crop_size)

Resample crop_size x crop_size crops of windows (inference path).

Reproduces the training crops of crop_and_resize_skia for the same window: one bilinear sample per output pixel at its center, clamped to the frame's edge pixels, then weighted -- like Skia's antialiased draw -- by how much of the output pixel's footprint lies inside the frame, so a window that overhangs the frame fades to zero there.

Parameters:

Name Type Description Default
images Tensor

(B, C, H, W) frames.

required
windows Tensor

(N, 4) edge-coordinate windows (:func:bbox_crop_windows).

required
sample_inds Tensor

(N,) frame index of each window.

required
crop_size int

Output crop side.

required

Returns:

Type Description
Tensor

(N, C, crop_size, crop_size) float32 crops.

Source code in sleap_nn/data/instance_cropping.py
def crop_windows(
    images: torch.Tensor,
    windows: torch.Tensor,
    sample_inds: torch.Tensor,
    crop_size: int,
) -> torch.Tensor:
    """Resample ``crop_size`` x ``crop_size`` crops of ``windows`` (inference path).

    Reproduces the training crops of ``crop_and_resize_skia`` for the same window:
    one bilinear sample per output pixel at its center, clamped to the frame's edge
    pixels, then weighted -- like Skia's antialiased draw -- by how much of the
    output pixel's footprint lies inside the frame, so a window that overhangs the
    frame fades to zero there.

    Args:
        images: ``(B, C, H, W)`` frames.
        windows: ``(N, 4)`` edge-coordinate windows (:func:`bbox_crop_windows`).
        sample_inds: ``(N,)`` frame index of each window.
        crop_size: Output crop side.

    Returns:
        ``(N, C, crop_size, crop_size)`` float32 crops.
    """
    import torch.nn.functional as F

    n, c = windows.shape[0], images.shape[1]
    h, w = images.shape[-2:]
    out = torch.zeros(n, c, crop_size, crop_size, device=images.device)
    if n == 0:
        return out
    windows = windows.to(images.device, torch.float32)
    sample_inds = sample_inds.to(images.device, torch.long)
    i = torch.arange(crop_size, device=images.device, dtype=torch.float32) + 0.5
    step = (windows[:, 2:] - windows[:, :2]) / crop_size  # (N, 2)
    xs = windows[:, 0:1] + i[None, :] * step[:, 0:1]  # (N, crop) edge coords
    ys = windows[:, 1:2] + i[None, :] * step[:, 1:2]
    grid = torch.stack(
        [
            (2 * xs / w - 1)[:, None, :].expand(n, crop_size, crop_size),
            (2 * ys / h - 1)[:, :, None].expand(n, crop_size, crop_size),
        ],
        dim=-1,
    )  # (N, crop, crop, 2)
    frames = images.float()
    for b in torch.unique(sample_inds).tolist():
        rows = (sample_inds == b).nonzero(as_tuple=False).flatten()
        g = grid[rows].reshape(1, -1, crop_size, 2)  # stack this frame's crops
        sampled = F.grid_sample(
            frames[b : b + 1],
            g,
            mode="bilinear",
            padding_mode="border",
            align_corners=False,
        )  # (1, C, k * crop, crop)
        out[rows] = sampled.reshape(c, len(rows), crop_size, crop_size).transpose(0, 1)
    cov_x = _frame_coverage(windows[:, 0], windows[:, 2], crop_size, w)
    cov_y = _frame_coverage(windows[:, 1], windows[:, 3], crop_size, h)
    return out * (cov_y[:, None, :, None] * cov_x[:, None, None, :])

find_instance_crop_size(labels, padding=0, maximum_stride=2, min_crop_size=None, max_hw=None, anchor_ind=None, centroid_method=None, centroid_fallback=None, user_instances_only=False, superseded_only=False)

Compute a crop size that contains every labeled instance.

The size is measured the way the crop is actually taken: centered on each instance's centroid (iter_required_crop_sizes) and, when max_hw is given, in the size-matched pixel space that cropping happens in. Both matter -- a bounding-box measurement in native pixels under-sizes the crop whenever the centroid is off-center or the video is rescaled to max_hw.

Parameters:

Name Type Description Default
labels Labels

A sio.Labels containing user-labeled instances.

required
padding int

Integer number of pixels to add to the bounds as margin padding.

0
maximum_stride int

Ensure that the returned crop size is divisible by this value. Useful for ensuring that the crop size will not be truncated in a given architecture.

2
min_crop_size Optional[int]

A floor for the returned crop size, before padding.

None
max_hw Optional[Tuple[Optional[int], Optional[int]]]

The configured (max_height, max_width), so instances are measured in the space the size matcher puts them in. None (default) measures native pixels.

None
anchor_ind Optional[int]

Index of the anchor node the crop is centered on, or None for a centroid reduce method.

None
centroid_method Optional[str]

The resolved centroid method, or None to infer from anchor_ind.

None
centroid_fallback Optional[str]

The reduce method used when the anchor node is not visible on an instance.

None
user_instances_only bool

When True, ignore predicted instances -- pass data_config.user_instances_only so the crop is not sized from instances that are excluded from training.

False
superseded_only bool

See :func:_iter_frame_instances.

False

Returns:

Type Description
int

An integer crop size denoting the length of the side of the boxes that will contain the instances when cropped. The returned crop size will be larger or equal to the input min_crop_size.

This accounts for stride and padding when ensuring divisibility.

Source code in sleap_nn/data/instance_cropping.py
def find_instance_crop_size(
    labels: sio.Labels,
    padding: int = 0,
    maximum_stride: int = 2,
    min_crop_size: Optional[int] = None,
    max_hw: Optional[Tuple[Optional[int], Optional[int]]] = None,
    anchor_ind: Optional[int] = None,
    centroid_method: Optional[str] = None,
    centroid_fallback: Optional[str] = None,
    user_instances_only: bool = False,
    superseded_only: bool = False,
) -> int:
    """Compute a crop size that contains every labeled instance.

    The size is measured the way the crop is actually taken: centered on each
    instance's centroid (`iter_required_crop_sizes`) and, when ``max_hw`` is
    given, in the size-matched pixel space that cropping happens in. Both
    matter -- a bounding-box measurement in native pixels under-sizes the crop
    whenever the centroid is off-center or the video is rescaled to ``max_hw``.

    Args:
        labels: A `sio.Labels` containing user-labeled instances.
        padding: Integer number of pixels to add to the bounds as margin padding.
        maximum_stride: Ensure that the returned crop size is divisible by this value.
            Useful for ensuring that the crop size will not be truncated in a given
            architecture.
        min_crop_size: A floor for the returned crop size, before padding.
        max_hw: The configured ``(max_height, max_width)``, so instances are
            measured in the space the size matcher puts them in. ``None``
            (default) measures native pixels.
        anchor_ind: Index of the anchor node the crop is centered on, or
            ``None`` for a centroid reduce method.
        centroid_method: The resolved centroid method, or ``None`` to infer from
            ``anchor_ind``.
        centroid_fallback: The reduce method used when the anchor node is not
            visible on an instance.
        user_instances_only: When ``True``, ignore predicted instances -- pass
            ``data_config.user_instances_only`` so the crop is not sized from
            instances that are excluded from training.
        superseded_only: See :func:`_iter_frame_instances`.

    Returns:
        An integer crop size denoting the length of the side of the boxes that
        will contain the instances when cropped. The returned crop size will be
        larger or equal to the input `min_crop_size`.

        This accounts for stride and padding when ensuring divisibility.
    """
    min_crop_size = 0 if min_crop_size is None else min_crop_size

    # `min_crop_size` is a floor, applied before padding is added.
    max_length = float(min_crop_size - padding)
    for required in iter_required_crop_sizes(
        labels,
        max_hw=max_hw,
        anchor_ind=anchor_ind,
        centroid_method=centroid_method,
        centroid_fallback=centroid_fallback,
        user_instances_only=user_instances_only,
        superseded_only=superseded_only,
    ):
        max_length = max(max_length, required)

    max_length = max(max_length, 0.0) + float(padding)
    crop_size = math.ceil(max_length / float(maximum_stride)) * maximum_stride

    return int(crop_size)

find_mask_crop_size(labels, padding=0, maximum_stride=2, min_crop_size=None, max_hw=None, crop_centering='auto', user_instances_only=False, min_mask_area=1.0)

Compute a crop size that contains every mask (the mask analog of find_instance_crop_size).

Parameters:

Name Type Description Default
labels Labels

A sio.Labels whose frames carry masks.

required
padding int

Pixels of margin to add.

0
maximum_stride int

The returned size is divisible by this.

2
min_crop_size Optional[int]

A floor for the size, before padding.

None
max_hw Optional[Tuple[Optional[int], Optional[int]]]

The configured (max_height, max_width).

None
crop_centering str

Where the crop is centered on each mask (see :func:iter_mask_extents).

'auto'
user_instances_only bool

When True, ignore predicted masks.

False
min_mask_area float

Ignore masks smaller than this (image pixels).

1.0

Returns:

Type Description
int

The crop side length, a multiple of maximum_stride.

Source code in sleap_nn/data/instance_cropping.py
def find_mask_crop_size(
    labels: sio.Labels,
    padding: int = 0,
    maximum_stride: int = 2,
    min_crop_size: Optional[int] = None,
    max_hw: Optional[Tuple[Optional[int], Optional[int]]] = None,
    crop_centering: str = "auto",
    user_instances_only: bool = False,
    min_mask_area: float = 1.0,
) -> int:
    """Compute a crop size that contains every mask (the mask analog of `find_instance_crop_size`).

    Args:
        labels: A `sio.Labels` whose frames carry masks.
        padding: Pixels of margin to add.
        maximum_stride: The returned size is divisible by this.
        min_crop_size: A floor for the size, before padding.
        max_hw: The configured ``(max_height, max_width)``.
        crop_centering: Where the crop is centered on each mask (see
            :func:`iter_mask_extents`).
        user_instances_only: When ``True``, ignore predicted masks.
        min_mask_area: Ignore masks smaller than this (image pixels).

    Returns:
        The crop side length, a multiple of ``maximum_stride``.
    """
    min_crop_size = 0 if min_crop_size is None else min_crop_size
    max_length = float(min_crop_size - padding)
    for required, _ in iter_mask_extents(
        labels,
        max_hw=max_hw,
        crop_centering=crop_centering,
        user_instances_only=user_instances_only,
        min_mask_area=min_mask_area,
    ):
        max_length = max(max_length, required)
    max_length = max(max_length, 0.0) + float(padding)
    return int(math.ceil(max_length / float(maximum_stride)) * maximum_stride)

find_max_instance_bbox_size(labels, max_hw=None, user_instances_only=False, superseded_only=False)

Find the maximum bounding box dimension across all instances in labels.

Parameters:

Name Type Description Default
labels Labels

A sio.Labels containing user-labeled instances.

required
max_hw Optional[Tuple[Optional[int], Optional[int]]]

The configured (max_height, max_width). When given, each instance is measured in the size-matched space it will actually be cropped in, rather than in its video's native pixels. None (default) measures native pixels, the historical behavior.

None
user_instances_only bool

When True, ignore predicted instances.

False
superseded_only bool

See :func:_iter_frame_instances.

False

Returns:

Type Description
float

The maximum bounding box dimension (max of width or height) across all instances.

Source code in sleap_nn/data/instance_cropping.py
def find_max_instance_bbox_size(
    labels: sio.Labels,
    max_hw: Optional[Tuple[Optional[int], Optional[int]]] = None,
    user_instances_only: bool = False,
    superseded_only: bool = False,
) -> float:
    """Find the maximum bounding box dimension across all instances in labels.

    Args:
        labels: A `sio.Labels` containing user-labeled instances.
        max_hw: The configured ``(max_height, max_width)``. When given, each
            instance is measured in the size-matched space it will actually be
            cropped in, rather than in its video's native pixels. ``None``
            (default) measures native pixels, the historical behavior.
        user_instances_only: When ``True``, ignore predicted instances.
        superseded_only: See :func:`_iter_frame_instances`.

    Returns:
        The maximum bounding box dimension (max of width or height) across all instances.
    """
    max_length = 0.0
    for lf, pts in _iter_frame_instances(labels, user_instances_only, superseded_only):
        eff_scale = _frame_eff_scale(lf, max_hw)
        diff_x = np.nanmax(pts[:, 0]) - np.nanmin(pts[:, 0])
        diff_x = 0 if np.isnan(diff_x) else diff_x * eff_scale
        max_length = np.maximum(max_length, diff_x)
        diff_y = np.nanmax(pts[:, 1]) - np.nanmin(pts[:, 1])
        diff_y = 0 if np.isnan(diff_y) else diff_y * eff_scale
        max_length = np.maximum(max_length, diff_y)
    return float(max_length)

generate_crops(image, instance, centroid, crop_size)

Generate cropped image for the given centroid.

Parameters:

Name Type Description Default
image Tensor

Input source image. (n_samples, C, H, W)

required
instance Tensor

Keypoints for the instance to be cropped. (n_nodes, 2)

required
centroid Tensor

Centroid of the instance to be cropped. (2)

required
crop_size Tuple[int]

(height, width) of the crop to be generated.

required

Returns:

Type Description
Dict[str, Tensor]

A dictionary with cropped images, bounding box for the cropped instance, keypoints and centroids adjusted to the crop.

Source code in sleap_nn/data/instance_cropping.py
def generate_crops(
    image: torch.Tensor,
    instance: torch.Tensor,
    centroid: torch.Tensor,
    crop_size: Tuple[int],
) -> Dict[str, torch.Tensor]:
    """Generate cropped image for the given centroid.

    Args:
        image: Input source image. (n_samples, C, H, W)
        instance: Keypoints for the instance to be cropped. (n_nodes, 2)
        centroid: Centroid of the instance to be cropped. (2)
        crop_size: (height, width) of the crop to be generated.

    Returns:
        A dictionary with cropped images, bounding box for the cropped instance, keypoints and
        centroids adjusted to the crop.
    """
    box_size = crop_size

    # Generate bounding boxes from centroid.
    instance_bbox = torch.unsqueeze(
        make_centered_bboxes(centroid, box_size[0], box_size[1]), 0
    )  # (n_samples=1, 4, 2)

    # Generate cropped image of shape (n_samples, C, crop_H, crop_W)
    instance_image = crop_and_resize(
        image,
        boxes=instance_bbox,
        size=box_size,
    )

    # Access top left point (x,y) of bounding box and subtract this offset from
    # position of nodes.
    point = instance_bbox[0][0]
    center_instance = (instance - point).unsqueeze(0)  # (n_samples=1, n_nodes, 2)
    centered_centroid = (centroid - point).unsqueeze(0)  # (n_samples=1, 2)

    cropped_sample = {
        "instance_image": instance_image,
        "instance_bbox": instance_bbox,
        "instance": center_instance,
        "centroid": centered_centroid,
    }

    return cropped_sample

iter_mask_extents(labels, max_hw=None, crop_centering='auto', user_instances_only=False, min_mask_area=1.0)

Yield (required_crop_size, bbox_size) for every mask, in size-matched pixels.

The mask analog of :func:iter_required_crop_sizes, for the embedding model's mask mode, whose crop is centered on the mask's center of mass (crop_centering auto / mask_com) or on its bounding-box midpoint (bbox), as EmbeddingDataset._crop_mask centers it. required_crop_size is the side a crop around that center needs to hold the whole mask (twice the largest distance from the center to a mask pixel's far edge); bbox_size is the larger side of the mask's bounding box.

Parameters:

Name Type Description Default
labels Labels

A sio.Labels whose frames carry masks (lf.masks).

required
max_hw Optional[Tuple[Optional[int], Optional[int]]]

The configured (max_height, max_width); masks are measured in the space the size matcher puts the frame in.

None
crop_centering str

auto | mask_com | bbox.

'auto'
user_instances_only bool

When True, skip the predicted masks a user mask on the same frame supersedes, as the embedding dataset does (:func:sleap_nn.data.utils.embedding_detections).

False
min_mask_area float

Skip masks with fewer foreground pixels (image pixels) than this, as the embedding dataset does.

1.0
Source code in sleap_nn/data/instance_cropping.py
def iter_mask_extents(
    labels: sio.Labels,
    max_hw: Optional[Tuple[Optional[int], Optional[int]]] = None,
    crop_centering: str = "auto",
    user_instances_only: bool = False,
    min_mask_area: float = 1.0,
) -> Iterator[Tuple[float, float]]:
    """Yield ``(required_crop_size, bbox_size)`` for every mask, in size-matched pixels.

    The mask analog of :func:`iter_required_crop_sizes`, for the ``embedding`` model's
    mask mode, whose crop is centered on the mask's center of mass (``crop_centering``
    ``auto`` / ``mask_com``) or on its bounding-box midpoint (``bbox``), as
    ``EmbeddingDataset._crop_mask`` centers it. ``required_crop_size`` is the side a
    crop around that center needs to hold the whole mask (twice the largest distance
    from the center to a mask pixel's far edge); ``bbox_size`` is the larger side of
    the mask's bounding box.

    Args:
        labels: A `sio.Labels` whose frames carry masks (``lf.masks``).
        max_hw: The configured ``(max_height, max_width)``; masks are measured in the
            space the size matcher puts the frame in.
        crop_centering: ``auto`` | ``mask_com`` | ``bbox``.
        user_instances_only: When ``True``, skip the predicted masks a user mask on
            the same frame supersedes, as the embedding dataset does
            (:func:`sleap_nn.data.utils.embedding_detections`).
        min_mask_area: Skip masks with fewer foreground pixels (image pixels) than
            this, as the embedding dataset does.
    """
    for lf in labels:
        eff_scale = None
        for mask in embedding_detections(lf, user_instances_only)[1]:
            sx, sy = tuple(getattr(mask, "scale", (1.0, 1.0)) or (1.0, 1.0))
            if mask.area / max(sx * sy, 1e-12) < max(min_mask_area, 1.0):
                continue
            # From the run-length encoding: no full-frame decode per mask.
            rects = mask_foreground_rects(mask)
            if rects.shape[0] == 0:
                continue
            y0, y1 = int(rects[:, 0].min()), int(rects[:, 1].max()) - 1
            x0, x1 = int(rects[:, 2].min()), int(rects[:, 3].max()) - 1
            if crop_centering == "bbox":
                cx, cy = (x0 + x1) / 2.0, (y0 + y1) / 2.0
            else:
                # Center of mass, as offsets from the box corner (the arithmetic
                # these extents were always computed with, so crop sizing does not
                # move by a rounding step).
                n, sum_x, sum_y = mask_rects_moments(rects)
                cx = x0 + float(sum_x - n * x0) / float(n)
                cy = y0 + float(sum_y - n * y0) / float(n)
            if eff_scale is None:
                eff_scale = _frame_eff_scale(lf, max_hw)
            # Pixel centers sit at integer coordinates; +1 reaches their far edges.
            reach = max(cx - x0, x1 - cx, cy - y0, y1 - cy)
            required = (2.0 * reach + 1.0) * eff_scale
            bbox = (max(x1 - x0, y1 - y0) + 1.0) * eff_scale
            yield float(required), float(bbox)

iter_required_crop_sizes(labels, max_hw=None, anchor_ind=None, centroid_method=None, centroid_fallback=None, user_instances_only=False, superseded_only=False)

Yield the crop size each labeled instance needs to avoid being clipped.

Crops are centered on the instance's centroid, not on its bounding-box midpoint (see generate_crops / make_centered_bboxes), so the size an instance requires is twice its greatest node offset from that centroid -- not its bounding-box extent. The two agree only when the centroid happens to sit at the middle of the bounding box. For an anchor node near one end of the animal (a mouse anchored on the thorax, with a long tail) the required size approaches twice the extent, which is why a crop sized to the bounding box still clips the far nodes.

The centroid is derived through generate_centroids, the same op that positions the crop at training time, so this can't disagree with the actual crop center.

Parameters:

Name Type Description Default
labels Labels

A sio.Labels containing user-labeled instances.

required
max_hw Optional[Tuple[Optional[int], Optional[int]]]

The configured (max_height, max_width). When given, points are measured in the size-matched space the crop is taken in.

None
anchor_ind Optional[int]

Index of the anchor node, or None for a centroid reduce method. Passed through to generate_centroids.

None
centroid_method Optional[str]

The resolved centroid method, or None to infer from anchor_ind. See resolve_centroid_method.

None
centroid_fallback Optional[str]

The reduce method used when the anchor node is not visible on an instance.

None
user_instances_only bool

When True, ignore predicted instances.

False
superseded_only bool

See :func:_iter_frame_instances.

False

Yields:

Type Description
float

The required crop side length, in size-matched pixels, per instance. Instances whose centroid is undefined are skipped.

Source code in sleap_nn/data/instance_cropping.py
def iter_required_crop_sizes(
    labels: sio.Labels,
    max_hw: Optional[Tuple[Optional[int], Optional[int]]] = None,
    anchor_ind: Optional[int] = None,
    centroid_method: Optional[str] = None,
    centroid_fallback: Optional[str] = None,
    user_instances_only: bool = False,
    superseded_only: bool = False,
) -> Iterator[float]:
    """Yield the crop size each labeled instance needs to avoid being clipped.

    Crops are **centered on the instance's centroid**, not on its bounding-box
    midpoint (see `generate_crops` / `make_centered_bboxes`), so the size an
    instance requires is twice its greatest node offset from that centroid --
    *not* its bounding-box extent. The two agree only when the centroid happens
    to sit at the middle of the bounding box. For an anchor node near one end of
    the animal (a mouse anchored on the thorax, with a long tail) the required
    size approaches twice the extent, which is why a crop sized to the bounding
    box still clips the far nodes.

    The centroid is derived through `generate_centroids`, the same op that
    positions the crop at training time, so this can't disagree with the actual
    crop center.

    Args:
        labels: A `sio.Labels` containing user-labeled instances.
        max_hw: The configured ``(max_height, max_width)``. When given, points
            are measured in the size-matched space the crop is taken in.
        anchor_ind: Index of the anchor node, or ``None`` for a centroid reduce
            method. Passed through to `generate_centroids`.
        centroid_method: The resolved centroid method, or ``None`` to infer from
            ``anchor_ind``. See `resolve_centroid_method`.
        centroid_fallback: The reduce method used when the anchor node is not
            visible on an instance.
        user_instances_only: When ``True``, ignore predicted instances.
        superseded_only: See :func:`_iter_frame_instances`.

    Yields:
        The required crop side length, in size-matched pixels, per instance.
        Instances whose centroid is undefined are skipped.
    """
    for lf, pts in _iter_frame_instances(labels, user_instances_only, superseded_only):
        eff_scale = _frame_eff_scale(lf, max_hw)
        scaled = torch.from_numpy(pts.astype("float32")) * eff_scale
        centroid = generate_centroids(
            scaled,
            anchor_ind=anchor_ind,
            method=centroid_method,
            fallback=centroid_fallback,
        )
        if torch.isnan(centroid).any():
            continue
        # Largest per-axis offset from the crop center; the crop spans
        # centroid +/- size/2 on each axis, so it must be twice this to reach it.
        offsets = (scaled - centroid).abs()
        reach = torch.nan_to_num(offsets, nan=0.0).max()
        yield float(2.0 * reach)

letterbox_crops(crops, keep_half)

Zero the pixels of (N, C, h, w) crops outside the centered padded box.

Parameters:

Name Type Description Default
crops Tensor

Crops centered on their boxes.

required
keep_half Tensor

(N, 2) half-extent (x, y) of the region to keep, in crop pixels, around the crop center (from :func:bbox_crop_windows).

required
Source code in sleap_nn/data/instance_cropping.py
def letterbox_crops(crops: torch.Tensor, keep_half: torch.Tensor) -> torch.Tensor:
    """Zero the pixels of ``(N, C, h, w)`` crops outside the centered padded box.

    Args:
        crops: Crops centered on their boxes.
        keep_half: ``(N, 2)`` half-extent ``(x, y)`` of the region to keep, in crop
            pixels, around the crop center (from :func:`bbox_crop_windows`).
    """
    n, _, h, w = crops.shape
    if n == 0:
        return crops
    xs = torch.arange(w, device=crops.device, dtype=torch.float32) + 0.5 - w / 2
    ys = torch.arange(h, device=crops.device, dtype=torch.float32) + 0.5 - h / 2
    kh = keep_half.to(crops.device, torch.float32)
    keep_x = xs.abs()[None, :] <= kh[:, 0:1]  # (N, w)
    keep_y = ys.abs()[None, :] <= kh[:, 1:2]  # (N, h)
    keep = keep_y[:, :, None] & keep_x[:, None, :]  # (N, h, w)
    return crops * keep[:, None].to(crops.dtype)

make_centered_bboxes(centroids, box_height, box_width)

Create centered bounding boxes around centroid.

To be used with kornia.geometry.transform.crop_and_resizein the following (clockwise) order: top-left, top-right, bottom-right and bottom-left.

Parameters:

Name Type Description Default
centroids Tensor

A tensor of centroids with shape (n_centroids, 2), where n_centroids is the number of centroids, and the last dimension represents x and y coordinates.

required
box_height int

The desired height of the bounding boxes.

required
box_width int

The desired width of the bounding boxes.

required

Returns:

Type Description
Tensor

torch.Tensor: A tensor containing bounding box coordinates for each centroid. The output tensor has shape (n_centroids, 4, 2), where n_centroids is the number of centroids, and the second dimension represents the four corner points of the bounding boxes, each with x and y coordinates. The order of the corners follows a clockwise arrangement: top-left, top-right, bottom-right, and bottom-left.

Source code in sleap_nn/data/instance_cropping.py
def make_centered_bboxes(
    centroids: torch.Tensor, box_height: int, box_width: int
) -> torch.Tensor:
    """Create centered bounding boxes around centroid.

    To be used with `kornia.geometry.transform.crop_and_resize`in the following
    (clockwise) order: top-left, top-right, bottom-right and bottom-left.

    Args:
        centroids: A tensor of centroids with shape (n_centroids, 2), where n_centroids is the
            number of centroids, and the last dimension represents x and y coordinates.
        box_height: The desired height of the bounding boxes.
        box_width: The desired width of the bounding boxes.

    Returns:
        torch.Tensor: A tensor containing bounding box coordinates for each centroid.
            The output tensor has shape (n_centroids, 4, 2), where n_centroids is the number
            of centroids, and the second dimension represents the four corner points of
            the bounding boxes, each with x and y coordinates. The order of the corners
            follows a clockwise arrangement: top-left, top-right, bottom-right, and
            bottom-left.
    """
    half_h = box_height / 2
    half_w = box_width / 2

    # Get x and y values from the centroids tensor.
    x = centroids[..., 0]
    y = centroids[..., 1]

    # Calculate the corner points.
    top_left = torch.stack([x - half_w, y - half_h], dim=-1)
    top_right = torch.stack([x + half_w, y - half_h], dim=-1)
    bottom_left = torch.stack([x - half_w, y + half_h], dim=-1)
    bottom_right = torch.stack([x + half_w, y + half_h], dim=-1)

    # Get bounding box.
    corners = torch.stack([top_left, top_right, bottom_right, bottom_left], dim=-2)

    offset = torch.tensor([[+0.5, +0.5], [-0.5, +0.5], [-0.5, -0.5], [+0.5, -0.5]]).to(
        corners.device
    )

    return corners + offset

mask_foreground_rects(mask)

A mask's foreground on the image-pixel grid, as disjoint rectangles.

The union of the returned rectangles is exactly the foreground of :func:sleap_nn.inference.segmentation_convert.decode_mask_to_image_res: the stored grid, resampled nearest-neighbor to its image extent when scale is not 1, placed at the rounded offset, with pixels an offset puts above or left of the image dropped. It is read off the run-length encoding, so the cost grows with the number of runs, not with the frame; no full-grid decode is made.

Parameters:

Name Type Description Default
mask SegmentationMask

A sio.SegmentationMask.

required

Returns:

Type Description
ndarray

An (n, 4) int64 array of half-open [y0, y1, x0, x1) rectangles in image pixels, empty when the mask has no foreground pixel on the image.

Source code in sleap_nn/data/instance_cropping.py
def mask_foreground_rects(mask: "sio.SegmentationMask") -> np.ndarray:
    """A mask's foreground on the image-pixel grid, as disjoint rectangles.

    The union of the returned rectangles is exactly the foreground of
    :func:`sleap_nn.inference.segmentation_convert.decode_mask_to_image_res`: the
    stored grid, resampled nearest-neighbor to its image extent when ``scale`` is
    not 1, placed at the rounded ``offset``, with pixels an offset puts above or left
    of the image dropped. It is read off the run-length encoding, so the cost grows
    with the number of runs, not with the frame; no full-grid decode is made.

    Args:
        mask: A ``sio.SegmentationMask``.

    Returns:
        An ``(n, 4)`` int64 array of half-open ``[y0, y1, x0, x1)`` rectangles in
        image pixels, empty when the mask has no foreground pixel on the image.
    """
    height, width = int(mask.height), int(mask.width)
    counts = np.asarray(mask.rle_counts, dtype=np.int64)
    total = height * width
    if total == 0 or counts.size < 2:
        return _NO_RECTS
    ends = np.cumsum(counts)
    # Odd runs are foreground (the encoding starts with a background run); the
    # decoder truncates the runs at the grid size.
    run_start = np.minimum((ends - counts)[1::2], total)
    run_end = np.minimum(ends[1::2], total)
    keep = run_end > run_start
    run_start, run_end = run_start[keep], run_end[keep]
    if run_start.size == 0:
        return _NO_RECTS

    # Split each run (row-major) into one segment per row it covers.
    first_row = run_start // width
    n_rows = (run_end - 1) // width - first_row + 1
    run = np.repeat(np.arange(run_start.size), n_rows)
    rows = first_row[run] + (
        np.arange(run.size) - np.repeat(np.cumsum(n_rows) - n_rows, n_rows)
    )
    col0 = np.maximum(run_start[run], rows * width) - rows * width
    col1 = np.minimum(run_end[run], (rows + 1) * width) - rows * width

    scale = tuple(getattr(mask, "scale", (1.0, 1.0)))
    offset = tuple(getattr(mask, "offset", (0.0, 0.0)))
    if scale == (1.0, 1.0):
        y0, y1, x0, x1 = rows, rows + 1, col0, col1
    else:
        # `SegmentationMask.resampled`'s nearest-neighbor maps (image row i reads
        # stored row row_idx[i]). They never decrease, so the image rows reading a
        # stored row, and the image columns reading a stored column range, are
        # contiguous.
        img_h, img_w = (int(v) for v in mask.image_extent)
        row_idx = (np.arange(img_h) * height / img_h).astype(int).clip(0, height - 1)
        col_idx = (np.arange(img_w) * width / img_w).astype(int).clip(0, width - 1)
        y0 = np.searchsorted(row_idx, rows, side="left")
        y1 = np.searchsorted(row_idx, rows, side="right")
        x0 = np.searchsorted(col_idx, col0, side="left")
        x1 = np.searchsorted(col_idx, col1, side="left")
    # The offset is baked in the way `decode_mask_to_image_res` bakes it.
    ox, oy = int(round(offset[0])), int(round(offset[1]))
    rects = np.stack(
        [
            np.maximum(y0 + oy, 0),
            np.maximum(y1 + oy, 0),
            np.maximum(x0 + ox, 0),
            np.maximum(x1 + ox, 0),
        ],
        axis=1,
    ).astype(np.int64)
    return rects[(rects[:, 1] > rects[:, 0]) & (rects[:, 3] > rects[:, 2])]

mask_rects_center(rects, crop_centering='auto', frame_hw=None)

The (cx, cy) a mask crop centers on, from the mask's rectangles.

auto / mask_com: the center of mass; bbox: the bounding-box midpoint. Both are the values _compute_mask_centroids / _mask_bbox_midpoint compute on the decoded mask, to the last bit: the coordinate sums are exact integers, and the mean is their one rounded quotient, as np.mean of the pixel coordinates is.

Parameters:

Name Type Description Default
rects ndarray

:func:mask_foreground_rects output (clipped to the frame the crop is taken from, see :func:clip_rects).

required
crop_centering str

auto | mask_com | bbox.

'auto'
frame_hw Optional[Tuple[int, int]]

The frame (height, width), for an empty mask, which centers on the frame like the decoded-mask helpers do. Required when rects can be empty.

None
Source code in sleap_nn/data/instance_cropping.py
def mask_rects_center(
    rects: np.ndarray,
    crop_centering: str = "auto",
    frame_hw: Optional[Tuple[int, int]] = None,
) -> Tuple[float, float]:
    """The ``(cx, cy)`` a mask crop centers on, from the mask's rectangles.

    ``auto`` / ``mask_com``: the center of mass; ``bbox``: the bounding-box
    midpoint. Both are the values ``_compute_mask_centroids`` /
    ``_mask_bbox_midpoint`` compute on the decoded mask, to the last bit: the
    coordinate sums are exact integers, and the mean is their one rounded quotient,
    as ``np.mean`` of the pixel coordinates is.

    Args:
        rects: :func:`mask_foreground_rects` output (clipped to the frame the crop
            is taken from, see :func:`clip_rects`).
        crop_centering: ``auto`` | ``mask_com`` | ``bbox``.
        frame_hw: The frame ``(height, width)``, for an empty mask, which centers on
            the frame like the decoded-mask helpers do. Required when ``rects`` can
            be empty.
    """
    if rects.shape[0] == 0:
        if frame_hw is None:
            raise ValueError("An empty mask has no center without the frame size.")
        return float(frame_hw[1]) / 2.0, float(frame_hw[0]) / 2.0
    y0, y1, x0, x1 = rects[:, 0], rects[:, 1], rects[:, 2], rects[:, 3]
    if crop_centering == "bbox":
        return (
            (float(x0.min()) + float(x1.max() - 1)) / 2.0,
            (float(y0.min()) + float(y1.max() - 1)) / 2.0,
        )
    n, sum_x, sum_y = mask_rects_moments(rects)
    return float(sum_x) / float(n), float(sum_y) / float(n)

mask_rects_moments(rects)

(n, sum_x, sum_y): a mask's pixel count and coordinate sums, exactly.

Python integers, from the :func:mask_foreground_rects rectangles: no pixel is enumerated.

Source code in sleap_nn/data/instance_cropping.py
def mask_rects_moments(rects: np.ndarray) -> Tuple[int, int, int]:
    """``(n, sum_x, sum_y)``: a mask's pixel count and coordinate sums, exactly.

    Python integers, from the :func:`mask_foreground_rects` rectangles: no pixel is
    enumerated.
    """
    y0, y1, x0, x1 = rects[:, 0], rects[:, 1], rects[:, 2], rects[:, 3]
    heights, widths = y1 - y0, x1 - x0
    n = int((heights * widths).sum())
    # Sum of the integers a..b-1 is (a + b - 1) * (b - a) / 2, always whole.
    sum_x = int((heights * ((x0 + x1 - 1) * widths // 2)).sum())
    sum_y = int((widths * ((y0 + y1 - 1) * heights // 2)).sum())
    return n, sum_x, sum_y

points_from_crop(points, windows, scales)

Map crop keypoints (N, ..., 2) back to the frame.

The inverse of :func:points_to_crop.

Source code in sleap_nn/data/instance_cropping.py
def points_from_crop(
    points: torch.Tensor, windows: torch.Tensor, scales: torch.Tensor
) -> torch.Tensor:
    """Map crop keypoints ``(N, ..., 2)`` back to the frame.

    The inverse of :func:`points_to_crop`.
    """
    shape = (points.shape[0],) + (1,) * (points.ndim - 2) + (2,)
    origin = windows[:, :2].reshape(shape).to(points)
    s = scales.reshape(shape).to(points)
    return (points + 0.5) / s + origin - 0.5

points_to_crop(points, windows, scales)

Map frame keypoints (N, ..., 2) into the crops of windows.

Source code in sleap_nn/data/instance_cropping.py
def points_to_crop(
    points: torch.Tensor, windows: torch.Tensor, scales: torch.Tensor
) -> torch.Tensor:
    """Map frame keypoints ``(N, ..., 2)`` into the crops of ``windows``."""
    shape = (points.shape[0],) + (1,) * (points.ndim - 2) + (2,)
    origin = windows[:, :2].reshape(shape).to(points)
    s = scales.reshape(shape).to(points)
    return (points + 0.5 - origin) * s - 0.5

pre_crop_size(crop_size)

Side of the training pre-crop that leaves room for rotation augmentation.

At least crop_size * sqrt(2), and of the same parity as crop_size so the final re-crop around the pre-crop center lands on whole pixels (no half-pixel resampling between the pre-crop and the final crop).

Source code in sleap_nn/data/instance_cropping.py
def pre_crop_size(crop_size: int) -> int:
    """Side of the training pre-crop that leaves room for rotation augmentation.

    At least ``crop_size * sqrt(2)``, and of the same parity as ``crop_size`` so the
    final re-crop around the pre-crop center lands on whole pixels (no half-pixel
    resampling between the pre-crop and the final crop).
    """
    m = int(math.ceil(crop_size * math.sqrt(2)))
    return m + ((m - crop_size) % 2)

rasterize_rects(rects, y0, y1, x0, x1)

The [y0:y1, x0:x1] window of the mask rects cover, as a bool array.

Source code in sleap_nn/data/instance_cropping.py
def rasterize_rects(
    rects: np.ndarray, y0: int, y1: int, x0: int, x1: int
) -> np.ndarray:
    """The ``[y0:y1, x0:x1]`` window of the mask ``rects`` cover, as a bool array."""
    height, width = max(y1 - y0, 0), max(x1 - x0, 0)
    if height == 0 or width == 0 or rects.shape[0] == 0:
        return np.zeros((height, width), dtype=bool)
    local = rects - np.array([y0, y0, x0, x0], dtype=np.int64)
    local[:, 0:2] = np.clip(local[:, 0:2], 0, height)
    local[:, 2:4] = np.clip(local[:, 2:4], 0, width)
    local = local[(local[:, 1] > local[:, 0]) & (local[:, 3] > local[:, 2])]
    # One segment per row a rectangle covers, then a per-row difference array: +1
    # where a segment starts, -1 where it ends, a prefix sum along the row fills it.
    n_rows = local[:, 1] - local[:, 0]
    rect = np.repeat(np.arange(local.shape[0]), n_rows)
    rows = local[rect, 0] + (
        np.arange(rect.size) - np.repeat(np.cumsum(n_rows) - n_rows, n_rows)
    )
    diff = np.zeros((height, width + 1), dtype=np.int32)
    np.add.at(diff, (rows, local[rect, 2]), 1)
    np.add.at(diff, (rows, local[rect, 3]), -1)
    return np.cumsum(diff, axis=1, dtype=np.int32)[:, :width] > 0

scale_windows(windows, factor)

Grow (or shrink) (N, 4) edge-coordinate windows about their centers.

Source code in sleap_nn/data/instance_cropping.py
def scale_windows(windows: torch.Tensor, factor: float) -> torch.Tensor:
    """Grow (or shrink) ``(N, 4)`` edge-coordinate windows about their centers."""
    center = (windows[:, :2] + windows[:, 2:]) / 2
    half = (windows[:, 2:] - windows[:, :2]) / 2 * float(factor)
    return torch.cat([center - half, center + half], dim=-1)

window_corners(windows)

(N, 4) edge-coordinate windows -> (N, 4, 2) TL, TR, BR, BL corners.

Source code in sleap_nn/data/instance_cropping.py
def window_corners(windows: torch.Tensor) -> torch.Tensor:
    """``(N, 4)`` edge-coordinate windows -> ``(N, 4, 2)`` TL, TR, BR, BL corners."""
    x1, y1, x2, y2 = windows.unbind(-1)
    return torch.stack(
        [
            torch.stack([x1, y1], -1),
            torch.stack([x2, y1], -1),
            torch.stack([x2, y2], -1),
            torch.stack([x1, y2], -1),
        ],
        dim=-2,
    )