pose2d

Turn video into (x, y) keypoint peaks per (view, point). The subsystem is modular: swappable backbones and heads are composed by an assembler into a detector, with the field geometry and confidence shared across families. The networks -- which run externally-trained checkpoints -- stay behind a torch-free seam so import deeperfly.pose2d never imports torch; the seam and the networks install with the detect extra.

The 2D keypoint detectors and their heatmap decode.

Backbones and heads are composed by detect into a detector (ADR-0008, ADR-0010, ADR-0011), with the field geometry and the confidence shared across families. Torch stays behind a seam, imported lazily, so import deeperfly.pose2d needs no torch (ADR-0009); only the torch-free surface is re-exported here. A window, a resize and the input-normalized space are how a detector reads a frame and stay inside the stage: every point it hands out is in the raw frame of its video (ADR-0076).

FieldGeometry dataclass

The affine between a heatmap field and input-normalized (x, y).

Field cell c is input pixel c * stride - origin, divided by the input size.

Attributes:

Name	Type	Description
input_hw	tuple[int, int]	The reported frame's (height, width), in input pixels.
stride	float	Input pixels per field cell.
origin_xy	tuple[float, float]	The input-pixel location (x, y) of field cell (0, 0).
field_hw	tuple[int, int]	Padded field size in cells, (0, 0) where no fixed grid is tracked.
pad_px	tuple[int, int]	(y, x) input pixels padded onto each side.

Source code in src/deeperfly/pose2d/geometry.py

@dataclass(frozen=True)
class FieldGeometry:
    """The affine between a heatmap field and input-normalized `(x, y)`.

    Field cell `c` is input pixel `c * stride - origin`, divided by the input size.

    Attributes:
        input_hw: The reported frame's `(height, width)`, in input pixels.
        stride: Input pixels per field cell.
        origin_xy: The input-pixel location `(x, y)` of field cell `(0, 0)`.
        field_hw: Padded field size in cells, `(0, 0)` where no fixed grid is tracked.
        pad_px: `(y, x)` input pixels padded onto each side.
    """

    input_hw: tuple[int, int]
    stride: float
    origin_xy: tuple[float, float]
    field_hw: tuple[int, int] = (0, 0)
    pad_px: tuple[int, int] = (0, 0)

    @classmethod
    def padded(
        cls, input_hw: tuple[int, int], *, margin: float, stride: int
    ) -> FieldGeometry:
        """The grid for an input of `input_hw` extended by `margin` on every side.

        `margin` is a fraction of the input, rounded to whole cells.
        """
        h, w = int(input_hw[0]), int(input_hw[1])
        pad = (round(margin * h / stride) * stride, round(margin * w / stride) * stride)
        field = ((h + 2 * pad[0]) // stride, (w + 2 * pad[1]) // stride)
        return cls((h, w), stride, (pad[1], pad[0]), field, pad)

    @classmethod
    def padded_px(
        cls, input_hw: tuple[int, int], *, pad_px: tuple[int, int], stride: int
    ) -> FieldGeometry:
        """The grid for an input extended by `pad_px = (y, x)` pixels on every side.

        Pixels rather than a fraction, so the two axes can pad the same amount; each pad
        must be a whole number of cells.
        """
        h, w = int(input_hw[0]), int(input_hw[1])
        py, px = int(pad_px[0]), int(pad_px[1])
        if py % stride or px % stride:
            raise ValueError(f"pad_px {pad_px} must be whole cells of stride {stride}")
        field = ((h + 2 * py) // stride, (w + 2 * px) // stride)
        return cls((h, w), stride, (px, py), field, (py, px))

    def to_input_normalized(self, cells):
        """Field cells `(..., 2)` in `(x, y)` -> input-normalized; torch or NumPy."""
        ox, oy = self.origin_xy
        h, w = self.input_hw
        x = (cells[..., 0] * self.stride - ox) / w
        y = (cells[..., 1] * self.stride - oy) / h
        if "torch" in type(cells).__module__:
            import torch

            return torch.stack([x, y], dim=-1)
        return np.stack([np.asarray(x), np.asarray(y)], axis=-1)

    def to_field_cells(self, xy):
        """Input-normalized `(..., 2)` -> field cells in `(x, y)`: the inverse.

        The map training targets render against; torch or NumPy.
        """
        ox, oy = self.origin_xy
        h, w = self.input_hw
        cx = (xy[..., 0] * w + ox) / self.stride
        cy = (xy[..., 1] * h + oy) / self.stride
        if "torch" in type(xy).__module__:
            import torch

            return torch.stack([cx, cy], dim=-1)
        return np.stack([np.asarray(cx), np.asarray(cy)], axis=-1)

padded classmethod

padded(
    input_hw: tuple[int, int], *, margin: float, stride: int
) -> FieldGeometry

The grid for an input of input_hw extended by margin on every side.

margin is a fraction of the input, rounded to whole cells.

Source code in src/deeperfly/pose2d/geometry.py

@classmethod
def padded(
    cls, input_hw: tuple[int, int], *, margin: float, stride: int
) -> FieldGeometry:
    """The grid for an input of `input_hw` extended by `margin` on every side.

    `margin` is a fraction of the input, rounded to whole cells.
    """
    h, w = int(input_hw[0]), int(input_hw[1])
    pad = (round(margin * h / stride) * stride, round(margin * w / stride) * stride)
    field = ((h + 2 * pad[0]) // stride, (w + 2 * pad[1]) // stride)
    return cls((h, w), stride, (pad[1], pad[0]), field, pad)

padded_px classmethod

padded_px(
    input_hw: tuple[int, int],
    *,
    pad_px: tuple[int, int],
    stride: int,
) -> FieldGeometry

The grid for an input extended by pad_px = (y, x) pixels on every side.

Pixels rather than a fraction, so the two axes can pad the same amount; each pad must be a whole number of cells.

Source code in src/deeperfly/pose2d/geometry.py

@classmethod
def padded_px(
    cls, input_hw: tuple[int, int], *, pad_px: tuple[int, int], stride: int
) -> FieldGeometry:
    """The grid for an input extended by `pad_px = (y, x)` pixels on every side.

    Pixels rather than a fraction, so the two axes can pad the same amount; each pad
    must be a whole number of cells.
    """
    h, w = int(input_hw[0]), int(input_hw[1])
    py, px = int(pad_px[0]), int(pad_px[1])
    if py % stride or px % stride:
        raise ValueError(f"pad_px {pad_px} must be whole cells of stride {stride}")
    field = ((h + 2 * py) // stride, (w + 2 * px) // stride)
    return cls((h, w), stride, (px, py), field, (py, px))

to_field_cells

to_field_cells(xy)

Input-normalized (..., 2) -> field cells in (x, y): the inverse.

The map training targets render against; torch or NumPy.

Source code in src/deeperfly/pose2d/geometry.py

def to_field_cells(self, xy):
    """Input-normalized `(..., 2)` -> field cells in `(x, y)`: the inverse.

    The map training targets render against; torch or NumPy.
    """
    ox, oy = self.origin_xy
    h, w = self.input_hw
    cx = (xy[..., 0] * w + ox) / self.stride
    cy = (xy[..., 1] * h + oy) / self.stride
    if "torch" in type(xy).__module__:
        import torch

        return torch.stack([cx, cy], dim=-1)
    return np.stack([np.asarray(cx), np.asarray(cy)], axis=-1)

to_input_normalized

to_input_normalized(cells)

Field cells (..., 2) in (x, y) -> input-normalized; torch or NumPy.

Source code in src/deeperfly/pose2d/geometry.py

def to_input_normalized(self, cells):
    """Field cells `(..., 2)` in `(x, y)` -> input-normalized; torch or NumPy."""
    ox, oy = self.origin_xy
    h, w = self.input_hw
    x = (cells[..., 0] * self.stride - ox) / w
    y = (cells[..., 1] * self.stride - oy) / h
    if "torch" in type(cells).__module__:
        import torch

        return torch.stack([x, y], dim=-1)
    return np.stack([np.asarray(x), np.asarray(y)], axis=-1)

Device, dtype and tensor plumbing shared by every detector class.

It imports torch at module scope, behind the torch-free seam, so import deeperfly.pose2d never imports it.

accumulation

accumulation(dtype: str = 'float32')

Float16 matrix products accumulate in dtype inside the block (ADR-0150).

"float16" sets torch's process-wide allow_fp16_accumulation and puts back what it was on the way out; "float32", or a torch without the switch, changes nothing.

Source code in src/deeperfly/pose2d/runtime.py

@contextlib.contextmanager
def accumulation(dtype: str = "float32"):
    """Float16 matrix products accumulate in `dtype` inside the block (ADR-0150).

    `"float16"` sets torch's process-wide `allow_fp16_accumulation` and puts back what
    it was on the way out; `"float32"`, or a torch without the switch, changes nothing.
    """
    matmul = torch.backends.cuda.matmul
    before = getattr(matmul, "allow_fp16_accumulation", None)
    if dtype != "float16" or before is None:
        yield
        return
    matmul.allow_fp16_accumulation = True
    try:
        yield
    finally:
        matmul.allow_fp16_accumulation = before

as_torch

as_torch(inputs) -> torch.Tensor

Network inputs as a torch tensor, on-device when possible.

A torch tensor passes through, another DLPack-capable array is bridged zero-copy, and host NumPy is copied to a writable tensor.

Source code in src/deeperfly/pose2d/runtime.py

def as_torch(inputs) -> "torch.Tensor":
    """Network inputs as a torch tensor, on-device when possible.

    A torch tensor passes through, another DLPack-capable array is bridged zero-copy,
    and host NumPy is copied to a writable tensor.
    """
    if isinstance(inputs, torch.Tensor):
        return inputs
    if hasattr(inputs, "__dlpack__"):  # other on-device array -- zero-copy
        return torch.from_dlpack(inputs)
    return torch.from_numpy(np.array(inputs))  # host array -> writable copy

autocast_dtype

autocast_dtype(model, dev: device)

The autocast dtype for the forward, or None to run in float32.

Only on CUDA, CPU and MPS staying on the float32 reference; "auto" is bfloat16 with tensor cores and float16 below.

Source code in src/deeperfly/pose2d/runtime.py

def autocast_dtype(model, dev: "torch.device"):
    """The autocast dtype for the forward, or `None` to run in float32.

    Only on CUDA, CPU and MPS staying on the float32 reference; `"auto"` is bfloat16
    with tensor cores and float16 below.
    """
    if dev.type != "cuda":
        return None
    precision = getattr(model, "_deeperfly_precision", "float32")
    if precision == "auto":
        precision = "bfloat16" if has_bf16_tensor_cores() else "float16"
    return _PRECISIONS.get(precision)

detector_device

detector_device(model) -> str

The device the detector's parameters live on ("cuda:0", "cpu").

Read off the parameters, so torch is never imported; "cpu" for a parameterless model.

Source code in src/deeperfly/pose2d/runtime.py

def detector_device(model) -> str:
    """The device the detector's parameters live on (`"cuda:0"`, `"cpu"`).

    Read off the parameters, so torch is never imported; `"cpu"` for a parameterless
    model.
    """
    params = getattr(model, "parameters", None)
    if params is None:
        return "cpu"
    try:
        return str(next(params()).device)
    except StopIteration:
        return "cpu"

device

device() -> str

The best torch device available: "cuda", "mps" (Apple Metal) or "cpu".

Source code in src/deeperfly/pose2d/runtime.py

def device() -> str:
    """The best torch device available: `"cuda"`, `"mps"` (Apple Metal) or `"cpu"`."""
    if torch.cuda.is_available():
        return "cuda"
    if torch.backends.mps.is_available():
        return "mps"
    return "cpu"

has_bf16_tensor_cores

has_bf16_tensor_cores() -> bool

Whether the CUDA device runs bfloat16 on tensor cores (capability 8.0 and up).

Not is_bf16_supported, which answers true on Turing, where bfloat16 is emulated (ADR-0069).

Source code in src/deeperfly/pose2d/runtime.py

def has_bf16_tensor_cores() -> bool:
    """Whether the CUDA device runs bfloat16 on tensor cores (capability 8.0 and up).

    Not `is_bf16_supported`, which answers true on Turing, where bfloat16 is emulated
    (ADR-0069).
    """
    return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8

set_precision

set_precision(model, precision: str = 'float32') -> None

Record the forward precision on the detector, where the forward reads it.

"float32" (the reference), "float16" or "bfloat16" (CUDA autocast), or "auto", resolved per card at forward time (ADR-0069).

Source code in src/deeperfly/pose2d/runtime.py

def set_precision(model, precision: str = "float32") -> None:
    """Record the forward precision on the detector, where the forward reads it.

    `"float32"` (the reference), `"float16"` or `"bfloat16"` (CUDA autocast), or
    `"auto"`, resolved per card at forward time (ADR-0069).
    """
    precision = (precision or "float32").lower()
    if precision not in PRECISIONS:
        opts = ", ".join(repr(p) for p in PRECISIONS)
        raise ValueError(f"unknown detector precision {precision!r}; use one of {opts}")
    # A real detector (nn.Module) carries a __dict__; bare test stubs do not (and never
    # run the real forward), so there is nothing to set on them.
    if hasattr(model, "__dict__"):
        model._deeperfly_precision = precision

The released detectors: which exist, and each one's copy in the cache (ADR-0151).

RELEASED is the catalog this deeperfly ships with (data/released.toml): the detectors and locators published on Hugging Face, each pinned to a revision of the repository and to the SHA-256 of its checkpoint. A config names one where it would otherwise give a path, and artifact_file turns either into a file, downloading a model the cache does not hold yet (ADR-0153). Only download reaches the network; deeperfly models download calls it ahead of a run, and deeperfly models remove deletes what it fetched (remove). Torch-free.

Released dataclass

One released model: a folder of the repository and the checkpoint in it.

Source code in src/deeperfly/pose2d/released.py

@dataclass(frozen=True)
class Released:
    """One released model: a folder of the repository and the checkpoint in it."""

    name: str
    rig: str
    role: Literal["detector", "locator"]
    file: str
    size: int
    sha256: str

    @property
    def path(self) -> Path:
        """Where the checkpoint is once downloaded: `models/<name>/<file>`."""
        return cache_dir() / "models" / self.name / self.file

    @property
    def downloaded(self) -> bool:
        """Whether the cache holds the checkpoint."""
        return self.path.is_file()

downloaded property

downloaded: bool

Whether the cache holds the checkpoint.

path property

path: Path

Where the checkpoint is once downloaded: models/<name>/<file>.

artifact_file

artifact_file(ref: str | Path) -> Path

The file a detector reference names: a released model's, or the path itself.

A released model the cache does not hold yet is downloaded first (download), so the first use of one needs a network.

Raises:

Type	Description
RuntimeError	If it names a released model that cannot be downloaded.

Source code in src/deeperfly/pose2d/released.py

def artifact_file(ref: str | Path) -> Path:
    """The file a detector reference names: a released model's, or the path itself.

    A released model the cache does not hold yet is downloaded first
    ([download][deeperfly.pose2d.released.download]), so the first use of one needs
    a network.

    Raises:
        RuntimeError: If it names a released model that cannot be downloaded.
    """
    model = RELEASED.get(str(ref))
    if model is None:
        return Path(ref).expanduser()
    if not model.downloaded:
        download([model.name])
    return model.path

download

download(
    names: Iterable[str],
    *,
    log: Callable[[str], object] = print,
) -> list[Path]

Fetch the released models named into the cache; their files.

A model's whole folder comes with it: the checkpoint, its README.md, the training recipe and report. One already there is not fetched again, and every checkpoint is checked against the catalog's SHA-256 either way.

Raises:

Type	Description
KeyError	If a name is not a released model.
RuntimeError	If the detect extra is missing, the download fails, or a checkpoint is not the released file.

Source code in src/deeperfly/pose2d/released.py

def download(
    names: Iterable[str], *, log: Callable[[str], object] = print
) -> list[Path]:
    """Fetch the released models named into the cache; their files.

    A model's whole folder comes with it: the checkpoint, its `README.md`, the
    training recipe and report. One already there is not fetched again, and every
    checkpoint is checked against the catalog's SHA-256 either way.

    Raises:
        KeyError: If a name is not a released model.
        RuntimeError: If the `detect` extra is missing, the download fails, or a
            checkpoint is not the released file.
    """
    models = [RELEASED[name] for name in names]
    root = cache_dir() / "models"
    for model in models:
        if model.downloaded:
            log(f"{model.name}: already downloaded")
        else:
            log(f"{model.name}: downloading {model.size / 1e6:.0f} MB to {root}")
            _fetch(model, root)
        if _sha256(model.path) != model.sha256:
            raise RuntimeError(
                f"{model.path} is not the released file (its SHA-256 differs); "
                "delete it and download again"
            )
    return [model.path for model in models]

remove

remove(names: Iterable[str]) -> dict[str, int]

Delete the released models named from the cache; the bytes each held.

A model's whole folder goes, and a model the cache does not hold is left out of the result. A models/ that is a link to a shared copy is refused, as clear leaves what a link points at alone; that is also what deletes every model at once.

Raises:

Type	Description
KeyError	If a name is not a released model.
RuntimeError	If models/ is a link.

Source code in src/deeperfly/pose2d/released.py

def remove(names: Iterable[str]) -> dict[str, int]:
    """Delete the released models named from the cache; the bytes each held.

    A model's whole folder goes, and a model the cache does not hold is left out of
    the result. A `models/` that is a link to a shared copy is refused, as
    [clear][deeperfly.cache.clear] leaves what a link points at alone; that is also
    what deletes every model at once.

    Raises:
        KeyError: If a name is not a released model.
        RuntimeError: If `models/` is a link.
    """
    models = [RELEASED[name] for name in names]
    root = cache_dir() / "models"
    if root.is_symlink():
        raise RuntimeError(
            f"{root} is a link to {root.resolve()}, and a model is not deleted "
            "through a link"
        )
    freed = {}
    for model in models:
        folder = model.path.parent
        if folder.is_dir():
            freed[model.name] = tree_size(folder)
            shutil.rmtree(folder)
    return freed

The padded-field affine: heatmap cells <-> input-normalized coordinates.

A heatmap field may extend past the input by a margin, so a joint pushed off the frame still has a cell to peak in. FieldGeometry is the one place that padding is undone, for the monocular grid (stride 4) and the multiview net's upsampled one (stride 1). Coordinates are unclamped: outside [0, 1] is an off-frame joint, which is signal.

FieldGeometry dataclass

The affine between a heatmap field and input-normalized (x, y).

Field cell c is input pixel c * stride - origin, divided by the input size.

Attributes:

Name	Type	Description
input_hw	tuple[int, int]	The reported frame's (height, width), in input pixels.
stride	float	Input pixels per field cell.
origin_xy	tuple[float, float]	The input-pixel location (x, y) of field cell (0, 0).
field_hw	tuple[int, int]	Padded field size in cells, (0, 0) where no fixed grid is tracked.
pad_px	tuple[int, int]	(y, x) input pixels padded onto each side.

Source code in src/deeperfly/pose2d/geometry.py

@dataclass(frozen=True)
class FieldGeometry:
    """The affine between a heatmap field and input-normalized `(x, y)`.

    Field cell `c` is input pixel `c * stride - origin`, divided by the input size.

    Attributes:
        input_hw: The reported frame's `(height, width)`, in input pixels.
        stride: Input pixels per field cell.
        origin_xy: The input-pixel location `(x, y)` of field cell `(0, 0)`.
        field_hw: Padded field size in cells, `(0, 0)` where no fixed grid is tracked.
        pad_px: `(y, x)` input pixels padded onto each side.
    """

    input_hw: tuple[int, int]
    stride: float
    origin_xy: tuple[float, float]
    field_hw: tuple[int, int] = (0, 0)
    pad_px: tuple[int, int] = (0, 0)

    @classmethod
    def padded(
        cls, input_hw: tuple[int, int], *, margin: float, stride: int
    ) -> FieldGeometry:
        """The grid for an input of `input_hw` extended by `margin` on every side.

        `margin` is a fraction of the input, rounded to whole cells.
        """
        h, w = int(input_hw[0]), int(input_hw[1])
        pad = (round(margin * h / stride) * stride, round(margin * w / stride) * stride)
        field = ((h + 2 * pad[0]) // stride, (w + 2 * pad[1]) // stride)
        return cls((h, w), stride, (pad[1], pad[0]), field, pad)

    @classmethod
    def padded_px(
        cls, input_hw: tuple[int, int], *, pad_px: tuple[int, int], stride: int
    ) -> FieldGeometry:
        """The grid for an input extended by `pad_px = (y, x)` pixels on every side.

        Pixels rather than a fraction, so the two axes can pad the same amount; each pad
        must be a whole number of cells.
        """
        h, w = int(input_hw[0]), int(input_hw[1])
        py, px = int(pad_px[0]), int(pad_px[1])
        if py % stride or px % stride:
            raise ValueError(f"pad_px {pad_px} must be whole cells of stride {stride}")
        field = ((h + 2 * py) // stride, (w + 2 * px) // stride)
        return cls((h, w), stride, (px, py), field, (py, px))

    def to_input_normalized(self, cells):
        """Field cells `(..., 2)` in `(x, y)` -> input-normalized; torch or NumPy."""
        ox, oy = self.origin_xy
        h, w = self.input_hw
        x = (cells[..., 0] * self.stride - ox) / w
        y = (cells[..., 1] * self.stride - oy) / h
        if "torch" in type(cells).__module__:
            import torch

            return torch.stack([x, y], dim=-1)
        return np.stack([np.asarray(x), np.asarray(y)], axis=-1)

    def to_field_cells(self, xy):
        """Input-normalized `(..., 2)` -> field cells in `(x, y)`: the inverse.

        The map training targets render against; torch or NumPy.
        """
        ox, oy = self.origin_xy
        h, w = self.input_hw
        cx = (xy[..., 0] * w + ox) / self.stride
        cy = (xy[..., 1] * h + oy) / self.stride
        if "torch" in type(xy).__module__:
            import torch

            return torch.stack([cx, cy], dim=-1)
        return np.stack([np.asarray(cx), np.asarray(cy)], axis=-1)

padded classmethod

padded(
    input_hw: tuple[int, int], *, margin: float, stride: int
) -> FieldGeometry

The grid for an input of input_hw extended by margin on every side.

margin is a fraction of the input, rounded to whole cells.

Source code in src/deeperfly/pose2d/geometry.py

@classmethod
def padded(
    cls, input_hw: tuple[int, int], *, margin: float, stride: int
) -> FieldGeometry:
    """The grid for an input of `input_hw` extended by `margin` on every side.

    `margin` is a fraction of the input, rounded to whole cells.
    """
    h, w = int(input_hw[0]), int(input_hw[1])
    pad = (round(margin * h / stride) * stride, round(margin * w / stride) * stride)
    field = ((h + 2 * pad[0]) // stride, (w + 2 * pad[1]) // stride)
    return cls((h, w), stride, (pad[1], pad[0]), field, pad)

padded_px classmethod

padded_px(
    input_hw: tuple[int, int],
    *,
    pad_px: tuple[int, int],
    stride: int,
) -> FieldGeometry

The grid for an input extended by pad_px = (y, x) pixels on every side.

Pixels rather than a fraction, so the two axes can pad the same amount; each pad must be a whole number of cells.

Source code in src/deeperfly/pose2d/geometry.py

@classmethod
def padded_px(
    cls, input_hw: tuple[int, int], *, pad_px: tuple[int, int], stride: int
) -> FieldGeometry:
    """The grid for an input extended by `pad_px = (y, x)` pixels on every side.

    Pixels rather than a fraction, so the two axes can pad the same amount; each pad
    must be a whole number of cells.
    """
    h, w = int(input_hw[0]), int(input_hw[1])
    py, px = int(pad_px[0]), int(pad_px[1])
    if py % stride or px % stride:
        raise ValueError(f"pad_px {pad_px} must be whole cells of stride {stride}")
    field = ((h + 2 * py) // stride, (w + 2 * px) // stride)
    return cls((h, w), stride, (px, py), field, (py, px))

to_field_cells

to_field_cells(xy)

Input-normalized (..., 2) -> field cells in (x, y): the inverse.

The map training targets render against; torch or NumPy.

Source code in src/deeperfly/pose2d/geometry.py

def to_field_cells(self, xy):
    """Input-normalized `(..., 2)` -> field cells in `(x, y)`: the inverse.

    The map training targets render against; torch or NumPy.
    """
    ox, oy = self.origin_xy
    h, w = self.input_hw
    cx = (xy[..., 0] * w + ox) / self.stride
    cy = (xy[..., 1] * h + oy) / self.stride
    if "torch" in type(xy).__module__:
        import torch

        return torch.stack([cx, cy], dim=-1)
    return np.stack([np.asarray(cx), np.asarray(cy)], axis=-1)

to_input_normalized

to_input_normalized(cells)

Field cells (..., 2) in (x, y) -> input-normalized; torch or NumPy.

Source code in src/deeperfly/pose2d/geometry.py

def to_input_normalized(self, cells):
    """Field cells `(..., 2)` in `(x, y)` -> input-normalized; torch or NumPy."""
    ox, oy = self.origin_xy
    h, w = self.input_hw
    x = (cells[..., 0] * self.stride - ox) / w
    y = (cells[..., 1] * self.stride - oy) / h
    if "torch" in type(cells).__module__:
        import torch

        return torch.stack([x, y], dim=-1)
    return np.stack([np.asarray(x), np.asarray(y)], axis=-1)

Per-point confidence strategies, and the name-keyed registry over them.

A detector records its strategy as confidence_kind. peak_amax is the raw peak value, cheap but blind to a broad peak; softmax_window_mass the softmax mass near the peak (the soft-argmax decode's); peak_mass the same on a probability field itself (the DARK decode's), about 0.85 for one blob of the training sigma. Torch is imported lazily.

peak_amax

peak_amax(heatmaps)

(..., H, W) field -> (...) peak value: the maximum over the last two axes.

Source code in src/deeperfly/pose2d/confidence.py

def peak_amax(heatmaps):
    """`(..., H, W)` field -> `(...)` peak value: the maximum over the last two axes."""
    return heatmaps.flatten(-2).amax(dim=-1)

peak_mass

peak_mass(prob, *, radius: int = PEAK_MASS_RADIUS)

Probability mass within radius cells of each channel's peak, in [0, 1].

Source code in src/deeperfly/pose2d/confidence.py

def peak_mass(prob, *, radius: int = PEAK_MASS_RADIUS):
    """Probability mass within `radius` cells of each channel's peak, in `[0, 1]`."""
    *lead, h, w = prob.shape
    flat = prob.reshape(-1, h, w)
    idx = flat.reshape(flat.shape[0], -1).argmax(dim=-1)
    mass = _window_mass(flat[:, None], (idx // w)[:, None], (idx % w)[:, None], radius)
    return mass.reshape(*lead)

peak_mass_at

peak_mass_at(
    prob, cells, *, radius: int = PEAK_MASS_RADIUS
)

peak_mass at the given cells, for the candidate path.

prob is (N, H, W) and cells integer (N, M, 2) in (x, y); the result is (N, M). One box filter over the field, gathered at the cells.

Source code in src/deeperfly/pose2d/confidence.py

def peak_mass_at(prob, cells, *, radius: int = PEAK_MASS_RADIUS):
    """`peak_mass` at the given cells, for the candidate path.

    `prob` is `(N, H, W)` and `cells` integer `(N, M, 2)` in `(x, y)`; the result is
    `(N, M)`. One box filter over the field, gathered at the cells.
    """
    import torch
    import torch.nn.functional as F

    k = 2 * radius + 1
    mass = F.avg_pool2d(prob.float()[:, None], k, stride=1, padding=radius)[:, 0]
    mass = mass * float(k * k)
    ar = torch.arange(prob.shape[0], device=prob.device)[:, None]
    return mass[ar, cells[..., 1], cells[..., 0]]

softmax_window_mass

softmax_window_mass(
    prob,
    locs,
    *,
    sigma: float = WINDOW_SIGMA,
    num_stds: float = WINDOW_NUM_STDS,
)

The softmax mass within floor(sigma * num_stds) cells of each peak.

prob is a (B, C, H, W) spatial softmax and locs the (B, C, 2) peaks in (x, y) cells, fractional for an expectation; the window is centered on the nearest cell, since a cell's coordinate is its center (ADR-0055). The result is (B, C) in [0, 1].

Source code in src/deeperfly/pose2d/confidence.py

def softmax_window_mass(
    prob, locs, *, sigma: float = WINDOW_SIGMA, num_stds: float = WINDOW_NUM_STDS
):
    """The softmax mass within `floor(sigma * num_stds)` cells of each peak.

    `prob` is a `(B, C, H, W)` spatial softmax and `locs` the `(B, C, 2)` peaks in
    `(x, y)` cells, fractional for an expectation; the window is centered on the
    nearest cell, since a cell's coordinate is its center (ADR-0055). The result is
    `(B, C)` in `[0, 1]`.
    """
    reach = int(np.floor(sigma * num_stds))
    cells = locs.round().long()
    return _window_mass(prob, cells[..., 1], cells[..., 0], reach)

The backbone-head coupling: feature maps selected by output stride.

A head names its levels by stride, not index: HRNet returns a stride-2 stem at index 0 where HGNetV2 and ResNet start at stride 4, so an index would feed the head the wrong maps and evaluate a plausible, wrong model.

select_by_stride

select_by_stride(
    feature_info,
    want_reductions: tuple[int, ...],
    *,
    pinned: tuple[int, ...] | None = None,
    name: str = "",
) -> tuple[tuple[int, ...], tuple[int, ...]]

(indices, channels) of feature_info's levels at want_reductions.

pinned checks the widths trained weights assume (None accepts timm's, which the strict=True load then confirms); name is for the error message.

Raises:

Type	Description
RuntimeError	If a wanted stride is absent, or pinned disagrees with timm.

Source code in src/deeperfly/pose2d/backbones/base.py

def select_by_stride(
    feature_info,
    want_reductions: tuple[int, ...],
    *,
    pinned: tuple[int, ...] | None = None,
    name: str = "",
) -> tuple[tuple[int, ...], tuple[int, ...]]:
    """`(indices, channels)` of `feature_info`'s levels at `want_reductions`.

    `pinned` checks the widths trained weights assume (`None` accepts `timm`'s, which
    the `strict=True` load then confirms); `name` is for the error message.

    Raises:
        RuntimeError: If a wanted stride is absent, or `pinned` disagrees with `timm`.
    """
    red = tuple(feature_info.reduction())
    missing = [r for r in want_reductions if r not in red]
    if missing:
        raise RuntimeError(
            f"{name} exposes feature strides {red}; the head needs {want_reductions} "
            f"and {missing} are absent"
        )
    sel = tuple(red.index(r) for r in want_reductions)
    all_chs = tuple(feature_info.channels())
    chs = tuple(all_chs[i] for i in sel)
    if pinned is not None and chs != pinned:
        raise RuntimeError(
            f"{name} feature channels {chs}, expected {pinned}; the trained weights "
            "assume the latter"
        )
    return sel, chs

The CNN pyramid backbone: a timm features_only net at strides 4 to 32.

create returns the raw model, bound unwrapped for strict=True loading, and the levels a head needs; forward_features returns them finest first, NCHW.

create

create(
    backbone: str,
    *,
    input_hw: tuple[int, int],
    pretrained: bool = False,
    in_chans: int = 1,
)

Build the timm backbone and its stride selection.

in_chans is 1 for grayscale video. Returns (model, sel, channels): the model, the feature_info indices for WANT_REDUCTIONS, and each level's width.

Source code in src/deeperfly/pose2d/backbones/cnn_pyramid.py

def create(
    backbone: str,
    *,
    input_hw: tuple[int, int],
    pretrained: bool = False,
    in_chans: int = 1,
):
    """Build the `timm` backbone and its stride selection.

    `in_chans` is 1 for grayscale video. Returns `(model, sel, channels)`: the model,
    the `feature_info` indices for `WANT_REDUCTIONS`, and each level's width.
    """
    import timm

    kw = {"pretrained": pretrained, "features_only": True, "in_chans": in_chans}
    try:
        model = timm.create_model(backbone, img_size=input_hw, **kw)
    except TypeError:
        model = timm.create_model(backbone, **kw)
    sel, chs = select_by_stride(
        model.feature_info,
        WANT_REDUCTIONS,
        pinned=PINNED_CHANNELS.get(backbone),
        name=backbone,
    )
    return model, sel, chs

forward_features

forward_features(
    model,
    sel: tuple[int, ...],
    channels: tuple[int, ...],
    x,
)

Run model on x: the selected feature maps, finest first, NCHW.

NHWC backbones (swin, maxvit) are detected from the tensor and permuted.

Source code in src/deeperfly/pose2d/backbones/cnn_pyramid.py

def forward_features(model, sel: tuple[int, ...], channels: tuple[int, ...], x):
    """Run `model` on `x`: the selected feature maps, finest first, NCHW.

    NHWC backbones (swin, maxvit) are detected from the tensor and permuted.
    """
    raw = list(model(x))
    fs = [raw[i] for i in sel]
    if fs[0].shape[1] != channels[0]:
        fs = [f.permute(0, 3, 1, 2).contiguous() for f in fs]
    return fs

The multiview ViT backbone: every view of a frame encoded together.

A timm ViT whose blocks follow a cross-view attention schedule, a regrouping of the token axis between blocks with no per-view parameters: view-local ("l", tokens (B*V, P + N, dim)) or global ("g", (B, V*P + V*N, dim), every view's prefix tokens first). A rope backbone (DINOv3) gets its table tiled V times, so a patch has one position in every view. With no per-view parameters the function is permutation-equivariant over views and V is free; padded views (view_valid) are masked out of every global block's keys.

create

create(
    backbone: str,
    *,
    input_hw: tuple[int, int],
    attn_scopes: str,
    pretrained: bool = False,
    drop_path_rate: float = 0.0,
)

Build the timm ViT and validate the schedule against its block count.

The backbone sees the reported frame only; the margin is added in the head. The pretrained prefix tokens (class token, registers) are kept, since timm's pretrained load assumes them, and dropped in encode. Returns (model, dim, patch): the model, its embed dim and its (ph, pw) patch.

Source code in src/deeperfly/pose2d/backbones/vit_multiview.py

def create(
    backbone: str,
    *,
    input_hw: tuple[int, int],
    attn_scopes: str,
    pretrained: bool = False,
    drop_path_rate: float = 0.0,
):
    """Build the `timm` ViT and validate the schedule against its block count.

    The backbone sees the reported frame only; the margin is added in the head. The
    pretrained prefix tokens (class token, registers) are kept, since `timm`'s
    pretrained load assumes them, and dropped in `encode`. Returns
    `(model, dim, patch)`: the model, its embed dim and its `(ph, pw)` patch.
    """
    import timm

    model = timm.create_model(
        backbone,
        pretrained=pretrained,
        img_size=tuple(input_hw),
        in_chans=1,
        num_classes=0,
        global_pool="",
        dynamic_img_size=True,
        drop_path_rate=float(drop_path_rate),
    )
    dim = int(model.embed_dim)
    n_blocks = len(model.blocks)
    if len(attn_scopes) != n_blocks:
        raise RuntimeError(
            f"attn_scopes has {len(attn_scopes)} entries but {backbone} has "
            f"{n_blocks} blocks"
        )
    ph, pw = model.patch_embed.patch_size
    _memoize_rope(model)
    return model, dim, (int(ph), int(pw))

encode

encode(
    model,
    attn_scopes,
    images,
    dim,
    *,
    norm_mean=0.0,
    norm_std=1.0,
    view_valid=None,
    mask_token=None,
    mask_ratio: float = 0.0,
    pad_tokens: tuple[int, int] = (0, 0),
    pad_token=None,
    global_rope: bool = True,
    checkpoint: bool = False,
)

(B, V, 1, H, W) synchronized views -> (B*V, dim, Hp, Wp), batch-major in V.

The patch and positional embeds run per view, then the block schedule regroups the tokens at each "l"/"g" transition; the prefix tokens are dropped after the final norm. mask_ratio (a number or a (B*V,) tensor) replaces that fraction of each view's patches with mask_token before the positional embed, so a masked patch keeps its position and loses its content. pad_tokens (py, px) of pad_token surround each view's grid before the blocks, giving the margin cells attention can write. global_rope=False withholds the rope from the global blocks, and checkpoint recomputes each block in the backward pass.

Source code in src/deeperfly/pose2d/backbones/vit_multiview.py

def encode(
    model,
    attn_scopes,
    images,
    dim,
    *,
    norm_mean=0.0,
    norm_std=1.0,
    view_valid=None,
    mask_token=None,
    mask_ratio: float = 0.0,
    pad_tokens: tuple[int, int] = (0, 0),
    pad_token=None,
    global_rope: bool = True,
    checkpoint: bool = False,
):
    """`(B, V, 1, H, W)` synchronized views -> `(B*V, dim, Hp, Wp)`, batch-major in V.

    The patch and positional embeds run per view, then the block schedule regroups the
    tokens at each `"l"`/`"g"` transition; the prefix tokens are dropped after the final
    norm. `mask_ratio` (a number or a `(B*V,)` tensor) replaces that fraction of each
    view's patches with `mask_token` before the positional embed, so a masked patch
    keeps its position and loses its content. `pad_tokens` `(py, px)` of `pad_token`
    surround each view's grid before the blocks, giving the margin cells attention can
    write. `global_rope=False` withholds the rope from the global blocks, and
    `checkpoint` recomputes each block in the backward pass.
    """
    b, v = images.shape[0], images.shape[1]
    flat = images.reshape(-1, *images.shape[2:])
    g = (flat - norm_mean) / norm_std
    tok = model.patch_embed(g)  # (B*V, Hp, Wp, dim)
    h_tok, w_tok = tok.shape[1], tok.shape[2]
    if mask_token is not None and mask_ratio is not None:
        ratio = torch.as_tensor(mask_ratio, device=tok.device, dtype=torch.float32)
        ratio = ratio.reshape(-1)
        if ratio.numel() == 1:
            ratio = ratio.expand(tok.shape[0])
        n_patch_tokens = h_tok * w_tok
        n_mask = (ratio * n_patch_tokens).round().long()  # (B*V,)
        if bool((n_mask > 0).any()):
            scores = torch.rand(tok.shape[0], n_patch_tokens, device=tok.device)
            ranks = scores.argsort(dim=1).argsort(dim=1)
            hide = (ranks < n_mask[:, None]).view(tok.shape[0], h_tok, w_tok, 1)
            tok = torch.where(hide, mask_token.to(tok.dtype).view(1, 1, 1, -1), tok)
    py, px = int(pad_tokens[0]), int(pad_tokens[1])
    if (py or px) and pad_token is not None:
        fill = pad_token.to(tok.dtype).view(1, 1, 1, -1)
        tok = torch.nn.functional.pad(tok, (0, 0, px, px, py, py))
        keep = torch.zeros(
            tok.shape[1], tok.shape[2], dtype=torch.bool, device=tok.device
        )
        keep[py : py + h_tok, px : px + w_tok] = True
        tok = torch.where(keep[None, :, :, None], tok, fill)
        h_tok, w_tok = tok.shape[1], tok.shape[2]
    pos = model._pos_embed(tok)  # (B*V, P + N, dim), the P prefix tokens first
    if isinstance(pos, tuple):
        h, rope = pos
    else:
        h, rope = pos, None
    h = model.norm_pre(h)
    n_tok = h.shape[1]
    n_prefix = int(getattr(model, "num_prefix_tokens", 0))
    n_patch = n_tok - n_prefix
    mask = None
    if view_valid is not None and not bool(view_valid.all()):
        keep = view_valid.to(device=h.device, dtype=torch.bool)  # (B, V)
        # Key mask in the global layout, [V*P prefixes, V*N patches]: one per token.
        mask = torch.cat(
            [
                keep.repeat_interleave(n_prefix, dim=1),
                keep.repeat_interleave(n_patch, dim=1),
            ],
            dim=1,
        )[:, None, None, :]  # (B, 1, 1, T): broadcast over heads and queries
    rope_g = None if rope is None else rope.repeat(v, *([1] * (rope.ndim - 1)))
    is_global = False
    for block, scope in zip(model.blocks, attn_scopes, strict=True):
        want = scope == "g"
        if want != is_global:
            h = _to_global(h, b, v, n_prefix) if want else _to_local(h, b, v, n_prefix)
            is_global = want
        if not is_global:
            kw = {"rope": rope} if rope is not None else {}
            h = _run_block(block, h, kw, None, n_prefix, checkpoint)
            continue
        kw = {}
        if rope_g is not None and global_rope:
            kw["rope"] = rope_g
        if mask is not None:
            kw["attn_mask"] = mask
        prefix = n_prefix * v if "rope" in kw else None
        h = _run_block(block, h, kw, prefix, n_prefix, checkpoint)
    if is_global:
        h = _to_local(h, b, v, n_prefix)
    h = model.norm(h)[:, n_prefix:]
    return h.reshape(b * v, h_tok, w_tok, dim).permute(0, 3, 1, 2)

The heatmap heads: a 2D field per point, and its sub-pixel readouts.

concat (multi-scale laterals at stride 4, one 3x3 block) and unet (a coarse-to-fine decoder) are the monocular heads, read out by argmax_parabola; upsample_head is the probability-field head of mvt and fast, read out by dark_decode. The factories return the modules the assembler binds, so a checkpoint loads strict=True with no key translation.

argmax_parabola

argmax_parabola(heatmaps)

Hard peak and parabolic sub-pixel refinement, (..., 2) (x, y) in cells.

The second difference is clamped negative and the shift to half a cell, so a flat neighborhood gives no shift and the refinement never leaves the peak's cell.

Source code in src/deeperfly/pose2d/heads/heatmap.py

def argmax_parabola(heatmaps):
    """Hard peak and parabolic sub-pixel refinement, `(..., 2)` `(x, y)` in cells.

    The second difference is clamped negative and the shift to half a cell, so a flat
    neighborhood gives no shift and the refinement never leaves the peak's cell.
    """
    *lead, h, w = heatmaps.shape
    hm = heatmaps.reshape(-1, h, w)
    idx = hm.reshape(hm.shape[0], h * w).argmax(dim=1)
    cells = parabola_at(hm, (idx // w)[:, None], (idx % w)[:, None])
    return cells[:, 0].reshape(*lead, 2)

build_head

build_head(
    kind: str,
    channels: tuple[int, ...],
    n_keypoints: int,
    *,
    lat: int = 96,
    mid: int = 128,
)

(laterals, head, decoder) for head kind, to bind on the assembled detector.

kind is "concat" or "unet"; channels are the backbone's widths, finest first; lat and mid are the lateral and hidden widths. An empty piece is still a real module, so the state-dict keys are stable.

Source code in src/deeperfly/pose2d/heads/heatmap.py

def build_head(
    kind: str,
    channels: tuple[int, ...],
    n_keypoints: int,
    *,
    lat: int = 96,
    mid: int = 128,
):
    """`(laterals, head, decoder)` for head `kind`, to bind on the assembled detector.

    `kind` is `"concat"` or `"unet"`; `channels` are the backbone's widths, finest
    first; `lat` and `mid` are the lateral and hidden widths. An empty piece is still a
    real module, so the state-dict keys are stable.
    """
    if kind == "concat":
        laterals = nn.ModuleList([nn.Conv2d(c, lat, kernel_size=1) for c in channels])
        head = nn.Sequential(
            nn.Conv2d(lat * len(channels), mid, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(mid),
            nn.ReLU(inplace=True),
            nn.Conv2d(mid, n_keypoints, kernel_size=1),
        )
        return laterals, head, None
    if kind == "unet":
        return nn.ModuleList(), nn.Identity(), _UNetDecoder(channels, n_keypoints, mid)
    raise RuntimeError(f"unknown head {kind!r}; use 'concat' or 'unet'")

coarse_reach

coarse_reach(lat: int, reach: int) -> nn.Sequential

reach 3x3 convs (GELU between) on the coarsest lateral, for the margin's reach.

Empty for reach = 0, so the shipped concat head's parameters are unchanged.

Source code in src/deeperfly/pose2d/heads/heatmap.py

def coarse_reach(lat: int, reach: int) -> nn.Sequential:
    """`reach` 3x3 convs (GELU between) on the coarsest lateral, for the margin's reach.

    Empty for `reach = 0`, so the shipped concat head's parameters are unchanged.
    """
    layers = []
    for i in range(int(reach)):
        if i:
            layers.append(nn.GELU())
        layers.append(nn.Conv2d(lat, lat, kernel_size=3, padding=1))
    return nn.Sequential(*layers)

dark_at

dark_at(heatmaps, iy, ix)

dark_refine's step at M given cells per field: (N, M, 2) in cells.

Source code in src/deeperfly/pose2d/heads/heatmap.py

def dark_at(heatmaps, iy, ix):
    """`dark_refine`'s step at `M` given cells per field: `(N, M, 2)` in cells."""
    n, h, w = heatmaps.shape
    lg = heatmaps.float().clamp_min(1e-10).log()
    ar = torch.arange(n, device=lg.device)[:, None]

    def at(dy, dx):
        return lg[ar, (iy + dy).clamp(0, h - 1), (ix + dx).clamp(0, w - 1)]

    c = at(0, 0)
    gx = 0.5 * (at(0, 1) - at(0, -1))
    gy = 0.5 * (at(1, 0) - at(-1, 0))
    hxx = at(0, 1) - 2 * c + at(0, -1)
    hyy = at(1, 0) - 2 * c + at(-1, 0)
    hxy = 0.25 * (at(1, 1) - at(1, -1) - at(-1, 1) + at(-1, -1))
    det = hxx * hyy - hxy * hxy
    interior = (ix >= 1) & (ix <= w - 2) & (iy >= 1) & (iy <= h - 2)
    ok = interior & (det > 0) & (hxx < 0)
    det = torch.where(ok, det, torch.ones_like(det))
    zero = torch.zeros_like(gx)
    dx = torch.where(ok, -(hyy * gx - hxy * gy) / det, zero).clamp(-1.0, 1.0)
    dy = torch.where(ok, -(hxx * gy - hxy * gx) / det, zero).clamp(-1.0, 1.0)
    return torch.stack([ix + dx, iy + dy], dim=-1)

dark_decode

dark_decode(prob, *, geom: FieldGeometry)

A probability field -> input-normalized (..., 2) peaks and confidence.

dark_modulate then dark_refine, inverted through geom, and peak_mass for the confidence: the share of probability near the peak.

Source code in src/deeperfly/pose2d/heads/heatmap.py

def dark_decode(prob, *, geom: FieldGeometry):
    """A probability field -> input-normalized `(..., 2)` peaks and confidence.

    `dark_modulate` then `dark_refine`, inverted through `geom`, and `peak_mass` for the
    confidence: the share of probability near the peak.
    """
    prob = prob.float()
    cells = dark_refine(dark_modulate(prob))
    return geom.to_input_normalized(cells), peak_mass(prob)

dark_modulate

dark_modulate(
    field,
    *,
    sigma: float = DARK_SIGMA,
    kernel: int = DARK_KERNEL,
)

Gaussian-smooth each channel of a (..., H, W) field: DARK's modulation.

The training kernel restores the single Gaussian shape dark_refine's Taylor step assumes; the reference's rescale to the original peak is skipped, since the log-derivatives are scale-free.

Source code in src/deeperfly/pose2d/heads/heatmap.py

def dark_modulate(field, *, sigma: float = DARK_SIGMA, kernel: int = DARK_KERNEL):
    """Gaussian-smooth each channel of a `(..., H, W)` field: DARK's modulation.

    The training kernel restores the single Gaussian shape `dark_refine`'s Taylor step
    assumes; the reference's rescale to the original peak is skipped, since the
    log-derivatives are scale-free.
    """
    if kernel % 2 == 0:
        raise ValueError(f"kernel must be odd, got {kernel}")
    *lead, h, w = field.shape
    x = torch.arange(kernel, device=field.device, dtype=torch.float32)
    g = torch.exp(-((x - (kernel - 1) / 2) ** 2) / (2 * sigma * sigma))
    g = g / g.sum()
    flat = field.reshape(1, -1, h, w).float()
    n = flat.shape[1]
    pad = kernel // 2
    # In fp32 under any autocast: the refinement takes log-derivatives of this, and a
    # bf16 blur is the slow path here. Separable and depthwise, with zero padding,
    # which biases a blob cut by the border least (0.2 cells two cells from the edge,
    # against 0.65 to 2 for reflect, replicate or normalized). A device autocast does
    # not know (a `meta` trace) has no state to disable, and asking raises.
    dev = field.device.type
    off = (
        torch.autocast(dev, enabled=False)
        if torch.amp.is_autocast_available(dev)
        else contextlib.nullcontext()
    )
    with off:
        ky = g.view(1, 1, kernel, 1).expand(n, 1, kernel, 1)
        kx = g.view(1, 1, 1, kernel).expand(n, 1, 1, kernel)
        out = F.conv2d(flat, ky, padding=(pad, 0), groups=n)
        out = F.conv2d(out, kx, padding=(0, pad), groups=n)
    return out.reshape(*lead, h, w)

dark_refine

dark_refine(heatmaps)

Hard peak and one Newton step on the log-heatmap: (..., 2) (x, y) in cells.

DARK (Zhang et al., CVPR 2020): the offset to the mode is -H^-1 g from 3x3 finite differences of the log-heatmap, exact for a Gaussian blob. A border peak, or one whose Hessian is not negative definite, keeps its cell, and the step is clamped to one cell.

Source code in src/deeperfly/pose2d/heads/heatmap.py

def dark_refine(heatmaps):
    """Hard peak and one Newton step on the log-heatmap: `(..., 2)` `(x, y)` in cells.

    DARK (Zhang et al., CVPR 2020): the offset to the mode is `-H^-1 g` from 3x3 finite
    differences of the log-heatmap, exact for a Gaussian blob. A border peak, or one
    whose Hessian is not negative definite, keeps its cell, and the step is clamped to
    one cell.
    """
    *lead, h, w = heatmaps.shape
    hm = heatmaps.reshape(-1, h, w)
    idx = hm.reshape(hm.shape[0], h * w).argmax(dim=1)
    cells = dark_at(hm, (idx // w)[:, None], (idx % w)[:, None])
    return cells[:, 0].reshape(*lead, 2)

forward_field

forward_field(
    laterals, head, decoder, feats, pad_cells, reach=None
)

Backbone features feats (finest first) -> the padded heatmap field.

The margin (pad_cells, (py, px)) is zero-padded: nothing is there, and reflecting would paste a mirrored fly where an off-image joint belongs. Padded on the fused stride-4 features it sees one cell of content; with reach the coarsest lateral is padded at its own stride, run through the reach convs and upsampled, so the whole margin is extrapolated from coarse features.

Source code in src/deeperfly/pose2d/heads/heatmap.py

def forward_field(laterals, head, decoder, feats, pad_cells, reach=None):
    """Backbone features `feats` (finest first) -> the padded heatmap field.

    The margin (`pad_cells`, `(py, px)`) is zero-padded: nothing is there, and
    reflecting would paste a mirrored fly where an off-image joint belongs. Padded on
    the fused stride-4 features it sees one cell of content; with `reach` the coarsest
    lateral is padded at its own stride, run through the reach convs and upsampled, so
    the whole margin is extrapolated from coarse features.
    """
    py, px = pad_cells
    if decoder is not None:
        return decoder(feats, (py, px))
    h, w = feats[0].shape[-2:]  # stride-4 == the UNPADDED heatmap grid
    red = [lateral(f) for lateral, f in zip(laterals, feats, strict=True)]
    if (py or px) and reach is not None and len(reach):
        coarse = red[-1]
        s = h // coarse.shape[-2]  # stride-4 cells per coarse cell
        cpy, cpx = -(-py // s), -(-px // s)  # whole coarse cells covering the pad
        c = reach(F.pad(coarse, (cpx, cpx, cpy, cpy)))
        up = F.interpolate(c, scale_factor=s, mode="bilinear", align_corners=False)
        oy, ox = cpy * s - py, cpx * s - px
        up = up[..., oy : oy + h + 2 * py, ox : ox + w + 2 * px]
        fine = [red[0]] + [
            F.interpolate(r, size=(h, w), mode="bilinear", align_corners=False)
            for r in red[1:-1]
        ]
        feat = F.pad(torch.cat(fine, dim=1), (px, px, py, py))
        return head(torch.cat([feat, up], dim=1))
    ups = [red[0]] + [
        F.interpolate(r, size=(h, w), mode="bilinear", align_corners=False)
        for r in red[1:]
    ]
    feat = torch.cat(ups, dim=1)
    if py or px:
        feat = F.pad(feat, (px, px, py, py))
    return head(feat)

masked_heatmap_loss

masked_heatmap_loss(
    pred,
    kp_cells,
    weight,
    *,
    sigma: float = 2.0,
    fg_weight: float = 30.0,
)

Weight-masked, foreground-weighted Gaussian-heatmap MSE.

Each pixel is weighted 1 + fg_weight * target, since a plain MSE on sparse targets is minimized by predicting background. weight (B, K) is continuous, and 0 removes a channel from numerator and denominator alike; kp_cells may be NaN there. Returns (loss, aux) with heatmap_mse and weight_sum.

Source code in src/deeperfly/pose2d/heads/heatmap.py

def masked_heatmap_loss(
    pred, kp_cells, weight, *, sigma: float = 2.0, fg_weight: float = 30.0
):
    """Weight-masked, foreground-weighted Gaussian-heatmap MSE.

    Each pixel is weighted `1 + fg_weight * target`, since a plain MSE on sparse targets
    is minimized by predicting background. `weight` `(B, K)` is continuous, and 0
    removes a channel from numerator and denominator alike; `kp_cells` may be NaN there.
    Returns `(loss, aux)` with `heatmap_mse` and `weight_sum`.
    """
    if pred.ndim != 4:
        raise ValueError(f"pred must be (B, K, H, W), got {tuple(pred.shape)}")
    _, _, h, w = pred.shape
    with torch.no_grad():
        ones = torch.ones(kp_cells.shape[:-1], device=pred.device, dtype=pred.dtype)
        target = render_gaussian_targets(kp_cells, ones, hw=(h, w), sigma=sigma).to(
            pred.dtype
        )
        fg = 1.0 + fg_weight * target
    per_kp = (fg * (pred - target) ** 2).mean(dim=(-1, -2))  # (B, K)
    w_sum = weight.sum()
    # clamp_min guards the all-unlabeled batch (loss 0, finite, zero grad).
    loss = (per_kp * weight.to(per_kp.dtype)).sum() / w_sum.clamp_min(1.0)
    return loss, {"heatmap_mse": loss.detach(), "weight_sum": w_sum.detach()}

parabola_at

parabola_at(heatmaps, iy, ix)

argmax_parabola's refinement at M given cells per field, (N, M, 2).

For the candidate path; iy and ix are integer (N, M) cell indices.

Source code in src/deeperfly/pose2d/heads/heatmap.py

def parabola_at(heatmaps, iy, ix):
    """`argmax_parabola`'s refinement at `M` given cells per field, `(N, M, 2)`.

    For the candidate path; `iy` and `ix` are integer `(N, M)` cell indices.
    """
    n, h, w = heatmaps.shape
    ar = torch.arange(n, device=heatmaps.device)[:, None]

    def gather(yy, xx):
        return heatmaps[ar, yy.clamp(0, h - 1), xx.clamp(0, w - 1)]

    c = gather(iy, ix)
    xl, xr = gather(iy, ix - 1), gather(iy, ix + 1)
    yu, yd = gather(iy - 1, ix), gather(iy + 1, ix)
    dx = 0.5 * (xl - xr) / (xl - 2 * c + xr).clamp(max=-1e-6)
    dy = 0.5 * (yu - yd) / (yu - 2 * c + yd).clamp(max=-1e-6)
    x = ix.to(heatmaps.dtype) + dx.clamp(-0.5, 0.5)
    y = iy.to(heatmaps.dtype) + dy.clamp(-0.5, 0.5)
    return torch.stack([x, y], dim=-1)

render_gaussian_targets

render_gaussian_targets(
    kp_cells,
    vis,
    *,
    hw: tuple[int, int],
    sigma: float = 2.0,
)

Peak-vis separable Gaussian targets from keypoints in field cells.

Rendered at the fractional cell, since rounding would bake in a half-cell bias. kp_cells (..., K, 2) may be non-finite where unlabeled, and vis (..., K) scales the amplitude, 0 giving an exactly zero channel. Returns (..., K, height, width).

Source code in src/deeperfly/pose2d/heads/heatmap.py

def render_gaussian_targets(kp_cells, vis, *, hw: tuple[int, int], sigma: float = 2.0):
    """Peak-`vis` separable Gaussian targets from keypoints in field cells.

    Rendered at the fractional cell, since rounding would bake in a half-cell bias.
    `kp_cells` `(..., K, 2)` may be non-finite where unlabeled, and `vis` `(..., K)`
    scales the amplitude, 0 giving an exactly zero channel. Returns
    `(..., K, height, width)`.
    """
    if kp_cells.shape[-1] != 2:
        raise ValueError(
            f"kp_cells must end in a size-2 axis, got {tuple(kp_cells.shape)}"
        )
    if tuple(vis.shape) != tuple(kp_cells.shape[:-1]):
        raise ValueError(
            f"vis {tuple(vis.shape)} must match kp_cells {tuple(kp_cells.shape[:-1])}"
        )
    height, width = hw
    device, dtype = kp_cells.device, kp_cells.dtype
    # NaN * 0 is NaN, so masking afterwards is not enough: neutralize before arithmetic.
    safe = torch.nan_to_num(kp_cells, nan=0.0, posinf=0.0, neginf=0.0)
    x = safe[..., 0:1].unsqueeze(-1)
    y = safe[..., 1:2].unsqueeze(-1)
    gx = torch.arange(width, device=device, dtype=dtype).view(1, width)
    gy = torch.arange(height, device=device, dtype=dtype).view(height, 1)
    gx_term = torch.exp(-((gx - x) ** 2) / (2 * sigma**2))
    gy_term = torch.exp(-((gy - y) ** 2) / (2 * sigma**2))
    return gx_term * gy_term * vis.to(dtype).unsqueeze(-1).unsqueeze(-1)

soft_argmax

soft_argmax(
    field, *, input_hw, downsample_factor: int, margin=0
)

(B, C, Hm, Wm) field -> input-normalized (B, C, 2) peaks and confidence.

Upsampled to the input (downsample_factor 2x steps), softmaxed sharp and reduced by expectation, with the padding inverted, so a coordinate outside [0, 1] is off the frame. The reference readout: its upsample costs as much as the forward, for a mode DARK reaches from a 3x3 neighborhood. margin is a scalar or a (my, mx) pair of off-frame input pixels.

Source code in src/deeperfly/pose2d/heads/heatmap.py

def soft_argmax(field, *, input_hw, downsample_factor: int, margin=0):
    """`(B, C, Hm, Wm)` field -> input-normalized `(B, C, 2)` peaks and confidence.

    Upsampled to the input (`downsample_factor` 2x steps), softmaxed sharp and reduced
    by expectation, with the padding inverted, so a coordinate outside `[0, 1]` is off
    the frame. The reference readout: its upsample costs as much as the forward, for a
    mode DARK reaches from a 3x3 neighborhood. `margin` is a scalar or a `(my, mx)` pair
    of off-frame input pixels.
    """
    if downsample_factor not in GRID_OFFSETS:
        raise SystemExit(
            f"downsample_factor {downsample_factor} has no grid-offset correction; "
            f"known: {sorted(GRID_OFFSETS)}"
        )
    my, mx = (margin, margin) if isinstance(margin, (int, float)) else margin
    my, mx = int(my), int(mx)
    if my < 0 or mx < 0:
        raise ValueError(f"margin must be >= 0, got {margin}")
    h_rep, w_rep = int(input_hw[0]), int(input_hw[1])
    h_in, w_in = h_rep + 2 * my, w_rep + 2 * mx
    off = GRID_OFFSETS[downsample_factor]
    geom = FieldGeometry((h_rep, w_rep), stride=1, origin_xy=(off + mx, off + my))
    xs, cs = [], []
    for lo in range(0, field.shape[1], DECODE_CHUNK):
        hm = field[:, lo : lo + DECODE_CHUNK].float()
        for _ in range(downsample_factor):
            hm = _pyr_up(hm)
        if hm.shape[-2:] != (h_in, w_in):
            raise RuntimeError(
                f"the upsampled field is {tuple(hm.shape[-2:])}, not the "
                f"{(h_in, w_in)} this reported frame {(h_rep, w_rep)} plus margin "
                f"{margin} describes"
            )
        prob = spatial_softmax(hm, SOFTARGMAX_TEMPERATURE)
        pts = _spatial_expectation(prob)
        # Confidence is read at the raw expectation (the cell the peak occupies), and
        # the geometry then applies the grid offset and margin shift.
        cs.append(softmax_window_mass(prob, pts))
        xs.append(pts)
    return geom.to_input_normalized(torch.cat(xs, dim=1)), torch.cat(cs, dim=1)

spatial_softmax

spatial_softmax(x, temperature: float)

Softmax over the spatial axes of (B, C, H, W), scaled by temperature.

Source code in src/deeperfly/pose2d/heads/heatmap.py

def spatial_softmax(x, temperature: float):
    """Softmax over the spatial axes of `(B, C, H, W)`, scaled by `temperature`."""
    b, c, h, w = x.shape
    return F.softmax(x.reshape(b, c, -1) * temperature, dim=-1).view(b, c, h, w)

upsample_head

upsample_head(
    dim: int,
    n_keypoints: int,
    *,
    final_softmax: bool = True,
    pad_tokens: tuple[int, int] = (0, 0),
    reach: int = 0,
    kind: str = "upsample",
    width: int | None = None,
) -> nn.Module

A token-to-field head: PixelShuffle and a transposed conv, or ViTPose's deconvs.

The field is the token grid times 4, a temperature-1 spatial softmax with final_softmax. pad_tokens (py, px) are zero-padded around the tokens before the upsampling, so the field extends past the frame, and reach adds that many 3x3 token convs after the pad, each letting the margin see one more token of content. kind is "upsample" (the shipped head) or "deconv" (two stride-2 transposed convs at DECONV_WIDTH). width adds a 1x1 projection before the pad, for a conv backbone 512 wide at stride 16, where one reach conv costs 2.26 GMAC a view.

Source code in src/deeperfly/pose2d/heads/heatmap.py

def upsample_head(
    dim: int,
    n_keypoints: int,
    *,
    final_softmax: bool = True,
    pad_tokens: tuple[int, int] = (0, 0),
    reach: int = 0,
    kind: str = "upsample",
    width: int | None = None,
) -> nn.Module:
    """A token-to-field head: PixelShuffle and a transposed conv, or ViTPose's deconvs.

    The field is the token grid times 4, a temperature-1 spatial softmax with
    `final_softmax`. `pad_tokens` `(py, px)` are zero-padded around the tokens before
    the upsampling, so the field extends past the frame, and `reach` adds that many 3x3
    token convs after the pad, each letting the margin see one more token of content.
    `kind` is `"upsample"` (the shipped head) or `"deconv"` (two stride-2 transposed
    convs at `DECONV_WIDTH`). `width` adds a 1x1 projection before the pad, for a conv
    backbone 512 wide at stride 16, where one reach conv costs 2.26 GMAC a view.
    """
    if kind not in ("upsample", "deconv"):
        raise RuntimeError(f"unknown mvt head {kind!r}; use 'upsample' or 'deconv'")
    w = int(width) if width else int(dim)
    if kind == "upsample" and w % 4:
        raise RuntimeError(
            f"head width {w} is not divisible by 4 (the PixelShuffle(2) quarters it)"
        )
    py, px = int(pad_tokens[0]), int(pad_tokens[1])

    class _UpsampleHead(nn.Module):
        def __init__(self) -> None:
            super().__init__()
            self.pad_tokens = (py, px)
            self.project = nn.Conv2d(dim, w, kernel_size=1) if width else nn.Identity()
            layers = []
            for i in range(int(reach)):
                if i:
                    layers.append(nn.GELU())
                layers.append(nn.Conv2d(w, w, kernel_size=3, padding=1))
            self.reach = nn.Sequential(*layers)
            if kind == "upsample":
                self.upsampling_layers = nn.Sequential(
                    nn.PixelShuffle(2),
                    nn.ConvTranspose2d(
                        w // 4,
                        n_keypoints,
                        kernel_size=(3, 3),
                        stride=(2, 2),
                        padding=(1, 1),
                        output_padding=(1, 1),
                    ),
                )
            else:
                d = DECONV_WIDTH
                self.upsampling_layers = nn.Sequential(
                    nn.ConvTranspose2d(w, d, 4, stride=2, padding=1, bias=False),
                    nn.BatchNorm2d(d),
                    nn.GELU(),
                    nn.ConvTranspose2d(d, d, 4, stride=2, padding=1, bias=False),
                    nn.BatchNorm2d(d),
                    nn.GELU(),
                    nn.Conv2d(d, n_keypoints, kernel_size=1),
                )

        def forward(self, x):
            x = self.project(x)
            if py or px:
                x = F.pad(x, (px, px, py, py))
            hm = self.upsampling_layers(self.reach(x))
            return spatial_softmax(hm, 1.0) if final_softmax else hm

    return _UpsampleHead()

The SimCC head: per-axis 1D coordinate classification (Li et al., ECCV 2022).

Two 1D distributions per point, over x and y bins at split_ratio times the input resolution, each read out by soft-argmax: sub-pixel resolution with no 2D field, and a differentiable readout. No checkpoint has shipped, so the key layout is ours, and predict_heatmaps refuses rather than fabricate a field.

bins_for

bins_for(
    input_hw: tuple[int, int],
    split_ratio: float,
    margin: int = 0,
) -> tuple[int, int]

(y_bins, x_bins) at split_ratio bins a pixel, with margin px each side.

With a margin the bins span [-margin, size + margin], so an off-frame coordinate has a bin.

Source code in src/deeperfly/pose2d/heads/simcc.py

def bins_for(
    input_hw: tuple[int, int], split_ratio: float, margin: int = 0
) -> tuple[int, int]:
    """`(y_bins, x_bins)` at `split_ratio` bins a pixel, with `margin` px each side.

    With a margin the bins span `[-margin, size + margin]`, so an off-frame coordinate
    has a bin.
    """
    return (
        round((input_hw[0] + 2 * margin) * split_ratio),
        round((input_hw[1] + 2 * margin) * split_ratio),
    )

simcc_decode

simcc_decode(
    x_logits,
    y_logits,
    *,
    split_ratio: float,
    input_hw: tuple[int, int],
    margin: int = 0,
)

Bin logits -> input-normalized (N, K, 2) peaks and (N, K) confidence.

Each axis is softmaxed and reduced by expectation; bin b is input pixel b / split_ratio - margin. The confidence is the geometric mean of the two axes' peak probabilities.

Source code in src/deeperfly/pose2d/heads/simcc.py

def simcc_decode(
    x_logits,
    y_logits,
    *,
    split_ratio: float,
    input_hw: tuple[int, int],
    margin: int = 0,
):
    """Bin logits -> input-normalized `(N, K, 2)` peaks and `(N, K)` confidence.

    Each axis is softmaxed and reduced by expectation; bin `b` is input pixel
    `b / split_ratio - margin`. The confidence is the geometric mean of the two axes'
    peak probabilities.
    """

    def axis(logits, size):
        prob = F.softmax(logits, dim=-1)
        idx = torch.arange(logits.shape[-1], dtype=prob.dtype, device=prob.device)
        ex = (prob * idx).sum(-1)
        return (ex / split_ratio - margin) / size, prob.amax(dim=-1)

    xn, cx = axis(x_logits, input_hw[1])
    yn, cy = axis(y_logits, input_hw[0])
    return torch.stack([xn, yn], dim=-1), torch.sqrt(cx * cy)

simcc_head

simcc_head(
    in_channels: int,
    n_keypoints: int,
    feat_hw: tuple[int, int],
    bins: tuple[int, int],
) -> nn.Module

A head mapping the finest feature map to per-axis bin logits.

(N, in_channels, Hf, Wf) -> ((N, K, x_bins), (N, K, y_bins)), with feat_hw the feature map's size in cells and bins (y_bins, x_bins).

Source code in src/deeperfly/pose2d/heads/simcc.py

def simcc_head(
    in_channels: int,
    n_keypoints: int,
    feat_hw: tuple[int, int],
    bins: tuple[int, int],
) -> nn.Module:
    """A head mapping the finest feature map to per-axis bin logits.

    `(N, in_channels, Hf, Wf) -> ((N, K, x_bins), (N, K, y_bins))`, with `feat_hw` the
    feature map's size in cells and `bins` `(y_bins, x_bins)`.
    """
    hf, wf = feat_hw
    hy, wx = bins

    class _SimCC(nn.Module):
        def __init__(self) -> None:
            super().__init__()
            self.stem = nn.Conv2d(in_channels, n_keypoints, kernel_size=1)
            self.mlp_x = nn.Linear(wf, wx)
            self.mlp_y = nn.Linear(hf, hy)

        def forward(self, feat):
            f = self.stem(feat)  # (N, K, Hf, Wf)
            # Marginalize each axis, then project to that axis' bins.
            return self.mlp_x(f.mean(dim=2)), self.mlp_y(f.mean(dim=3))

    return _SimCC()

simcc_kl_loss

simcc_kl_loss(
    x_logits, y_logits, x_target, y_target, weight
)

Weight-masked KL between the target bins and the logits' softmax, per axis.

weight == 0 drops a point from numerator and denominator. Returns (loss, aux).

Source code in src/deeperfly/pose2d/heads/simcc.py

def simcc_kl_loss(x_logits, y_logits, x_target, y_target, weight):
    """Weight-masked KL between the target bins and the logits' softmax, per axis.

    `weight == 0` drops a point from numerator and denominator. Returns `(loss, aux)`.
    """

    def kl(logits, target):
        return F.kl_div(F.log_softmax(logits, dim=-1), target, reduction="none").sum(-1)

    per = kl(x_logits, x_target) + kl(y_logits, y_target)  # (..., K)
    w_sum = weight.sum()
    loss = (per * weight.to(per.dtype)).sum() / w_sum.clamp_min(1.0)
    return loss, {"simcc_kl": loss.detach(), "weight_sum": w_sum.detach()}

simcc_peaks

simcc_peaks(
    x_logits,
    y_logits,
    k: int,
    *,
    split_ratio: float,
    input_hw: tuple[int, int],
    margin: int = 0,
    radius: int = 1,
)

The k strongest distinct peak pairs, (N, K, k, 2) and (N, K, k) scores.

The joint of two independent marginals is their outer product, whose local maxima are pairs of 1D maxima: each axis's maxima over 2 radius + 1 bins, refined by a parabola, paired and ranked by the geometric mean of their peak probabilities. Missing slots are NaN with score 0.

Source code in src/deeperfly/pose2d/heads/simcc.py

def simcc_peaks(
    x_logits,
    y_logits,
    k: int,
    *,
    split_ratio: float,
    input_hw: tuple[int, int],
    margin: int = 0,
    radius: int = 1,
):
    """The `k` strongest distinct peak pairs, `(N, K, k, 2)` and `(N, K, k)` scores.

    The joint of two independent marginals is their outer product, whose local maxima
    are pairs of 1D maxima: each axis's maxima over `2 radius + 1` bins, refined by a
    parabola, paired and ranked by the geometric mean of their peak probabilities.
    Missing slots are NaN with score 0.
    """

    def axis(logits, size):
        prob = F.softmax(logits.float(), dim=-1)
        n, kk, b = prob.shape
        flat = prob.reshape(n * kk, 1, b)
        highest = F.max_pool1d(flat, 2 * radius + 1, stride=1, padding=radius)
        lowest = -F.max_pool1d(-flat, 2 * radius + 1, stride=1, padding=radius)
        neg = torch.full_like(flat, float("-inf"))
        peaks = torch.where((flat == highest) & (flat > lowest), flat, neg)[:, 0]
        m = min(k, b)
        val, idx = peaks.topk(m, dim=-1)
        rows = torch.arange(n * kk, device=prob.device)[:, None]
        line = flat[:, 0]
        c = line[rows, idx]
        left = line[rows, (idx - 1).clamp(0, b - 1)]
        right = line[rows, (idx + 1).clamp(0, b - 1)]
        shift = 0.5 * (left - right) / (left - 2 * c + right).clamp(max=-1e-6)
        pos = idx.to(prob.dtype) + shift.clamp(-0.5, 0.5)
        coord = (pos / split_ratio - margin) / size
        valid = torch.isfinite(val)
        coord = torch.where(valid, coord, torch.full_like(coord, float("nan")))
        val = torch.where(valid, val, torch.zeros_like(val))
        return coord.reshape(n, kk, m), val.reshape(n, kk, m)

    xn, px = axis(x_logits, input_hw[1])
    yn, py = axis(y_logits, input_hw[0])
    n, kk, mx = xn.shape
    my = yn.shape[-1]
    score = torch.sqrt(px[..., :, None] * py[..., None, :]).reshape(n, kk, mx * my)
    xy = torch.stack(
        [
            xn[..., :, None].expand(n, kk, mx, my),
            yn[..., None, :].expand(n, kk, mx, my),
        ],
        dim=-1,
    ).reshape(n, kk, mx * my, 2)
    m = min(k, mx * my)
    top, order = score.topk(m, dim=-1)
    xy = xy.gather(2, order[..., None].expand(n, kk, m, 2))
    valid = top > 0
    xy = torch.where(valid[..., None], xy, torch.full_like(xy, float("nan")))
    top = torch.where(valid, top, torch.zeros_like(top))
    if m < k:
        pad_xy = torch.full((n, kk, k - m, 2), float("nan"), device=xy.device)
        xy = torch.cat([xy, pad_xy.to(xy.dtype)], dim=2)
        top = torch.cat([top, torch.zeros((n, kk, k - m), device=top.device)], dim=2)
    return xy, top

simcc_targets

simcc_targets(
    coords_norm,
    *,
    input_hw: tuple[int, int],
    split_ratio: float,
    sigma: float = 6.0,
    margin: int = 0,
)

Per-axis 1D Gaussian bin targets, each summing to 1, from normalized coords.

A label at input-normalized x centers on bin (x * W + margin) * split_ratio, with margin matching the decode's; a non-finite label goes to bin 0, masked by weight in the loss. sigma is in bins.

Source code in src/deeperfly/pose2d/heads/simcc.py

def simcc_targets(
    coords_norm,
    *,
    input_hw: tuple[int, int],
    split_ratio: float,
    sigma: float = 6.0,
    margin: int = 0,
):
    """Per-axis 1D Gaussian bin targets, each summing to 1, from normalized coords.

    A label at input-normalized `x` centers on bin `(x * W + margin) * split_ratio`,
    with `margin` matching the decode's; a non-finite label goes to bin 0, masked by
    weight in the loss. `sigma` is in bins.
    """
    hy, wx = bins_for(input_hw, split_ratio, margin)
    safe = torch.nan_to_num(coords_norm, nan=0.0, posinf=0.0, neginf=0.0)
    bx = (safe[..., 0] * input_hw[1] + margin) * split_ratio
    by = (safe[..., 1] * input_hw[0] + margin) * split_ratio

    def render(center, nbins):
        idx = torch.arange(nbins, device=center.device, dtype=center.dtype)
        g = torch.exp(-((idx - center.unsqueeze(-1)) ** 2) / (2 * sigma**2))
        return g / g.sum(-1, keepdim=True).clamp_min(1e-8)

    return render(bx, wx), render(by, hy)

The assembler: compose a backbone and a head into a Detector.

Four families share one Detector and one decode and predict surface: monocular (a CNN pyramid and a heatmap head, one image at a time, the deeperfly-dense38-1 checkpoint layout, a parabolic readout), simcc (the same pyramid and a per-axis coordinate-classification head, no 2D field), mvt (a multiview ViT over the V synchronized views of a frame, a probability field read out by DARK) and fast (the mvt's tensor contract and head over a per-view CNN pyramid). A monocular checkpoint's field margin, readout and normalization are silent wrong answers if they disagree with training, so each is recorded in the artifact or refused.

Detector

Bases: Module

A backbone and a head, one of the four families.

monocular binds backbone, laterals, head and decoder (the names a deeperfly-dense38-1 checkpoint expects) and returns a one-element list; simcc returns the (x_logits, y_logits) bin pair; mvt and fast return (B, V * K, Hm, Wm), view c // K and keypoint c % K. K is num_classes, whose last channel is the animal-center field with center_channel.

Source code in src/deeperfly/pose2d/detect.py

class Detector(nn.Module):
    """A backbone and a head, one of the four families.

    `monocular` binds `backbone`, `laterals`, `head` and `decoder` (the names a
    `deeperfly-dense38-1` checkpoint expects) and returns a one-element list; `simcc`
    returns the `(x_logits, y_logits)` bin pair; `mvt` and `fast` return
    `(B, V * K, Hm, Wm)`, view `c // K` and keypoint `c % K`. `K` is `num_classes`,
    whose last channel is the animal-center field with `center_channel`.
    """

    def __init__(
        self,
        family: str,
        backbone,
        n_keypoints: int,
        *,
        backbone_name: str | None = None,
        sel=None,
        channels=None,
        head_kind: str = "concat",
        readout: str = "argmax",
        geom: FieldGeometry | None = None,
        margin: float = 0.0,
        attn_scopes: str | None = None,
        dim: int | None = None,
        input_hw=None,
        downsample_factor: int | None = None,
        margin_px: int = 0,
        pad_tokens: tuple[int, int] = (0, 0),
        final_softmax: bool = True,
        reach: int = 0,
        split_ratio: float = 2.0,
        token_mask: float = 0.0,
        token_pad: int = 0,
        global_rope: bool = True,
        level: int = 0,
        feature_stride: int = 0,
        width: int | None = None,
    ):
        """Bind the backbone and head for `family` (see the `build_*` assemblers)."""
        super().__init__()
        self.family = family
        self.backbone = backbone
        self.backbone_name = backbone_name
        self.num_classes = int(n_keypoints)
        #: Whether the last channel is an animal-center field; set after assembly, so a
        #: state dict never depends on it.
        self.center_channel = False
        self.point_names: list[str] = []
        #: Both families normalize their own grayscale input; a spec declares mean 0.0.
        self.norm_mean = 0.0
        self.norm_std = 1.0
        self.accepts_gray = True
        if family == "monocular":
            self._sel = tuple(sel)
            self._channels = tuple(channels)
            self.head_kind = head_kind
            self.laterals, self.head, self.decoder = heatmap.build_head(
                head_kind, channels, n_keypoints
            )
            # Empty for the shipped heads; `reach` convs on the coarsest lateral
            # otherwise.
            lat = self.laterals[-1].out_channels if len(self.laterals) else 0
            self.coarse_reach = heatmap.coarse_reach(lat, reach)
            self.reach = int(reach)
            self.field_is_prob = False  # the field is raw logits
            self.geom = geom
            self.input_hw = geom.input_hw
            self.stride = geom.stride
            self.pad_px = geom.pad_px
            self.field_hw = geom.field_hw
            #: The margin as a fraction of the input, or `None` when built from `pad_px`
            #: (a pixel policy has no single fraction on a non-square input).
            self.margin = None if margin is None else float(margin)
            self.readout = readout
            # The field is stride-`s`, so the soft-argmax readout recovers the input in
            # log2(s) upsample steps and inverts the per-axis `(py, px)` pad as margin.
            self.downsample_factor = int(geom.stride).bit_length() - 1
            self.hm_margin_px = geom.pad_px
            if readout == "argmax":
                self.output_kind = "heatmap"
                self.confidence_kind = "peak_amax"
            elif readout == "softargmax":
                self.output_kind = "heatmap_softargmax"
                self.confidence_kind = "peak_mass"
            else:
                raise ValueError(
                    f"unknown monocular readout {readout!r}; "
                    "use 'argmax' or 'softargmax'"
                )
            self.padded_field = True
            self.joint_views = False
            self.owns_prepare = False
            self.has_heatmaps = True
        elif family == "simcc":
            self._sel = tuple(sel)
            self._channels = tuple(channels)
            self.input_hw = tuple(int(v) for v in input_hw)
            self.split_ratio = float(split_ratio)
            self.hm_margin_px = int(margin_px)
            # The backbone sees the unpadded frame; the margin only widens the bin
            # range, so an off-frame coordinate lands on a real bin, not clipped.
            feat_hw = (self.input_hw[0] // STRIDE, self.input_hw[1] // STRIDE)
            self.head = simcc.simcc_head(
                channels[0],
                n_keypoints,
                feat_hw,
                simcc.bins_for(self.input_hw, split_ratio, self.hm_margin_px),
            )
            self.output_kind = "simcc"
            self.confidence_kind = "simcc_axis"
            self.joint_views = False
            self.owns_prepare = False
            # No 2D field to pad; off-frame coordinates come from the widened bin range.
            self.padded_field = False
            self.has_heatmaps = False
        elif family == "mvt":
            #: Where the off-frame margin's cells come from: the innermost `token_pad`
            #: tokens per side are padded into the token grid before the transformer
            #: (learned pad token, attention-capable), the rest zero-padded in the head.
            self.token_pad = int(token_pad)
            if self.token_pad > min(int(pad_tokens[0]), int(pad_tokens[1])):
                raise RuntimeError(
                    f"token_pad {token_pad} exceeds the margin's {pad_tokens} tokens"
                )
            self.pad_grid = (self.token_pad, self.token_pad)
            self.pad_token = nn.Parameter(torch.zeros(int(dim)))
            self.head = heatmap.upsample_head(
                dim,
                n_keypoints,
                final_softmax=final_softmax,
                pad_tokens=(
                    int(pad_tokens[0]) - self.token_pad,
                    int(pad_tokens[1]) - self.token_pad,
                ),
                reach=reach,
                kind=head_kind,
            )
            self.head_kind = head_kind
            self.reach = int(reach)
            #: Rope in the global blocks (see `vit_multiview.encode`); `False` is the
            #: content-only cross-view attention ablation.
            self.global_rope = bool(global_rope)
            #: Recompute the blocks in the backward pass. Set by the training run, not
            #: by the `[model]` table the artifact records: it changes no number.
            self.grad_checkpoint = False
            #: With `final_softmax` the head's field is a per-channel probability map.
            self.field_is_prob = bool(final_softmax)
            self.attn_scopes = list(attn_scopes)
            self.dim = int(dim)
            #: Training-time patch masking (see `vit_multiview.encode`); 0 masks
            #: nothing. The token is always a parameter, so the state dict is fixed.
            self.token_mask = float(token_mask)
            self.mask_token = nn.Parameter(torch.zeros(self.dim))
            self.input_hw = tuple(int(v) for v in input_hw)
            self.downsample_factor = int(downsample_factor)
            self.hm_margin_px = int(margin_px)
            #: The stride-`2 ** downsample_factor` cell grid the head's field lives on,
            #: padded by the margin on every side: what the readout inverts.
            self.geom = FieldGeometry.padded_px(
                self.input_hw,
                pad_px=(self.hm_margin_px, self.hm_margin_px),
                stride=2**self.downsample_factor,
            )
            self.output_kind = "heatmap_softargmax"
            self.confidence_kind = "peak_mass"
            self.joint_views = True
            self.padded_field = int(margin_px) > 0
            self.owns_prepare = True
            self.has_heatmaps = True
        elif family == "fast":
            self._sel = tuple(sel)
            self._channels = tuple(channels)
            #: Which pyramid level the head reads, and its stride in input pixels.
            self._level = int(level)
            self.feature_stride = int(feature_stride)
            self.dim = int(dim)
            self.width = None if width is None else int(width)
            self.head = heatmap.upsample_head(
                self.dim,
                n_keypoints,
                final_softmax=final_softmax,
                # The whole margin is padded in the head: a conv has no attention to
                # read into the pad with, so there is no inner band as in the mvt.
                pad_tokens=(int(pad_tokens[0]), int(pad_tokens[1])),
                reach=reach,
                kind=head_kind,
                width=self.width,
            )
            self.head_kind = head_kind
            self.reach = int(reach)
            self.field_is_prob = bool(final_softmax)
            self.input_hw = tuple(int(v) for v in input_hw)
            self.downsample_factor = int(downsample_factor)
            self.hm_margin_px = int(margin_px)
            self.geom = FieldGeometry.padded_px(
                self.input_hw,
                pad_px=(self.hm_margin_px, self.hm_margin_px),
                stride=2**self.downsample_factor,
            )
            self.output_kind = "heatmap_softargmax"
            self.confidence_kind = "peak_mass"
            #: The tensor contract keeps `V`, which is what the predict path requires;
            #: the views are still computed apart (see the module docstring).
            self.joint_views = True
            self.padded_field = int(margin_px) > 0
            self.owns_prepare = True
            self.has_heatmaps = True
        else:
            raise ValueError(
                f"unknown detector family {family!r}; use 'monocular', 'simcc', 'mvt' "
                "or 'fast'"
            )

    @property
    def n_points(self) -> int:
        """Tracked points: every output channel but the center one, if there is one."""
        return self.num_classes - int(self.center_channel)

    def forward(self, x, view_valid=None, mask_ratio=None):
        """Image(s) -> the family's output: a list, an `(x, y)` bin pair or a tensor.

        `view_valid` (a `(B, V)` mask) and `mask_ratio` (a number or a `(B*V,)` tensor)
        are the mvt's: padded views stay out of the real views' attention, and
        `mask_ratio` overrides `token_mask` for this call. `fast` takes the same tensor
        and ignores both. Dispatch is on `family`: `fast` keeps `V` but is not
        cross-view.
        """
        if self.family == "mvt":
            v = x.shape[1]
            if mask_ratio is None:
                mask_ratio = self.token_mask
            masking = self.training and (
                torch.is_tensor(mask_ratio) or float(mask_ratio) > 0
            )
            feat = vit_multiview.encode(
                self.backbone,
                self.attn_scopes,
                x,
                self.dim,
                norm_mean=self.norm_mean,
                norm_std=self.norm_std,
                view_valid=view_valid,
                mask_token=self.mask_token if masking else None,
                mask_ratio=mask_ratio if masking else None,
                pad_tokens=self.pad_grid,
                pad_token=self.pad_token if self.token_pad else None,
                global_rope=self.global_rope,
                checkpoint=self.grad_checkpoint and self.training,
            )
            hm = self.head(feat)  # (B*V, K, Hm, Wm)
            return hm.reshape(x.shape[0], v * hm.shape[1], *hm.shape[-2:])
        if self.family == "fast":
            if x.ndim != 5:
                raise ValueError(
                    f"the {self.family!r} detector takes (B, V, 1, H, W); got "
                    f"{tuple(x.shape)}. Its views are independent, but the tensor "
                    "keeps the V axis so one moment is one call."
                )
            b, v = x.shape[0], x.shape[1]
            # Fold the views into the batch FIRST: on a 5-D tensor `x[:, :1]` would take
            # one view, not one channel.
            flat = x.reshape(b * v, *x.shape[2:])
            g = (flat[:, :1] - self.norm_mean) / self.norm_std
            feats = cnn_pyramid.forward_features(
                self.backbone, self._sel, self._channels, g
            )
            hm = self.head(feats[self._level])  # (B*V, K, Hm, Wm)
            return hm.reshape(b, v * hm.shape[1], *hm.shape[-2:])
        # `x[:, :1]` is the first channel, not a luminance mix: on monochrome
        # full-range video R == G == B == Y, so it is the luma the network trained on.
        g = (x[:, :1] - self.norm_mean) / self.norm_std
        feats = cnn_pyramid.forward_features(
            self.backbone, self._sel, self._channels, g
        )
        if self.family == "simcc":
            # `(x_logits, y_logits)` from the finest feature map.
            return self.head(feats[0])
        pad_cells = (self.pad_px[0] // STRIDE, self.pad_px[1] // STRIDE)
        return [
            heatmap.forward_field(
                self.laterals,
                self.head,
                self.decoder,
                feats,
                pad_cells,
                reach=self.coarse_reach,
            )
        ]

    def heatmaps_by_view(self, images, view_valid=None, mask_ratio=None):
        """`(B, V, 1, H, W)` -> `(B, V, K, Hm, Wm)` (mvt only)."""
        hm = self.forward(images, view_valid=view_valid, mask_ratio=mask_ratio)
        k = self.num_classes
        return hm.reshape(hm.shape[0], hm.shape[1] // k, k, *hm.shape[-2:])

n_points property

n_points: int

Tracked points: every output channel but the center one, if there is one.

__init__

__init__(
    family: str,
    backbone,
    n_keypoints: int,
    *,
    backbone_name: str | None = None,
    sel=None,
    channels=None,
    head_kind: str = "concat",
    readout: str = "argmax",
    geom: FieldGeometry | None = None,
    margin: float = 0.0,
    attn_scopes: str | None = None,
    dim: int | None = None,
    input_hw=None,
    downsample_factor: int | None = None,
    margin_px: int = 0,
    pad_tokens: tuple[int, int] = (0, 0),
    final_softmax: bool = True,
    reach: int = 0,
    split_ratio: float = 2.0,
    token_mask: float = 0.0,
    token_pad: int = 0,
    global_rope: bool = True,
    level: int = 0,
    feature_stride: int = 0,
    width: int | None = None,
)

Bind the backbone and head for family (see the build_* assemblers).

Source code in src/deeperfly/pose2d/detect.py

def __init__(
    self,
    family: str,
    backbone,
    n_keypoints: int,
    *,
    backbone_name: str | None = None,
    sel=None,
    channels=None,
    head_kind: str = "concat",
    readout: str = "argmax",
    geom: FieldGeometry | None = None,
    margin: float = 0.0,
    attn_scopes: str | None = None,
    dim: int | None = None,
    input_hw=None,
    downsample_factor: int | None = None,
    margin_px: int = 0,
    pad_tokens: tuple[int, int] = (0, 0),
    final_softmax: bool = True,
    reach: int = 0,
    split_ratio: float = 2.0,
    token_mask: float = 0.0,
    token_pad: int = 0,
    global_rope: bool = True,
    level: int = 0,
    feature_stride: int = 0,
    width: int | None = None,
):
    """Bind the backbone and head for `family` (see the `build_*` assemblers)."""
    super().__init__()
    self.family = family
    self.backbone = backbone
    self.backbone_name = backbone_name
    self.num_classes = int(n_keypoints)
    #: Whether the last channel is an animal-center field; set after assembly, so a
    #: state dict never depends on it.
    self.center_channel = False
    self.point_names: list[str] = []
    #: Both families normalize their own grayscale input; a spec declares mean 0.0.
    self.norm_mean = 0.0
    self.norm_std = 1.0
    self.accepts_gray = True
    if family == "monocular":
        self._sel = tuple(sel)
        self._channels = tuple(channels)
        self.head_kind = head_kind
        self.laterals, self.head, self.decoder = heatmap.build_head(
            head_kind, channels, n_keypoints
        )
        # Empty for the shipped heads; `reach` convs on the coarsest lateral
        # otherwise.
        lat = self.laterals[-1].out_channels if len(self.laterals) else 0
        self.coarse_reach = heatmap.coarse_reach(lat, reach)
        self.reach = int(reach)
        self.field_is_prob = False  # the field is raw logits
        self.geom = geom
        self.input_hw = geom.input_hw
        self.stride = geom.stride
        self.pad_px = geom.pad_px
        self.field_hw = geom.field_hw
        #: The margin as a fraction of the input, or `None` when built from `pad_px`
        #: (a pixel policy has no single fraction on a non-square input).
        self.margin = None if margin is None else float(margin)
        self.readout = readout
        # The field is stride-`s`, so the soft-argmax readout recovers the input in
        # log2(s) upsample steps and inverts the per-axis `(py, px)` pad as margin.
        self.downsample_factor = int(geom.stride).bit_length() - 1
        self.hm_margin_px = geom.pad_px
        if readout == "argmax":
            self.output_kind = "heatmap"
            self.confidence_kind = "peak_amax"
        elif readout == "softargmax":
            self.output_kind = "heatmap_softargmax"
            self.confidence_kind = "peak_mass"
        else:
            raise ValueError(
                f"unknown monocular readout {readout!r}; "
                "use 'argmax' or 'softargmax'"
            )
        self.padded_field = True
        self.joint_views = False
        self.owns_prepare = False
        self.has_heatmaps = True
    elif family == "simcc":
        self._sel = tuple(sel)
        self._channels = tuple(channels)
        self.input_hw = tuple(int(v) for v in input_hw)
        self.split_ratio = float(split_ratio)
        self.hm_margin_px = int(margin_px)
        # The backbone sees the unpadded frame; the margin only widens the bin
        # range, so an off-frame coordinate lands on a real bin, not clipped.
        feat_hw = (self.input_hw[0] // STRIDE, self.input_hw[1] // STRIDE)
        self.head = simcc.simcc_head(
            channels[0],
            n_keypoints,
            feat_hw,
            simcc.bins_for(self.input_hw, split_ratio, self.hm_margin_px),
        )
        self.output_kind = "simcc"
        self.confidence_kind = "simcc_axis"
        self.joint_views = False
        self.owns_prepare = False
        # No 2D field to pad; off-frame coordinates come from the widened bin range.
        self.padded_field = False
        self.has_heatmaps = False
    elif family == "mvt":
        #: Where the off-frame margin's cells come from: the innermost `token_pad`
        #: tokens per side are padded into the token grid before the transformer
        #: (learned pad token, attention-capable), the rest zero-padded in the head.
        self.token_pad = int(token_pad)
        if self.token_pad > min(int(pad_tokens[0]), int(pad_tokens[1])):
            raise RuntimeError(
                f"token_pad {token_pad} exceeds the margin's {pad_tokens} tokens"
            )
        self.pad_grid = (self.token_pad, self.token_pad)
        self.pad_token = nn.Parameter(torch.zeros(int(dim)))
        self.head = heatmap.upsample_head(
            dim,
            n_keypoints,
            final_softmax=final_softmax,
            pad_tokens=(
                int(pad_tokens[0]) - self.token_pad,
                int(pad_tokens[1]) - self.token_pad,
            ),
            reach=reach,
            kind=head_kind,
        )
        self.head_kind = head_kind
        self.reach = int(reach)
        #: Rope in the global blocks (see `vit_multiview.encode`); `False` is the
        #: content-only cross-view attention ablation.
        self.global_rope = bool(global_rope)
        #: Recompute the blocks in the backward pass. Set by the training run, not
        #: by the `[model]` table the artifact records: it changes no number.
        self.grad_checkpoint = False
        #: With `final_softmax` the head's field is a per-channel probability map.
        self.field_is_prob = bool(final_softmax)
        self.attn_scopes = list(attn_scopes)
        self.dim = int(dim)
        #: Training-time patch masking (see `vit_multiview.encode`); 0 masks
        #: nothing. The token is always a parameter, so the state dict is fixed.
        self.token_mask = float(token_mask)
        self.mask_token = nn.Parameter(torch.zeros(self.dim))
        self.input_hw = tuple(int(v) for v in input_hw)
        self.downsample_factor = int(downsample_factor)
        self.hm_margin_px = int(margin_px)
        #: The stride-`2 ** downsample_factor` cell grid the head's field lives on,
        #: padded by the margin on every side: what the readout inverts.
        self.geom = FieldGeometry.padded_px(
            self.input_hw,
            pad_px=(self.hm_margin_px, self.hm_margin_px),
            stride=2**self.downsample_factor,
        )
        self.output_kind = "heatmap_softargmax"
        self.confidence_kind = "peak_mass"
        self.joint_views = True
        self.padded_field = int(margin_px) > 0
        self.owns_prepare = True
        self.has_heatmaps = True
    elif family == "fast":
        self._sel = tuple(sel)
        self._channels = tuple(channels)
        #: Which pyramid level the head reads, and its stride in input pixels.
        self._level = int(level)
        self.feature_stride = int(feature_stride)
        self.dim = int(dim)
        self.width = None if width is None else int(width)
        self.head = heatmap.upsample_head(
            self.dim,
            n_keypoints,
            final_softmax=final_softmax,
            # The whole margin is padded in the head: a conv has no attention to
            # read into the pad with, so there is no inner band as in the mvt.
            pad_tokens=(int(pad_tokens[0]), int(pad_tokens[1])),
            reach=reach,
            kind=head_kind,
            width=self.width,
        )
        self.head_kind = head_kind
        self.reach = int(reach)
        self.field_is_prob = bool(final_softmax)
        self.input_hw = tuple(int(v) for v in input_hw)
        self.downsample_factor = int(downsample_factor)
        self.hm_margin_px = int(margin_px)
        self.geom = FieldGeometry.padded_px(
            self.input_hw,
            pad_px=(self.hm_margin_px, self.hm_margin_px),
            stride=2**self.downsample_factor,
        )
        self.output_kind = "heatmap_softargmax"
        self.confidence_kind = "peak_mass"
        #: The tensor contract keeps `V`, which is what the predict path requires;
        #: the views are still computed apart (see the module docstring).
        self.joint_views = True
        self.padded_field = int(margin_px) > 0
        self.owns_prepare = True
        self.has_heatmaps = True
    else:
        raise ValueError(
            f"unknown detector family {family!r}; use 'monocular', 'simcc', 'mvt' "
            "or 'fast'"
        )

forward

forward(x, view_valid=None, mask_ratio=None)

Image(s) -> the family's output: a list, an (x, y) bin pair or a tensor.

view_valid (a (B, V) mask) and mask_ratio (a number or a (B*V,) tensor) are the mvt's: padded views stay out of the real views' attention, and mask_ratio overrides token_mask for this call. fast takes the same tensor and ignores both. Dispatch is on family: fast keeps V but is not cross-view.

Source code in src/deeperfly/pose2d/detect.py

def forward(self, x, view_valid=None, mask_ratio=None):
    """Image(s) -> the family's output: a list, an `(x, y)` bin pair or a tensor.

    `view_valid` (a `(B, V)` mask) and `mask_ratio` (a number or a `(B*V,)` tensor)
    are the mvt's: padded views stay out of the real views' attention, and
    `mask_ratio` overrides `token_mask` for this call. `fast` takes the same tensor
    and ignores both. Dispatch is on `family`: `fast` keeps `V` but is not
    cross-view.
    """
    if self.family == "mvt":
        v = x.shape[1]
        if mask_ratio is None:
            mask_ratio = self.token_mask
        masking = self.training and (
            torch.is_tensor(mask_ratio) or float(mask_ratio) > 0
        )
        feat = vit_multiview.encode(
            self.backbone,
            self.attn_scopes,
            x,
            self.dim,
            norm_mean=self.norm_mean,
            norm_std=self.norm_std,
            view_valid=view_valid,
            mask_token=self.mask_token if masking else None,
            mask_ratio=mask_ratio if masking else None,
            pad_tokens=self.pad_grid,
            pad_token=self.pad_token if self.token_pad else None,
            global_rope=self.global_rope,
            checkpoint=self.grad_checkpoint and self.training,
        )
        hm = self.head(feat)  # (B*V, K, Hm, Wm)
        return hm.reshape(x.shape[0], v * hm.shape[1], *hm.shape[-2:])
    if self.family == "fast":
        if x.ndim != 5:
            raise ValueError(
                f"the {self.family!r} detector takes (B, V, 1, H, W); got "
                f"{tuple(x.shape)}. Its views are independent, but the tensor "
                "keeps the V axis so one moment is one call."
            )
        b, v = x.shape[0], x.shape[1]
        # Fold the views into the batch FIRST: on a 5-D tensor `x[:, :1]` would take
        # one view, not one channel.
        flat = x.reshape(b * v, *x.shape[2:])
        g = (flat[:, :1] - self.norm_mean) / self.norm_std
        feats = cnn_pyramid.forward_features(
            self.backbone, self._sel, self._channels, g
        )
        hm = self.head(feats[self._level])  # (B*V, K, Hm, Wm)
        return hm.reshape(b, v * hm.shape[1], *hm.shape[-2:])
    # `x[:, :1]` is the first channel, not a luminance mix: on monochrome
    # full-range video R == G == B == Y, so it is the luma the network trained on.
    g = (x[:, :1] - self.norm_mean) / self.norm_std
    feats = cnn_pyramid.forward_features(
        self.backbone, self._sel, self._channels, g
    )
    if self.family == "simcc":
        # `(x_logits, y_logits)` from the finest feature map.
        return self.head(feats[0])
    pad_cells = (self.pad_px[0] // STRIDE, self.pad_px[1] // STRIDE)
    return [
        heatmap.forward_field(
            self.laterals,
            self.head,
            self.decoder,
            feats,
            pad_cells,
            reach=self.coarse_reach,
        )
    ]

heatmaps_by_view

heatmaps_by_view(images, view_valid=None, mask_ratio=None)

(B, V, 1, H, W) -> (B, V, K, Hm, Wm) (mvt only).

Source code in src/deeperfly/pose2d/detect.py

def heatmaps_by_view(self, images, view_valid=None, mask_ratio=None):
    """`(B, V, 1, H, W)` -> `(B, V, K, Hm, Wm)` (mvt only)."""
    hm = self.forward(images, view_valid=view_valid, mask_ratio=mask_ratio)
    k = self.num_classes
    return hm.reshape(hm.shape[0], hm.shape[1] // k, k, *hm.shape[-2:])

build_fast

build_fast(
    num_keypoints: int,
    *,
    backbone: str = "efficientvit_b1",
    input_hw: tuple[int, int] = (256, 512),
    downsample_factor: int = 2,
    margin: int = 0,
    final_softmax: bool = True,
    reach: int = 0,
    width: int | None = None,
    pretrained: bool = False,
    head: str = "upsample",
    norm_mean: float | None = None,
    norm_std: float | None = None,
) -> Detector

Assemble a single-view convolutional detector, the fast family.

The mvt's tensor contract and head over a timm pyramid on each view alone. The views never inform each other, which costs a measured 0.23 px overall (0.52 on side-view legs), and a small conv net is 3 to 6 times cheaper per view than the smallest usable ViT.

Parameters:

Name	Type	Description	Default
num_keypoints	int	Heatmap channels per view.	required
backbone	str	A timm name exposing the stride-4..32 pyramid.	'efficientvit_b1'
input_hw	tuple[int, int]	The network input (height, width); the margin is added in the head.	(256, 512)
downsample_factor	int	log2 of the field's stride; the head upsamples 4x, so it reads the pyramid level at 4 * 2 ** this.	2
margin	int	Input pixels of field past the frame on every side, in whole feature cells.	0
final_softmax	bool	Emit a temperature-1 spatial-softmax field.	True
reach	int	3x3 convs at the feature stride between the pad and the upsampling, one feature cell of margin each; margin // feature_stride covers it, and 0 cannot localize an off-frame joint.	0
width	int | None	A 1x1 projection before the reach convs: on a 512-wide stride-16 map one reach conv costs more than the rest of a small net.	None
pretrained	bool	Load timm's ImageNet weights, this family's only free warm start: init_from cannot bring a ViT's tensors across (the chain: 1.9 px).	False
head	str	"upsample" or "deconv".	'upsample'
norm_mean	float | None	As for the mvt.	None
norm_std	float | None	As for the mvt.	None

Raises:

Type	Description
RuntimeError	If the backbone lacks the pyramid or the level, or margin or input_hw is not whole feature cells, which timm would floor silently.

Source code in src/deeperfly/pose2d/detect.py

def build_fast(
    num_keypoints: int,
    *,
    backbone: str = "efficientvit_b1",
    input_hw: tuple[int, int] = (256, 512),
    downsample_factor: int = 2,
    margin: int = 0,
    final_softmax: bool = True,
    reach: int = 0,
    width: int | None = None,
    pretrained: bool = False,
    head: str = "upsample",
    norm_mean: float | None = None,
    norm_std: float | None = None,
) -> Detector:
    """Assemble a single-view convolutional detector, the `fast` family.

    The mvt's tensor contract and head over a `timm` pyramid on each view alone. The
    views never inform each other, which costs a measured 0.23 px overall (0.52 on
    side-view legs), and a small conv net is 3 to 6 times cheaper per view than the
    smallest usable ViT.

    Args:
        num_keypoints: Heatmap channels per view.
        backbone: A `timm` name exposing the stride-4..32 pyramid.
        input_hw: The network input `(height, width)`; the margin is added in the head.
        downsample_factor: log2 of the field's stride; the head upsamples 4x, so it
            reads the pyramid level at `4 * 2 ** this`.
        margin: Input pixels of field past the frame on every side, in whole feature
            cells.
        final_softmax: Emit a temperature-1 spatial-softmax field.
        reach: 3x3 convs at the feature stride between the pad and the upsampling, one
            feature cell of margin each; `margin // feature_stride` covers it, and 0
            cannot localize an off-frame joint.
        width: A 1x1 projection before the reach convs: on a 512-wide stride-16 map one
            reach conv costs more than the rest of a small net.
        pretrained: Load `timm`'s ImageNet weights, this family's only free warm
            start: `init_from` cannot bring a ViT's tensors across (the chain: 1.9 px).
        head: `"upsample"` or `"deconv"`.
        norm_mean: As for the mvt.
        norm_std: As for the mvt.

    Raises:
        RuntimeError: If the backbone lacks the pyramid or the level, or `margin` or
            `input_hw` is not whole feature cells, which `timm` would floor silently.
    """
    stride = 4 * 2 ** int(downsample_factor)
    if stride not in cnn_pyramid.WANT_REDUCTIONS:
        raise RuntimeError(
            f"downsample_factor {downsample_factor} wants the stride-{stride} "
            f"level; the pyramid has {cnn_pyramid.WANT_REDUCTIONS} (the head "
            "upsamples by 4x)"
        )
    if margin % stride:
        raise RuntimeError(
            f"fast margin {margin} must be a whole number of stride-{stride} feature "
            "cells; the head pads the feature grid, so an off-cell margin has no cell "
            "to pad by"
        )
    if input_hw[0] % stride or input_hw[1] % stride:
        raise RuntimeError(
            f"fast input_hw {tuple(input_hw)} must be a whole number of "
            f"stride-{stride} cells; timm floors the feature grid instead of refusing, "
            "and the emitted field would not match the geometry the readout inverts"
        )
    model, sel, chs = cnn_pyramid.create(
        backbone, input_hw=input_hw, pretrained=pretrained, in_chans=1
    )
    level = cnn_pyramid.WANT_REDUCTIONS.index(stride)
    det = Detector(
        "fast",
        model,
        num_keypoints,
        backbone_name=backbone,
        sel=sel,
        channels=chs,
        level=level,
        feature_stride=stride,
        dim=chs[level],
        width=width,
        input_hw=input_hw,
        downsample_factor=downsample_factor,
        margin_px=margin,
        pad_tokens=(margin // stride, margin // stride),
        final_softmax=final_softmax,
        reach=reach,
        head_kind=head,
    )
    if norm_mean is None:
        norm_mean = 0.449 if pretrained else 0.0
    if norm_std is None:
        norm_std = 0.226 if pretrained else 1.0
    det.norm_mean, det.norm_std = float(norm_mean), float(norm_std)
    return det

build_monocular

build_monocular(
    n_keypoints: int,
    backbone: str = "hrnet_w32",
    *,
    pretrained: bool = False,
    head: str = "concat",
    readout: str = "argmax",
    input_hw: tuple[int, int] = (256, 512),
    margin: float | None = DEFAULT_MARGIN,
    margin_px: int | tuple[int, int] | None = None,
    reach: int = 0,
) -> Detector

Assemble a monocular Detector from named parts.

Parameters:

Name	Type	Description	Default
n_keypoints	int	Output heatmap channels (one per tracked point).	required
backbone	str	A timm model name ("hrnet_w32", "hgnetv2_b4", ...).	'hrnet_w32'
pretrained	bool	Load timm's ImageNet weights into the backbone.	False
head	str	"concat" or "unet" (see heatmap).	'concat'
readout	str	"argmax" (hard peak and parabola, the shipped default) or "softargmax" (the mvt's differentiable expectation): an ablation of the decode and the loss over the same field.	'argmax'
input_hw	tuple[int, int]	Network input (height, width).	(256, 512)
margin	float | None	Field margin as a fraction of the input, the shipped checkpoints' convention.	DEFAULT_MARGIN
margin_px	int | tuple[int, int] | None	Field margin in input pixels per side, (y, x) or one value, overriding margin: the training policy (64 px), in whole stride-4 cells.	None
reach	int	3x3 convs on the padded coarsest lateral so the margin sees image content (see forward_field); 0 is the shipped head.	0

Source code in src/deeperfly/pose2d/detect.py

def build_monocular(
    n_keypoints: int,
    backbone: str = "hrnet_w32",
    *,
    pretrained: bool = False,
    head: str = "concat",
    readout: str = "argmax",
    input_hw: tuple[int, int] = (256, 512),
    margin: float | None = DEFAULT_MARGIN,
    margin_px: int | tuple[int, int] | None = None,
    reach: int = 0,
) -> Detector:
    """Assemble a monocular [Detector][] from named parts.

    Args:
        n_keypoints: Output heatmap channels (one per tracked point).
        backbone: A `timm` model name (`"hrnet_w32"`, `"hgnetv2_b4"`, ...).
        pretrained: Load `timm`'s ImageNet weights into the backbone.
        head: `"concat"` or `"unet"` (see [heatmap][deeperfly.pose2d.heads.heatmap]).
        readout: `"argmax"` (hard peak and parabola, the shipped default) or
            `"softargmax"` (the mvt's differentiable expectation): an ablation of the
            decode and the loss over the same field.
        input_hw: Network input `(height, width)`.
        margin: Field margin as a fraction of the input, the shipped checkpoints'
            convention.
        margin_px: Field margin in input pixels per side, `(y, x)` or one value,
            overriding `margin`: the training policy (64 px), in whole stride-4 cells.
        reach: 3x3 convs on the padded coarsest lateral so the margin sees image
            content (see `forward_field`); 0 is the shipped head.
    """
    model, sel, chs = cnn_pyramid.create(
        backbone, input_hw=input_hw, pretrained=pretrained, in_chans=1
    )
    if margin_px is not None:
        py, px = (margin_px, margin_px) if isinstance(margin_px, int) else margin_px
        geom = FieldGeometry.padded_px(input_hw, pad_px=(py, px), stride=STRIDE)
        margin = None
    else:
        geom = FieldGeometry.padded(input_hw, margin=margin, stride=STRIDE)
    return Detector(
        "monocular",
        model,
        n_keypoints,
        backbone_name=backbone,
        sel=sel,
        channels=chs,
        head_kind=head,
        readout=readout,
        geom=geom,
        margin=margin,
        reach=reach,
    )

build_mvt

build_mvt(
    num_keypoints: int,
    *,
    backbone: str = "vit_base_patch16_224",
    input_hw: tuple[int, int] = (256, 512),
    attn_scopes: str = "llgllgllgllg",
    downsample_factor: int = 2,
    margin: int = 0,
    final_softmax: bool = True,
    reach: int = 0,
    pretrained: bool = False,
    drop_path: float = 0.0,
    token_mask: float = 0.0,
    token_pad: int = 0,
    norm_mean: float | None = None,
    norm_std: float | None = None,
    global_rope: bool = True,
    head: str = "upsample",
) -> Detector

Assemble a multiview-transformer Detector.

Parameters:

Name	Type	Description	Default
num_keypoints	int	Heatmap channels per view.	required
backbone	str	A timm ViT name whose blocks the schedule drives.	'vit_base_patch16_224'
input_hw	tuple[int, int]	The network input (height, width); the margin is added in the head.	(256, 512)
attn_scopes	str	One "l" or "g" per transformer block (see vit_multiview).	'llgllgllgllg'
downsample_factor	int	log2 of the field's stride in input pixels (2: stride 4).	2
margin	int	Model pixels of field past the frame on every side, in whole patches.	0
final_softmax	bool	Emit a temperature-1 spatial-softmax map from the head.	True
reach	int	Token-level 3x3 convs between the padding and the upsampling, so the margin sees image content; 0 is the shipped head, 2 the training policy.	0
pretrained	bool	Load timm's pretrained ViT weights.	False
drop_path	float	Stochastic depth rate across the blocks (training only).	0.0
token_mask	float	Fraction of each view's patch tokens replaced by a learned mask token in training; 0 masks nothing.	0.0
token_pad	int	Tokens per side of margin padded into the token grid before the transformer, so attention reaches them, at (1 + 2 t p / H) (1 + 2 t p / W) times the tokens (1.26x for 2 at 384x768).	0
norm_mean	float | None	Subtracted from the [0, 1] input before the patch embed; None is 0 for a random init and the ImageNet gray mean (0.449) with pretrained.	None
norm_std	float | None	The divisor after norm_mean; None is 1, 0.226 with pretrained.	None
global_rope	bool	Apply the tiled rope in the global blocks; False runs the cross-view attention on content alone.	True
head	str	"upsample" (the shipped head) or "deconv" (ViTPose's); see upsample_head.	'upsample'

Source code in src/deeperfly/pose2d/detect.py

def build_mvt(
    num_keypoints: int,
    *,
    backbone: str = "vit_base_patch16_224",
    input_hw: tuple[int, int] = (256, 512),
    attn_scopes: str = "llgllgllgllg",
    downsample_factor: int = 2,
    margin: int = 0,
    final_softmax: bool = True,
    reach: int = 0,
    pretrained: bool = False,
    drop_path: float = 0.0,
    token_mask: float = 0.0,
    token_pad: int = 0,
    norm_mean: float | None = None,
    norm_std: float | None = None,
    global_rope: bool = True,
    head: str = "upsample",
) -> Detector:
    """Assemble a multiview-transformer [Detector][deeperfly.pose2d.detect.Detector].

    Args:
        num_keypoints: Heatmap channels per view.
        backbone: A `timm` ViT name whose blocks the schedule drives.
        input_hw: The network input `(height, width)`; the margin is added in the head.
        attn_scopes: One `"l"` or `"g"` per transformer block (see
            [vit_multiview][deeperfly.pose2d.backbones.vit_multiview]).
        downsample_factor: log2 of the field's stride in input pixels (2: stride 4).
        margin: Model pixels of field past the frame on every side, in whole patches.
        final_softmax: Emit a temperature-1 spatial-softmax map from the head.
        reach: Token-level 3x3 convs between the padding and the upsampling, so the
            margin sees image content; 0 is the shipped head, 2 the training policy.
        pretrained: Load `timm`'s pretrained ViT weights.
        drop_path: Stochastic depth rate across the blocks (training only).
        token_mask: Fraction of each view's patch tokens replaced by a learned mask
            token in training; 0 masks nothing.
        token_pad: Tokens per side of `margin` padded into the token grid before the
            transformer, so attention reaches them, at
            `(1 + 2 t p / H) (1 + 2 t p / W)` times the tokens (1.26x for 2 at 384x768).
        norm_mean: Subtracted from the `[0, 1]` input before the patch embed; `None` is
            0 for a random init and the ImageNet gray mean (0.449) with `pretrained`.
        norm_std: The divisor after `norm_mean`; `None` is 1, 0.226 with `pretrained`.
        global_rope: Apply the tiled rope in the global blocks; `False` runs the
            cross-view attention on content alone.
        head: `"upsample"` (the shipped head) or `"deconv"` (ViTPose's); see
            [upsample_head][deeperfly.pose2d.heads.heatmap.upsample_head].
    """
    model, dim, patch = vit_multiview.create(
        backbone,
        input_hw=input_hw,
        attn_scopes=attn_scopes,
        pretrained=pretrained,
        drop_path_rate=drop_path,
    )
    ph, pw = patch
    if margin % ph or margin % pw:
        raise RuntimeError(
            f"mvt margin {margin} must be a whole number of patches "
            f"(patch {ph}x{pw}); the head pads the token grid, so an off-patch margin "
            "has no cell to pad by"
        )
    det = Detector(
        "mvt",
        model,
        num_keypoints,
        backbone_name=backbone,
        attn_scopes=attn_scopes,
        dim=dim,
        input_hw=input_hw,
        downsample_factor=downsample_factor,
        margin_px=margin,
        pad_tokens=(margin // ph, margin // pw),
        final_softmax=final_softmax,
        reach=reach,
        token_mask=token_mask,
        token_pad=token_pad,
        global_rope=global_rope,
        head_kind=head,
    )
    if norm_mean is None:
        norm_mean = 0.449 if pretrained else 0.0
    if norm_std is None:
        norm_std = 0.226 if pretrained else 1.0
    det.norm_mean, det.norm_std = float(norm_mean), float(norm_std)
    return det

build_simcc

build_simcc(
    n_keypoints: int,
    backbone: str = "hrnet_w32",
    *,
    pretrained: bool = False,
    input_hw: tuple[int, int] = (256, 512),
    split_ratio: float = 2.0,
    margin: int = 0,
) -> Detector

Assemble a SimCC Detector on a CNN pyramid.

Parameters:

Name	Type	Description	Default
n_keypoints	int	Output points, one 1D distribution pair each.	required
backbone	str	A timm model name ("hrnet_w32", "hgnetv2_b4", ...).	'hrnet_w32'
pretrained	bool	Load timm's ImageNet weights into the backbone.	False
input_hw	tuple[int, int]	Network input (height, width).	(256, 512)
split_ratio	float	Bins per input pixel per axis.	2.0
margin	int	Input pixels the bin range extends past the frame on every side.	0

Source code in src/deeperfly/pose2d/detect.py

def build_simcc(
    n_keypoints: int,
    backbone: str = "hrnet_w32",
    *,
    pretrained: bool = False,
    input_hw: tuple[int, int] = (256, 512),
    split_ratio: float = 2.0,
    margin: int = 0,
) -> Detector:
    """Assemble a SimCC [Detector][deeperfly.pose2d.detect.Detector] on a CNN pyramid.

    Args:
        n_keypoints: Output points, one 1D distribution pair each.
        backbone: A `timm` model name (`"hrnet_w32"`, `"hgnetv2_b4"`, ...).
        pretrained: Load `timm`'s ImageNet weights into the backbone.
        input_hw: Network input `(height, width)`.
        split_ratio: Bins per input pixel per axis.
        margin: Input pixels the bin range extends past the frame on every side.
    """
    model, sel, chs = cnn_pyramid.create(
        backbone, input_hw=input_hw, pretrained=pretrained, in_chans=1
    )
    return Detector(
        "simcc",
        model,
        n_keypoints,
        backbone_name=backbone,
        sel=sel,
        channels=chs,
        input_hw=input_hw,
        split_ratio=split_ratio,
        margin_px=margin,
    )

export_monocular

export_monocular(
    model: Detector, path: str | Path, *, point_names=None
) -> dict

Save a monocular model as a deeperfly-dense38-1 checkpoint; return the dict.

point_names defaults to the model's own.

Raises:

Type	Description
ValueError	If model is not monocular, or its backbone name is unknown.

Source code in src/deeperfly/pose2d/detect.py

def export_monocular(model: Detector, path: str | Path, *, point_names=None) -> dict:
    """Save a monocular `model` as a `deeperfly-dense38-1` checkpoint; return the dict.

    `point_names` defaults to the model's own.

    Raises:
        ValueError: If `model` is not monocular, or its backbone name is unknown.
    """
    if model.family != "monocular":
        raise ValueError(
            f"only the monocular family has the {ARTIFACT_FORMAT!r} format; "
            f"got {model.family!r}"
        )
    if not model.backbone_name:
        raise ValueError(
            "the detector has no backbone_name; build it via build_monocular so the "
            "loader can reconstruct the architecture"
        )
    ck = {
        "format": ARTIFACT_FORMAT,
        "state_dict": model.state_dict(),
        "backbone": model.backbone_name,
        "head": model.head_kind,
        # Parameter-free, so `strict=True` cannot catch its loss: a soft-argmax model
        # would silently reload as hard-argmax. `load_monocular` reads it back.
        "readout": model.readout,
        "n_keypoints": int(model.num_classes),
        "input_size": list(model.input_hw),
        "hm_margin": None if model.margin is None else float(model.margin),
        # The field pad in input pixels `(y, x)`; the loader prefers it, so a model
        # built by the pixel policy reloads with the same field whatever the fraction.
        "hm_margin_px": [int(v) for v in model.pad_px],
        "reach": int(model.reach),
        "point_names": list(
            point_names if point_names is not None else model.point_names
        ),
        "mean": float(model.norm_mean),
        "std": float(model.norm_std),
        # Every label it trained on is at a pixel center (ADR-0055, ADR-0078).
        "pixel_grid": "center",
    }
    torch.save(ck, Path(path))
    return ck

load_monocular

load_monocular(
    weights: str | Path,
    *,
    dev: str | None = None,
    mean: float = 0.0,
)

Load a deeperfly-dense38-1 checkpoint as an eval-mode detector on dev.

The architecture comes from the artifact, loaded strict=True; the checkpoint carries its own normalization, so mean must be 0.0.

Raises:

Type	Description
SystemExit	On a nonzero mean, a missing file, a wrong format, or a cross-view checkpoint.

Source code in src/deeperfly/pose2d/detect.py

def load_monocular(weights: str | Path, *, dev: str | None = None, mean: float = 0.0):
    """Load a `deeperfly-dense38-1` checkpoint as an eval-mode detector on `dev`.

    The architecture comes from the artifact, loaded `strict=True`; the checkpoint
    carries its own normalization, so `mean` must be 0.0.

    Raises:
        SystemExit: On a nonzero `mean`, a missing file, a wrong `format`, or a
            cross-view checkpoint.
    """
    if float(mean) != 0.0:
        raise SystemExit(
            f"a monocular model must declare mean = 0.0 (got {mean}); the checkpoint "
            "carries its own mean/std and applies them itself"
        )
    path = Path(weights)
    if not path.exists():
        raise SystemExit(f"no monocular checkpoint at {path}")
    ck = torch.load(path, map_location="cpu", weights_only=False)
    fmt = ck.get("format") if isinstance(ck, dict) else None
    if fmt != ARTIFACT_FORMAT:
        raise SystemExit(
            f"{path}: format {fmt!r}, expected {ARTIFACT_FORMAT!r}. Not an exported "
            "monocular artifact."
        )
    sd = ck["state_dict"]
    if any(k.startswith(("fusion.", "fuse_in.", "fuse_out.")) for k in sd):
        raise SystemExit(
            f"{path} is a CROSS-VIEW checkpoint: its residual is not the identity once "
            "trained, so running it one view at a time evaluates a different function. "
            "Run it through the grouped multiview path, not a per-view pathway."
        )
    names = list(ck.get("point_names") or [])
    head = str(ck.get("head") or ck.get("args", {}).get("head", "concat"))
    # The readout is parameter-free too, so a checkpoint trained for the soft-argmax
    # decode must record it; the shipped artifacts predate it and take the hard argmax.
    readout = str(ck.get("readout") or ck.get("args", {}).get("readout") or "argmax")
    n_kp = int(ck.get("n_keypoints") or len(names))
    input_hw = tuple(int(v) for v in ck.get("input_size", (256, 512)))
    # The margin is read from the artifact, else the ship default: parameter-free, so
    # `strict=True` cannot catch a disagreement, and recording it is the only guard. A
    # pixel margin (`hm_margin_px`, the training policy) wins over the fraction.
    margin_px = ck.get("hm_margin_px")
    margin = ck.get("hm_margin")
    if margin is None:
        margin = ck.get("args", {}).get("hm_margin")
    margin = DEFAULT_MARGIN if margin is None else float(margin)
    model = build_monocular(
        n_kp,
        ck["backbone"],
        head=head,
        readout=readout,
        input_hw=input_hw,
        margin=margin,
        margin_px=None if margin_px is None else tuple(int(v) for v in margin_px),
        reach=int(ck.get("reach", 0)),
    )
    # Strict, so a three-channel checkpoint fails here, on the stem, not at run time.
    model.load_state_dict(sd, strict=True)
    model.norm_mean = float(ck.get("mean", 0.0))
    model.norm_std = float(ck.get("std", 1.0))
    model.point_names = names
    model.pixel_grid = artifact_grid(ck, path)
    target = dev or device()
    model = model.eval().to(target)
    # A margin built from pixels has no fraction to print.
    margin_text = "in pixels" if model.margin is None else f"{model.margin:.2f}"
    log.info(
        "monocular detector %s: %s/%s, %d channels, input %dx%d, field %dx%d "
        "(margin %s), mean %.4f std %.4f, on %s",
        path.name,
        ck["backbone"],
        head,
        n_kp,
        *model.input_hw,
        *model.field_hw,
        margin_text,
        model.norm_mean,
        model.norm_std,
        target,
    )
    return model

Running an assembled detector: input prep, decode dispatch, and the predict surface.

detect is what a detector is; this is how to run it, with the readout dispatched on the model's output_kind. The monocular family folds every leading axis into the batch, and the multiview transformer keeps V.

decode_candidates

decode_candidates(
    model,
    outputs,
    k: int,
    *,
    radius_px: float = CANDIDATE_RADIUS_PX,
)

Forward outputs -> the k strongest distinct peaks per cell, by output_kind.

(N, K, k, 2) input-normalized peaks, NaN where a cell has fewer, and (N, K, k) scores in the family's confidence units, each refined and scored as its argmax readout is. radius_px is the separation, in input pixels, below which two maxima are one.

Source code in src/deeperfly/pose2d/infer.py

def decode_candidates(
    model, outputs, k: int, *, radius_px: float = CANDIDATE_RADIUS_PX
):
    """Forward outputs -> the `k` strongest distinct peaks per cell, by `output_kind`.

    `(N, K, k, 2)` input-normalized peaks, NaN where a cell has fewer, and `(N, K, k)`
    scores in the family's confidence units, each refined and scored as its argmax
    readout is. `radius_px` is the separation, in input pixels, below which two maxima
    are one.
    """
    kind = getattr(model, "output_kind", "heatmap")
    if kind == "simcc":
        x_logits, y_logits = outputs
        return simcc.simcc_peaks(
            x_logits,
            y_logits,
            k,
            split_ratio=model.split_ratio,
            input_hw=model.input_hw,
            margin=getattr(model, "hm_margin_px", 0),
            radius=max(1, round(radius_px * model.split_ratio)),
        )
    if kind not in ("heatmap", "heatmap_softargmax"):
        raise NotImplementedError(
            f"no candidate readout for output_kind {kind!r}; add a branch beside the "
            "heatmap readouts"
        )
    field = (outputs[-1] if isinstance(outputs, list) else outputs).float()
    n, c, h, w = field.shape
    radius = max(1, round(radius_px / model.geom.stride))
    if kind == "heatmap":
        flat = field.reshape(n * c, h, w)
        cells, val = field_peaks(flat, k, radius=radius)
        refined = heatmap.parabola_at(flat, cells[..., 1], cells[..., 0])
        score = val
    else:
        prob = field if getattr(model, "field_is_prob", False) else None
        prob = heatmap.spatial_softmax(field, 1.0) if prob is None else prob
        prob = prob.reshape(n * c, h, w)
        # Peaks and their refinement on the modulated field, as the argmax readout; the
        # score on the raw probability, as `peak_mass` is.
        mod = heatmap.dark_modulate(prob)
        cells, val = field_peaks(mod, k, radius=radius)
        refined = heatmap.dark_at(mod, cells[..., 1], cells[..., 0])
        score = peak_mass_at(prob, cells)
    valid = torch.isfinite(val)
    xy = model.geom.to_input_normalized(refined)
    xy = torch.where(valid[..., None], xy, torch.full_like(xy, float("nan")))
    score = torch.where(valid, score, torch.zeros_like(score))
    return xy.reshape(n, c, k, 2), score.reshape(n, c, k)

decode_points

decode_points(model, outputs)

Forward outputs -> input-normalized peaks and confidence, by output_kind.

"heatmap" by parabolic sub-pixel refinement over the padded field, "heatmap_softargmax" by DARK refinement (dark_decode), "simcc" by per-axis soft-argmax; an off-frame coordinate falls outside [0, 1].

Source code in src/deeperfly/pose2d/infer.py

def decode_points(model, outputs):
    """Forward outputs -> input-normalized peaks and confidence, by `output_kind`.

    `"heatmap"` by parabolic sub-pixel refinement over the padded field,
    `"heatmap_softargmax"` by DARK refinement
    ([dark_decode][deeperfly.pose2d.heads.heatmap.dark_decode]), `"simcc"` by per-axis
    soft-argmax; an off-frame coordinate falls outside `[0, 1]`.
    """
    kind = getattr(model, "output_kind", "heatmap")
    if kind == "heatmap":
        field = (outputs[-1] if isinstance(outputs, list) else outputs).float()
        return model.geom.to_input_normalized(
            heatmap.argmax_parabola(field)
        ), peak_amax(field)
    if kind == "heatmap_softargmax":
        field = (outputs[-1] if isinstance(outputs, list) else outputs).float()
        if not getattr(model, "field_is_prob", False):
            field = heatmap.spatial_softmax(field, 1.0)
        return heatmap.dark_decode(field, geom=model.geom)
    if kind == "simcc":
        x_logits, y_logits = outputs
        return simcc.simcc_decode(
            x_logits,
            y_logits,
            split_ratio=model.split_ratio,
            input_hw=model.input_hw,
            margin=getattr(model, "hm_margin_px", 0),
        )
    raise NotImplementedError(
        f"no decode for output_kind {kind!r}; add a branch beside the heatmap readouts"
    )

joint_channels

joint_channels(model, outputs)

A flat (N, C, h, w) forward output without its center channel: (N, K, h, w).

The center field is the last channel, and every joint readout drops it first; a model without one gets its output back untouched, whatever the family wraps it in.

Source code in src/deeperfly/pose2d/infer.py

def joint_channels(model, outputs):
    """A flat `(N, C, h, w)` forward output without its center channel: `(N, K, h, w)`.

    The center field is the last channel, and every joint readout drops it first; a
    model without one gets its output back untouched, whatever the family wraps it in.
    """
    if not getattr(model, "center_channel", False):
        return outputs
    k = model.n_points
    if isinstance(outputs, list):
        return [out[:, :k] for out in outputs]
    return outputs[:, :k]

predict_heatmaps

predict_heatmaps(model, inputs) -> np.ndarray

Final heatmaps as host NumPy, (..., K, Hf, Wf) or mvt's (B, V, K, Hm, Wm).

Raises:

Type	Description
NotImplementedError	For a family with no 2D field (SimCC).

Source code in src/deeperfly/pose2d/infer.py

def predict_heatmaps(model, inputs) -> np.ndarray:
    """Final heatmaps as host NumPy, `(..., K, Hf, Wf)` or mvt's `(B, V, K, Hm, Wm)`.

    Raises:
        NotImplementedError: For a family with no 2D field (SimCC).
    """
    if not getattr(model, "has_heatmaps", True):
        raise NotImplementedError(
            f"the {model.family!r} detector emits no 2D heatmap field; "
            "use predict_points"
        )
    dev = next(model.parameters()).device
    if model.joint_views:
        x = inputs if torch.is_tensor(inputs) else torch.as_tensor(np.asarray(inputs))
        x = x.to(device=dev, dtype=torch.float32)
        if x.ndim == 4:
            x = x.unsqueeze(0)
        with torch.inference_mode():
            hm = model.heatmaps_by_view(x)
        return hm.float().cpu().numpy()
    x = as_torch(inputs).float().to(dev)
    lead = x.shape[:-3]
    flat = x.reshape(-1, *x.shape[-3:])
    dtype = autocast_dtype(model, dev)
    with torch.inference_mode():
        if dtype is not None:
            with torch.autocast(dev.type, dtype=dtype):
                hm = model(flat)[-1]
        else:
            hm = model(flat)[-1]
    return hm.float().reshape(*lead, *hm.shape[-3:]).cpu().numpy()

predict_points

predict_points(
    model,
    inputs,
    *,
    method: str = "weighted",
    radius: int = 2,
    views=None,
)

Fused forward and decode -> input-normalized peaks and confidence.

method and radius are ignored: the readout is the one the targets trained for. For mvt the V axis is kept, and views decodes only those views.

Source code in src/deeperfly/pose2d/infer.py

def predict_points(
    model, inputs, *, method: str = "weighted", radius: int = 2, views=None
):
    """Fused forward and decode -> input-normalized peaks and confidence.

    `method` and `radius` are ignored: the readout is the one the targets trained for.
    For mvt the `V` axis is kept, and `views` decodes only those views.
    """
    if model.joint_views:
        return _predict_points_mvt(model, inputs, views=views)
    dev = next(model.parameters()).device
    x = as_torch(inputs).float().to(dev)
    lead = x.shape[:-3]
    flat = x.reshape(-1, *x.shape[-3:])
    dtype = autocast_dtype(model, dev)
    with torch.inference_mode():
        if dtype is not None:
            with torch.autocast(dev.type, dtype=dtype):
                out = model(flat)
        else:
            out = model(flat)
        xy, conf = decode_points(model, out)
    if dev.type == "cuda":
        torch.cuda.synchronize()
    c, k = model.num_classes, model.n_points
    return (
        xy.reshape(*lead, c, 2)[..., :k, :].cpu().numpy().astype(np.float32),
        conf.reshape(*lead, c)[..., :k].cpu().numpy().astype(np.float32),
    )

prepare_images

prepare_images(frames, input_hw, device)

(..., H, W[, C]) frames -> (..., 1, h, w) ViT input, on device.

A plain bilinear resize of one gray plane; the model normalizes it and adds the off-frame margin in its head.

Source code in src/deeperfly/pose2d/infer.py

def prepare_images(frames, input_hw, device):
    """`(..., H, W[, C])` frames -> `(..., 1, h, w)` ViT input, on `device`.

    A plain bilinear resize of one gray plane; the model normalizes it and adds the
    off-frame margin in its head.
    """
    import torch.nn.functional as F

    img = as_torch(frames)
    img = img.float() / 255.0 if not torch.is_floating_point(img) else img.float()
    if img.ndim >= 3 and img.shape[-1] in (1, 3, 4):
        img = img[..., :1]
    else:
        img = img.unsqueeze(-1)
    chw = img.movedim(-1, -3).contiguous().to(device)
    lead = chw.shape[:-3]
    flat = chw.reshape(-1, *chw.shape[-3:])
    out = F.interpolate(
        flat, size=tuple(input_hw), mode="bilinear", align_corners=False, antialias=True
    )
    return out.reshape(*lead, *out.shape[-3:])

Top-K distinct peaks per cell: the detector side of the candidate path (ADR-0038).

The single-peak readout commits a cell to its strongest peak, and when that is the wrong one (a crossed leg, a left/right confusion) nothing downstream gets the joint back. The candidate path keeps the K strongest distinct maxima, refined and scored as the head's own readout, so pictorial can prefer a weaker peak the other views agree with; slot 0 is the decoded point.

field_peaks

field_peaks(field, k: int, *, radius: int)

The k strongest local maxima per (N, H, W) field, as cells and values.

(N, k, 2) integer (x, y) cells and (N, k) values. A cell is a maximum when it equals its (2 radius + 1) square's maximum and exceeds its minimum, so a flat plateau keeps none; missing slots carry -inf at (0, 0).

Source code in src/deeperfly/pose2d/candidates.py

def field_peaks(field, k: int, *, radius: int):
    """The `k` strongest local maxima per `(N, H, W)` field, as cells and values.

    `(N, k, 2)` integer `(x, y)` cells and `(N, k)` values. A cell is a maximum when it
    equals its `(2 radius + 1)` square's maximum and exceeds its minimum, so a flat
    plateau keeps none; missing slots carry `-inf` at `(0, 0)`.
    """
    n, h, w = field.shape
    x = field.float()
    size = 2 * radius + 1
    highest = F.max_pool2d(x[:, None], size, stride=1, padding=radius)[:, 0]
    lowest = -F.max_pool2d(-x[:, None], size, stride=1, padding=radius)[:, 0]
    neg = torch.full_like(x, float("-inf"))
    peaks = torch.where((x == highest) & (x > lowest), x, neg).reshape(n, h * w)
    m = min(k, h * w)
    val, idx = peaks.topk(m, dim=-1)
    if m < k:
        val = torch.cat([val, val.new_full((n, k - m), float("-inf"))], dim=1)
        idx = torch.cat([idx, idx.new_zeros((n, k - m))], dim=1)
    return torch.stack([idx % w, idx // w], dim=-1), val

merge_candidates

merge_candidates(
    anchor_xy,
    anchor_score,
    pools: Sequence[tuple],
    k: int,
    *,
    radius: float,
)

One set of k distinct candidates per row: the anchor, then the pooled peaks.

The anchor (anchor_xy (M, 2), anchor_score (M,)) takes slot 0, a NaN one none; the pooled peaks follow strongest first, each dropped within radius of one already kept. Returns (M, k, 2) and (M, k), NaN at score 0 where fewer exist.

Source code in src/deeperfly/pose2d/candidates.py

def merge_candidates(
    anchor_xy, anchor_score, pools: Sequence[tuple], k: int, *, radius: float
):
    """One set of `k` distinct candidates per row: the anchor, then the pooled peaks.

    The anchor (`anchor_xy` `(M, 2)`, `anchor_score` `(M,)`) takes slot 0, a NaN one
    none; the pooled peaks follow strongest first, each dropped within `radius` of one
    already kept. Returns `(M, k, 2)` and `(M, k)`, NaN at score 0 where fewer exist.
    """
    xy = torch.cat([anchor_xy[:, None], *[p[0] for p in pools]], dim=1).float()
    sc = torch.cat([anchor_score[:, None], *[p[1] for p in pools]], dim=1).float()
    valid = torch.isfinite(xy).all(-1)
    key = torch.where(valid, sc, torch.full_like(sc, float("-inf")))
    key[:, 0] = torch.where(
        valid[:, 0], torch.full_like(key[:, 0], float("inf")), key[:, 0]
    )
    order = key.argsort(dim=1, descending=True)
    xy = xy.gather(1, order[..., None].expand(-1, -1, 2))
    sc = sc.gather(1, order)
    valid = valid.gather(1, order)
    # Invalid entries sit far away and never keep, so they neither suppress nor survive.
    far = xy.nan_to_num(nan=1e9)
    d = torch.cdist(far, far)
    keep = valid.clone()
    for j in range(1, xy.shape[1]):
        near = keep[:, :j] & (d[:, :j, j] < radius)
        keep[:, j] &= ~near.any(dim=1)
    rank = keep.cumsum(dim=1) - 1
    sel = keep & (rank < k)
    m = xy.shape[0]
    out_xy = xy.new_full((m, k, 2), float("nan"))
    out_sc = sc.new_zeros((m, k))
    rows, cols = sel.nonzero(as_tuple=True)
    out_xy[rows, rank[rows, cols]] = xy[rows, cols]
    out_sc[rows, rank[rows, cols]] = sc[rows, cols]
    return out_xy, out_sc

split_shared_peaks

split_shared_peaks(
    field: ndarray,
    cells: ndarray,
    scores: ndarray,
    animals: int,
    *,
    sigma: float = LOCATOR_SIGMA_CELLS,
    iterations: int = 20,
) -> tuple[np.ndarray, np.ndarray]

One view's locator peaks, a blob that holds several animals split among them.

A locator for several animals trains each view's field on one sigma-wide Gaussian per animal, renormalized (ADR-0127), so animals far apart are blobs of 1 / animals of the probability each, and animals touching in the view are one stretched blob holding all of it, whose single maximum is one peak. While fewer blobs than animals hold half an animal's share, the heaviest is fit by as many sigma-wide Gaussians as it holds animals (EM on its probability, started across its long axis), whose centers replace its peak, each with the peak's score over their count (ADR-0131). field is the view's (H, W) probability, cells its (k, 2) peaks in field cells and scores theirs, NaN and 0 where missing; the result keeps the shape, extra peaks dropping off the end.

Source code in src/deeperfly/pose2d/candidates.py

def split_shared_peaks(
    field: np.ndarray,
    cells: np.ndarray,
    scores: np.ndarray,
    animals: int,
    *,
    sigma: float = LOCATOR_SIGMA_CELLS,
    iterations: int = 20,
) -> tuple[np.ndarray, np.ndarray]:
    """One view's locator peaks, a blob that holds several animals split among them.

    A locator for several animals trains each view's field on one `sigma`-wide
    Gaussian per animal, renormalized (ADR-0127), so animals far apart are blobs of
    `1 / animals` of the probability each, and animals touching in the view are one
    stretched blob holding all of it, whose single maximum is one peak. While fewer
    blobs than `animals` hold half an animal's share, the heaviest is fit by as many
    `sigma`-wide Gaussians as it holds animals (EM on its probability, started across
    its long axis), whose centers replace its peak, each with the peak's score over
    their count (ADR-0131). `field` is the view's `(H, W)` probability, `cells` its
    `(k, 2)` peaks in field cells and `scores` theirs, NaN and 0 where missing; the
    result keeps the shape, extra peaks dropping off the end.
    """
    cells = np.array(cells, dtype=float)
    scores = np.array(scores, dtype=float)
    k = len(cells)
    if animals < 2 or k == 0:
        return cells, scores
    h, w = field.shape
    yy, xx = np.mgrid[:h, :w]
    valid = np.isfinite(cells).all(axis=-1) & (scores > 0)
    mass = np.zeros(k)
    for j in np.flatnonzero(valid):
        near = (np.abs(xx - cells[j, 0]) <= BLOB_RADIUS_CELLS) & (
            np.abs(yy - cells[j, 1]) <= BLOB_RADIUS_CELLS
        )
        mass[j] = float(field[near].sum())
    share = 1.0 / animals
    holding = int((mass >= 0.5 * share).sum())
    if holding >= animals or not valid.any():
        return cells, scores
    j = int(np.argmax(np.where(valid, mass, -1.0)))
    parts = min(animals - holding + 1, max(1, round(mass[j] / share)))
    if parts < 2:
        return cells, scores
    means = _fit_blob(field, cells[j], parts, sigma, iterations)
    out_c = np.concatenate([means, np.delete(cells, j, axis=0)])[:k]
    out_s = np.concatenate([np.full(parts, scores[j] / parts), np.delete(scores, j)])[
        :k
    ]
    return out_c, out_s

The optimizer schedule: a cosine-with-warmup LR factor and two parameter groups.

A backbone learning-rate multiplier is reachable, 1.0 (uniform) being the tested recipe. Ported from dfpose with the numerics unchanged.

apply_lr

apply_lr(
    optimizer, base_lrs: list[float], factor: float
) -> None

Set each param group's LR to its base times factor (zipped 1:1, in order).

Source code in src/deeperfly/pose2d/training/optim.py

def apply_lr(optimizer, base_lrs: list[float], factor: float) -> None:
    """Set each param group's LR to its base times `factor` (zipped 1:1, in order)."""
    for group, base in zip(optimizer.param_groups, base_lrs, strict=True):
        group["lr"] = base * factor

build_param_groups

build_param_groups(
    model, base_lr: float, *, backbone_lr_mult: float = 1.0
)

The head and backbone.* parameter groups and their base LRs.

The backbone's LR is scaled by backbone_lr_mult; frozen parameters are skipped.

Source code in src/deeperfly/pose2d/training/optim.py

def build_param_groups(model, base_lr: float, *, backbone_lr_mult: float = 1.0):
    """The head and `backbone.*` parameter groups and their base LRs.

    The backbone's LR is scaled by `backbone_lr_mult`; frozen parameters are skipped.
    """
    backbone, head = [], []
    for name, p in model.named_parameters():
        if not p.requires_grad:
            continue
        (backbone if name.startswith("backbone.") else head).append(p)
    base_lrs = [base_lr, base_lr * backbone_lr_mult]
    groups = [
        {"params": head, "lr": base_lrs[0]},
        {"params": backbone, "lr": base_lrs[1]},
    ]
    return groups, base_lrs

cosine_warmup_factor

cosine_warmup_factor(
    step: int,
    total_steps: int,
    warmup_steps: int,
    min_frac: float = 0.01,
) -> float

A linear warmup over warmup_steps, then a cosine decay to min_frac.

Source code in src/deeperfly/pose2d/training/optim.py

def cosine_warmup_factor(
    step: int, total_steps: int, warmup_steps: int, min_frac: float = 0.01
) -> float:
    """A linear warmup over `warmup_steps`, then a cosine decay to `min_frac`."""
    if warmup_steps > 0 and step < warmup_steps:
        return (step + 1) / warmup_steps
    if total_steps <= warmup_steps:
        return 1.0
    progress = (step - warmup_steps) / (total_steps - warmup_steps)
    cos = 0.5 * (1 + math.cos(math.pi * min(progress, 1.0)))
    return min_frac + (1 - min_frac) * cos

Losses over decoded coordinates, whatever head produced them.

coordinate_loss is the direct supervised term and triangulation_loss Lightning Pose 3D's geometric one (ADR-0072); both take decoded, augmentation-inverted coordinates, the geometric one in video pixels.

coordinate_loss

coordinate_loss(pred, gt, weight, *, huber: float = 0.0)

Weight-masked squared (or huber) error between (..., K, 2) coordinates.

A weight == 0 point, whose gt may be NaN, counts in neither numerator nor denominator. Returns (loss, aux) with coord_loss and weight_sum.

Source code in src/deeperfly/pose2d/training/losses.py

def coordinate_loss(pred, gt, weight, *, huber: float = 0.0):
    """Weight-masked squared (or `huber`) error between `(..., K, 2)` coordinates.

    A `weight == 0` point, whose `gt` may be NaN, counts in neither numerator nor
    denominator. Returns `(loss, aux)` with `coord_loss` and `weight_sum`.
    """
    safe_gt = torch.nan_to_num(gt, nan=0.0, posinf=0.0, neginf=0.0)
    d = (pred - safe_gt).norm(dim=-1)  # (..., K)
    if huber <= 0:
        per = d.pow(2)
    else:
        per = torch.where(d <= huber, d.pow(2), huber * (2 * d - huber))
    w_sum = weight.sum()
    loss = (per * weight.to(per.dtype)).sum() / w_sum.clamp_min(1.0)
    return loss, {"coord_loss": loss.detach(), "weight_sum": w_sum.detach()}

triangulate_pairs

triangulate_pairs(xy, proj, pairs, *, ridge: float = 1e-06)

DLT over view pairs -> (B, P, K, 3) world points.

With the homogeneous coordinate fixed at 1, each pair's four rows are a 4x3 least squares, solved through its ridge-regularized normal equations. xy (B, V, K, 2) and proj (B, V, 3, 4) are in video pixels; pairs is (P, 2).

Source code in src/deeperfly/pose2d/training/losses.py

def triangulate_pairs(xy, proj, pairs, *, ridge: float = 1e-6):
    """DLT over view pairs -> `(B, P, K, 3)` world points.

    With the homogeneous coordinate fixed at 1, each pair's four rows are a 4x3 least
    squares, solved through its ridge-regularized normal equations. `xy` `(B, V, K, 2)`
    and `proj` `(B, V, 3, 4)` are in video pixels; `pairs` is `(P, 2)`.
    """
    rows = []
    for side in (pairs[:, 0], pairs[:, 1]):
        o = xy[:, side]  # (B, P, K, 2)
        pm = proj[:, side]  # (B, P, 3, 4)
        p0, p1, p2 = pm[..., 0, :], pm[..., 1, :], pm[..., 2, :]  # (B, P, 4)
        rows.append(o[..., 0:1] * p2[:, :, None, :] - p0[:, :, None, :])
        rows.append(o[..., 1:2] * p2[:, :, None, :] - p1[:, :, None, :])
    a = torch.stack(rows, dim=-2)  # (B, P, K, 4, 4)
    lhs, rhs = a[..., :3], -a[..., 3]
    ata = lhs.transpose(-1, -2) @ lhs
    atb = lhs.transpose(-1, -2) @ rhs.unsqueeze(-1)
    eye = torch.eye(3, device=xy.device, dtype=xy.dtype).expand_as(ata)
    return torch.linalg.solve(ata + ridge * eye, atb).squeeze(-1)

triangulation_loss

triangulation_loss(
    pred_xy,
    gt_xy,
    weight,
    proj,
    *,
    view_valid=None,
    group_ok=None,
    huber: float = 0.0,
    max_world: float = 0.0,
)

LP3D's 3D loss: prediction and ground truth triangulated through the same pairs.

pred_xy and gt_xy are (B, V, K, 2) video pixels, weight (B, V, K) (a pair supervises where both views do), proj (B, V, 3, 4); view_valid and group_ok mask views and groups. huber replaces the square and max_world clips a point's world error (ADR-0072). Returns (loss, stats).

Source code in src/deeperfly/pose2d/training/losses.py

def triangulation_loss(
    pred_xy,
    gt_xy,
    weight,
    proj,
    *,
    view_valid=None,
    group_ok=None,
    huber: float = 0.0,
    max_world: float = 0.0,
):
    """LP3D's 3D loss: prediction and ground truth triangulated through the same pairs.

    `pred_xy` and `gt_xy` are `(B, V, K, 2)` video pixels, `weight` `(B, V, K)` (a pair
    supervises where both views do), `proj` `(B, V, 3, 4)`; `view_valid` and `group_ok`
    mask views and groups. `huber` replaces the square and `max_world` clips a point's
    world error (ADR-0072). Returns `(loss, stats)`.
    """
    _, v, _, _ = pred_xy.shape
    pairs = view_pairs(v).to(pred_xy.device)
    dtype = torch.float32
    # An unweighted cell may carry NaN ground truth; it must not poison the pairs that
    # share its solve (NaN * 0 is NaN), so both sides are made finite before the DLT.
    pred_f = torch.nan_to_num(pred_xy.to(dtype), nan=0.0, posinf=0.0, neginf=0.0)
    gt_f = torch.nan_to_num(gt_xy.to(dtype), nan=0.0, posinf=0.0, neginf=0.0)
    proj = proj.to(dtype)

    w = weight.to(dtype)
    if view_valid is not None:
        w = w * view_valid.to(dtype)[:, :, None]
    pw = w[:, pairs[:, 0]] * w[:, pairs[:, 1]]  # (B, P, K); both views must supervise
    pw = pw * _pair_is_usable(proj, pairs).to(dtype)[..., None]
    if group_ok is not None:
        pw = pw * group_ok.to(dtype)[:, None, None]

    zero = pred_xy.sum() * 0.0
    empty = {
        "tri_loss": zero.detach(),
        "tri_cells": torch.zeros((), device=pred_xy.device),
        "tri_rms": torch.zeros((), device=pred_xy.device),
        "tri_clipped": torch.zeros((), device=pred_xy.device),
    }
    if float(pw.sum()) == 0.0:
        return zero, empty

    # Autocast would recast these matmuls to bf16 whatever the inputs' dtype; the
    # normal equations need float32.
    with torch.autocast(pred_xy.device.type, enabled=False):
        x_pred = triangulate_pairs(pred_f, proj, pairs)
        with torch.no_grad():
            x_gt = triangulate_pairs(gt_f, proj, pairs)
        d = (x_pred - x_gt).norm(dim=-1)  # (B, P, K) world units
    d = torch.where(pw > 0, d, torch.zeros_like(d))

    clipped = torch.zeros((), device=d.device)
    if max_world > 0:
        over = (d > max_world) & (pw > 0)
        clipped = over.to(dtype).sum() / pw.gt(0).to(dtype).sum().clamp(min=1.0)
        d = d.clamp(max=max_world)
    per = (
        d.pow(2)
        if huber <= 0
        else torch.where(d <= huber, d.pow(2), huber * (2 * d - huber))
    )
    denom = pw.sum().clamp(min=1.0)
    loss = (per * pw).sum() / denom
    stats = {
        "tri_loss": loss.detach(),
        "tri_cells": pw.gt(0).to(dtype).sum().detach(),
        "tri_rms": ((d.pow(2) * pw).sum() / denom).sqrt().detach(),
        "tri_clipped": clipped.detach(),
    }
    return loss, stats

view_pairs

view_pairs(n_views: int)

(P, 2) long: every unordered pair of view indices, P = V(V-1)/2.

Source code in src/deeperfly/pose2d/training/losses.py

def view_pairs(n_views: int):
    """`(P, 2)` long: every unordered pair of view indices, `P = V(V-1)/2`."""
    i, j = torch.triu_indices(n_views, n_views, offset=1)
    return torch.stack([i, j], dim=-1)

A learned manifold of fly poses, and the loss that keeps predictions on it.

Ported from Lightning Pose's pca_singleview loss. A wrong-limb attachment is not a small perturbation of a real pose but a configuration off the manifold, which the residual penalizes while barely touching a few-pixel error; it needs no calibration and supervises unlabeled points too. Fitted in a frame normalized for translation and per-axis scale, so it encodes shape.

PosePCA

Bases: Module

Penalizes the part of a predicted pose that leaves the fitted subspace.

The mean squared reconstruction residual in the normalized frame, zero for any pose the labels could have produced.

Source code in src/deeperfly/pose2d/training/pose_pca.py

class PosePCA(nn.Module):
    """Penalizes the part of a predicted pose that leaves the fitted subspace.

    The mean squared reconstruction residual in the normalized frame, zero for any pose
    the labels could have produced.
    """

    def __init__(self, mean, components) -> None:
        """Store the fitted `mean (2K,)` and `components (C, 2K)` as buffers."""
        super().__init__()
        self.register_buffer("mean", torch.as_tensor(mean, dtype=torch.float32))
        self.register_buffer("comps", torch.as_tensor(components, dtype=torch.float32))

    def residual(self, kp, vis=None):
        """`(B, K, 2) -> (B,)` squared distance from the pose manifold."""
        x = normalize_pose(kp, vis).reshape(kp.shape[0], -1)
        centred = x - self.mean
        recon = centred @ self.comps.T @ self.comps
        return ((centred - recon) ** 2).mean(dim=-1)

    def forward(self, kp, vis=None):
        """Mean manifold-residual loss over the batch."""
        return self.residual(kp, vis).mean()

__init__

__init__(mean, components) -> None

Store the fitted mean (2K,) and components (C, 2K) as buffers.

Source code in src/deeperfly/pose2d/training/pose_pca.py

def __init__(self, mean, components) -> None:
    """Store the fitted `mean (2K,)` and `components (C, 2K)` as buffers."""
    super().__init__()
    self.register_buffer("mean", torch.as_tensor(mean, dtype=torch.float32))
    self.register_buffer("comps", torch.as_tensor(components, dtype=torch.float32))

forward

forward(kp, vis=None)

Mean manifold-residual loss over the batch.

Source code in src/deeperfly/pose2d/training/pose_pca.py

def forward(self, kp, vis=None):
    """Mean manifold-residual loss over the batch."""
    return self.residual(kp, vis).mean()

residual

residual(kp, vis=None)

(B, K, 2) -> (B,) squared distance from the pose manifold.

Source code in src/deeperfly/pose2d/training/pose_pca.py

def residual(self, kp, vis=None):
    """`(B, K, 2) -> (B,)` squared distance from the pose manifold."""
    x = normalize_pose(kp, vis).reshape(kp.shape[0], -1)
    centred = x - self.mean
    recon = centred @ self.comps.T @ self.comps
    return ((centred - recon) ** 2).mean(dim=-1)

fit_pose_pca

fit_pose_pca(
    poses,
    n_components: int | None = None,
    var_target: float = 0.99,
)

(N, K, 2) complete normalized poses -> (mean (2K,), comps (C, 2K), evr (C,)).

n_components fixes C, else the smallest C explaining var_target. Only complete poses: an imputed joint would teach the subspace the imputation.

Source code in src/deeperfly/pose2d/training/pose_pca.py

def fit_pose_pca(poses, n_components: int | None = None, var_target: float = 0.99):
    """`(N, K, 2)` complete normalized poses -> `(mean (2K,), comps (C, 2K), evr (C,))`.

    `n_components` fixes C, else the smallest C explaining `var_target`. Only complete
    poses: an imputed joint would teach the subspace the imputation.
    """
    x = np.asarray(poses, np.float64).reshape(len(poses), -1)
    if not np.isfinite(x).all():
        raise ValueError("fit_pose_pca needs complete poses; got non-finite entries")
    mean = x.mean(0)
    _, s, vt = np.linalg.svd(x - mean, full_matrices=False)
    evr = (s**2) / (s**2).sum()
    if n_components is None:
        n_components = int(np.searchsorted(np.cumsum(evr), var_target) + 1)
    n_components = max(1, min(n_components, vt.shape[0]))
    return mean, vt[:n_components], evr[:n_components]

normalize_pose

normalize_pose(kp, vis=None, eps: float = 1e-06)

(B, K, 2) poses centered and scaled to unit RMS radius per axis.

Per axis, since a window resized per axis stretches the animal; rotation stays.

Source code in src/deeperfly/pose2d/training/pose_pca.py

def normalize_pose(kp, vis=None, eps: float = 1e-6):
    """`(B, K, 2)` poses centered and scaled to unit RMS radius per axis.

    Per axis, since a window resized per axis stretches the animal; rotation stays.
    """
    if vis is None:
        w = torch.ones(kp.shape[:-1], dtype=kp.dtype, device=kp.device)
    else:
        w = (vis > 0).to(kp.dtype)
    n = w.sum(dim=-1, keepdim=True).clamp_min(1.0)
    centre = (kp * w[..., None]).sum(dim=-2) / n
    out = kp - centre[..., None, :]
    radius = torch.sqrt((out**2 * w[..., None]).sum(dim=-2) / n)
    return out / radius[..., None, :].clamp_min(eps)

Composable 2D affine augmentation for the detectors, plus photometric jitter.

Every geometric augmentation is one Affine2D: composed with @ into a single map, the image is resampled once, and its inverse() takes a prediction back to video pixels, which the triangulation loss needs. A flip also permutes the keypoint channels (Skeleton.flip_perm), so a mirrored left leg trains the right channel; inverse() undoes the geometry but not the channel order, which is consistent across a group that flips together. Photometric jitter touches no coordinate.

Affine2D dataclass

A 2x3 affine on 2D pixel coordinates: dst = matrix @ [x, y, 1].

a @ b applies b, then a. Pixel (x, y) sits at continuous coordinate (x, y), so a W-wide frame spans [-0.5, W - 0.5]: the align_corners=False convention of grid_sample and of the inference resize, which every constructor honors.

Source code in src/deeperfly/pose2d/training/augment.py

@dataclass(frozen=True)
class Affine2D:
    """A 2x3 affine on 2D pixel coordinates: `dst = matrix @ [x, y, 1]`.

    `a @ b` applies `b`, then `a`. Pixel `(x, y)` sits at continuous coordinate
    `(x, y)`, so a `W`-wide frame spans `[-0.5, W - 0.5]`: the `align_corners=False`
    convention of `grid_sample` and of the inference resize, which every constructor
    honors.
    """

    matrix: torch.Tensor  # (2, 3)

    def _h(self) -> torch.Tensor:
        """The 3x3 homogeneous form."""
        bottom = torch.tensor(
            [0.0, 0.0, 1.0], device=self.matrix.device, dtype=self.matrix.dtype
        )
        return torch.cat([self.matrix, bottom.view(1, 3)], dim=0)

    def __matmul__(self, other: Affine2D) -> Affine2D:
        """`self` after `other`: `(self @ other)(x) == self(other(x))`."""
        return Affine2D((self._h() @ other._h())[:2])

    def inverse(self) -> Affine2D:
        """The affine that undoes this one."""
        return Affine2D(torch.linalg.inv(self._h())[:2])

    def apply(self, points: torch.Tensor) -> torch.Tensor:
        """Transform `(..., 2)` `(x, y)` points, preserving the leading shape.

        In the higher of the point and matrix dtypes, so float64 video coordinates keep
        their precision; NaN propagates.
        """
        lin, off = self.matrix[:, :2], self.matrix[:, 2]
        dt = torch.promote_types(points.dtype, lin.dtype)
        return points.to(dt) @ lin.to(dt).transpose(0, 1) + off.to(dt)

    def to(self, *, device=None, dtype=None) -> Affine2D:
        """This affine on another device or dtype, to compose across sources."""
        return Affine2D(self.matrix.to(device=device, dtype=dtype))

    def warp_image(
        self,
        image: torch.Tensor,
        out_hw: tuple[int, int],
        *,
        mode: str = "bilinear",
        padding_mode: str = "zeros",
    ) -> torch.Tensor:
        """Resample `image` through this affine onto an `out_hw` canvas.

        `grid_sample` with `align_corners=False`; `"zeros"` padding leaves off-canvas
        regions black, as at inference. Returns the input's rank.
        """
        single = image.ndim == 3
        batched = image[None] if single else image
        b, _, h, w = batched.shape
        theta = self.grid_theta((h, w), out_hw).to(
            device=batched.device, dtype=batched.dtype
        )
        grid = F.affine_grid(
            theta.expand(b, 2, 3), (b, batched.shape[1], *out_hw), align_corners=False
        )
        out = F.grid_sample(
            batched, grid, mode=mode, padding_mode=padding_mode, align_corners=False
        )
        return out[0] if single else out

    def grid_theta(
        self, in_hw: tuple[int, int], out_hw: tuple[int, int]
    ) -> torch.Tensor:
        """`affine_grid` theta: output-normalized -> input-normalized coordinates.

        The inverse of this affine between the two images' `[-1, 1]` frames; `in_hw` is
        the canvas sampled, so a frame placed in a bigger buffer changes no pixel.
        """
        dev, dt = self.matrix.device, self.matrix.dtype
        to_in = _px_to_norm(in_hw, dev, dt)
        from_out = _px_to_norm(out_hw, dev, dt).inverse()
        return (to_in @ self.inverse() @ from_out).matrix

    # -- constructors --------------------------------------------------------

    @classmethod
    def translate(cls, tx: float, ty: float, *, device=None, dtype=torch.float32):
        """Shift by `(tx, ty)`."""
        m = torch.eye(2, 3, device=device, dtype=dtype)
        m[0, 2], m[1, 2] = tx, ty
        return cls(m)

    @classmethod
    def scale(
        cls,
        sx: float,
        sy: float,
        *,
        center: tuple[float, float] = (0.0, 0.0),
        device=None,
        dtype=torch.float32,
    ) -> Affine2D:
        """Scale by `(sx, sy)` about `center`."""
        cx, cy = center
        m = torch.tensor(
            [[sx, 0.0, cx * (1 - sx)], [0.0, sy, cy * (1 - sy)]],
            device=device,
            dtype=dtype,
        )
        return cls(m)

    @classmethod
    def rotate(
        cls,
        radians: float,
        *,
        center: tuple[float, float] = (0.0, 0.0),
        device=None,
        dtype=torch.float32,
    ) -> Affine2D:
        """Rotate by `radians` about `center`: +x toward +y, clockwise on screen."""
        c, s = math.cos(radians), math.sin(radians)
        cx, cy = center
        m = torch.tensor(
            [
                [c, -s, cx - c * cx + s * cy],
                [s, c, cy - s * cx - c * cy],
            ],
            device=device,
            dtype=dtype,
        )
        return cls(m)

    @classmethod
    def flip_x(cls, width: float, *, device=None, dtype=torch.float32) -> Affine2D:
        """Horizontal mirror of a `width`-wide frame: `x -> width - 1 - x`.

        Pixel 0 mirrors onto pixel `width - 1`; `width - x` would shift it a pixel.
        """
        m = torch.tensor(
            [[-1.0, 0.0, width - 1.0], [0.0, 1.0, 0.0]], device=device, dtype=dtype
        )
        return cls(m)

    @classmethod
    def resize(
        cls,
        in_hw: tuple[int, int],
        out_hw: tuple[int, int],
        *,
        device=None,
        dtype=torch.float32,
    ) -> Affine2D:
        """Map a full `in_hw` frame onto a full `out_hw` frame, as `F.interpolate` does.

        Pixel centers map as `(x + 0.5) * out / in - 0.5`; a plain scale about the
        origin would be `(out / in - 1) / 2` pixels off.
        """
        h, w = in_hw
        return cls.crop(
            (0.0, 0.0, float(w), float(h)), out_hw, device=device, dtype=dtype
        )

    @classmethod
    def crop(
        cls,
        box: tuple[float, float, float, float],
        out_hw: tuple[int, int],
        *,
        device=None,
        dtype=torch.float32,
    ) -> Affine2D:
        """Map the window `box = (x0, y0, x1, y1)`, slice bounds, onto `out_hw`.

        As `F.interpolate(image[..., y0:y1, x0:x1], out_hw, align_corners=False)`:
        pixel centers map as `(x - x0 + 0.5) * sx - 0.5`.
        """
        x0, y0, x1, y1 = box
        sx, sy = out_hw[1] / (x1 - x0), out_hw[0] / (y1 - y0)
        kw = {"device": device, "dtype": dtype}
        return (
            cls.translate(-0.5, -0.5, **kw)
            @ cls.scale(sx, sy, **kw)
            @ cls.translate(0.5 - x0, 0.5 - y0, **kw)
        )

__matmul__

__matmul__(other: Affine2D) -> Affine2D

self after other: (self @ other)(x) == self(other(x)).

Source code in src/deeperfly/pose2d/training/augment.py

def __matmul__(self, other: Affine2D) -> Affine2D:
    """`self` after `other`: `(self @ other)(x) == self(other(x))`."""
    return Affine2D((self._h() @ other._h())[:2])

apply

apply(points: Tensor) -> torch.Tensor

Transform (..., 2) (x, y) points, preserving the leading shape.

In the higher of the point and matrix dtypes, so float64 video coordinates keep their precision; NaN propagates.

Source code in src/deeperfly/pose2d/training/augment.py

def apply(self, points: torch.Tensor) -> torch.Tensor:
    """Transform `(..., 2)` `(x, y)` points, preserving the leading shape.

    In the higher of the point and matrix dtypes, so float64 video coordinates keep
    their precision; NaN propagates.
    """
    lin, off = self.matrix[:, :2], self.matrix[:, 2]
    dt = torch.promote_types(points.dtype, lin.dtype)
    return points.to(dt) @ lin.to(dt).transpose(0, 1) + off.to(dt)

crop classmethod

crop(
    box: tuple[float, float, float, float],
    out_hw: tuple[int, int],
    *,
    device=None,
    dtype=torch.float32,
) -> Affine2D

Map the window box = (x0, y0, x1, y1), slice bounds, onto out_hw.

As F.interpolate(image[..., y0:y1, x0:x1], out_hw, align_corners=False): pixel centers map as (x - x0 + 0.5) * sx - 0.5.

Source code in src/deeperfly/pose2d/training/augment.py

@classmethod
def crop(
    cls,
    box: tuple[float, float, float, float],
    out_hw: tuple[int, int],
    *,
    device=None,
    dtype=torch.float32,
) -> Affine2D:
    """Map the window `box = (x0, y0, x1, y1)`, slice bounds, onto `out_hw`.

    As `F.interpolate(image[..., y0:y1, x0:x1], out_hw, align_corners=False)`:
    pixel centers map as `(x - x0 + 0.5) * sx - 0.5`.
    """
    x0, y0, x1, y1 = box
    sx, sy = out_hw[1] / (x1 - x0), out_hw[0] / (y1 - y0)
    kw = {"device": device, "dtype": dtype}
    return (
        cls.translate(-0.5, -0.5, **kw)
        @ cls.scale(sx, sy, **kw)
        @ cls.translate(0.5 - x0, 0.5 - y0, **kw)
    )

flip_x classmethod

flip_x(
    width: float, *, device=None, dtype=torch.float32
) -> Affine2D

Horizontal mirror of a width-wide frame: x -> width - 1 - x.

Pixel 0 mirrors onto pixel width - 1; width - x would shift it a pixel.

Source code in src/deeperfly/pose2d/training/augment.py

@classmethod
def flip_x(cls, width: float, *, device=None, dtype=torch.float32) -> Affine2D:
    """Horizontal mirror of a `width`-wide frame: `x -> width - 1 - x`.

    Pixel 0 mirrors onto pixel `width - 1`; `width - x` would shift it a pixel.
    """
    m = torch.tensor(
        [[-1.0, 0.0, width - 1.0], [0.0, 1.0, 0.0]], device=device, dtype=dtype
    )
    return cls(m)

grid_theta

grid_theta(
    in_hw: tuple[int, int], out_hw: tuple[int, int]
) -> torch.Tensor

affine_grid theta: output-normalized -> input-normalized coordinates.

The inverse of this affine between the two images' [-1, 1] frames; in_hw is the canvas sampled, so a frame placed in a bigger buffer changes no pixel.

Source code in src/deeperfly/pose2d/training/augment.py

def grid_theta(
    self, in_hw: tuple[int, int], out_hw: tuple[int, int]
) -> torch.Tensor:
    """`affine_grid` theta: output-normalized -> input-normalized coordinates.

    The inverse of this affine between the two images' `[-1, 1]` frames; `in_hw` is
    the canvas sampled, so a frame placed in a bigger buffer changes no pixel.
    """
    dev, dt = self.matrix.device, self.matrix.dtype
    to_in = _px_to_norm(in_hw, dev, dt)
    from_out = _px_to_norm(out_hw, dev, dt).inverse()
    return (to_in @ self.inverse() @ from_out).matrix

inverse

inverse() -> Affine2D

The affine that undoes this one.

Source code in src/deeperfly/pose2d/training/augment.py

def inverse(self) -> Affine2D:
    """The affine that undoes this one."""
    return Affine2D(torch.linalg.inv(self._h())[:2])

resize classmethod

resize(
    in_hw: tuple[int, int],
    out_hw: tuple[int, int],
    *,
    device=None,
    dtype=torch.float32,
) -> Affine2D

Map a full in_hw frame onto a full out_hw frame, as F.interpolate does.

Pixel centers map as (x + 0.5) * out / in - 0.5; a plain scale about the origin would be (out / in - 1) / 2 pixels off.

Source code in src/deeperfly/pose2d/training/augment.py

@classmethod
def resize(
    cls,
    in_hw: tuple[int, int],
    out_hw: tuple[int, int],
    *,
    device=None,
    dtype=torch.float32,
) -> Affine2D:
    """Map a full `in_hw` frame onto a full `out_hw` frame, as `F.interpolate` does.

    Pixel centers map as `(x + 0.5) * out / in - 0.5`; a plain scale about the
    origin would be `(out / in - 1) / 2` pixels off.
    """
    h, w = in_hw
    return cls.crop(
        (0.0, 0.0, float(w), float(h)), out_hw, device=device, dtype=dtype
    )

rotate classmethod

rotate(
    radians: float,
    *,
    center: tuple[float, float] = (0.0, 0.0),
    device=None,
    dtype=torch.float32,
) -> Affine2D

Rotate by radians about center: +x toward +y, clockwise on screen.

Source code in src/deeperfly/pose2d/training/augment.py

@classmethod
def rotate(
    cls,
    radians: float,
    *,
    center: tuple[float, float] = (0.0, 0.0),
    device=None,
    dtype=torch.float32,
) -> Affine2D:
    """Rotate by `radians` about `center`: +x toward +y, clockwise on screen."""
    c, s = math.cos(radians), math.sin(radians)
    cx, cy = center
    m = torch.tensor(
        [
            [c, -s, cx - c * cx + s * cy],
            [s, c, cy - s * cx - c * cy],
        ],
        device=device,
        dtype=dtype,
    )
    return cls(m)

scale classmethod

scale(
    sx: float,
    sy: float,
    *,
    center: tuple[float, float] = (0.0, 0.0),
    device=None,
    dtype=torch.float32,
) -> Affine2D

Scale by (sx, sy) about center.

Source code in src/deeperfly/pose2d/training/augment.py

@classmethod
def scale(
    cls,
    sx: float,
    sy: float,
    *,
    center: tuple[float, float] = (0.0, 0.0),
    device=None,
    dtype=torch.float32,
) -> Affine2D:
    """Scale by `(sx, sy)` about `center`."""
    cx, cy = center
    m = torch.tensor(
        [[sx, 0.0, cx * (1 - sx)], [0.0, sy, cy * (1 - sy)]],
        device=device,
        dtype=dtype,
    )
    return cls(m)

to

to(*, device=None, dtype=None) -> Affine2D

This affine on another device or dtype, to compose across sources.

Source code in src/deeperfly/pose2d/training/augment.py

def to(self, *, device=None, dtype=None) -> Affine2D:
    """This affine on another device or dtype, to compose across sources."""
    return Affine2D(self.matrix.to(device=device, dtype=dtype))

translate classmethod

translate(
    tx: float,
    ty: float,
    *,
    device=None,
    dtype=torch.float32,
)

Shift by (tx, ty).

Source code in src/deeperfly/pose2d/training/augment.py

@classmethod
def translate(cls, tx: float, ty: float, *, device=None, dtype=torch.float32):
    """Shift by `(tx, ty)`."""
    m = torch.eye(2, 3, device=device, dtype=dtype)
    m[0, 2], m[1, 2] = tx, ty
    return cls(m)

warp_image

warp_image(
    image: Tensor,
    out_hw: tuple[int, int],
    *,
    mode: str = "bilinear",
    padding_mode: str = "zeros",
) -> torch.Tensor

Resample image through this affine onto an out_hw canvas.

grid_sample with align_corners=False; "zeros" padding leaves off-canvas regions black, as at inference. Returns the input's rank.

Source code in src/deeperfly/pose2d/training/augment.py

def warp_image(
    self,
    image: torch.Tensor,
    out_hw: tuple[int, int],
    *,
    mode: str = "bilinear",
    padding_mode: str = "zeros",
) -> torch.Tensor:
    """Resample `image` through this affine onto an `out_hw` canvas.

    `grid_sample` with `align_corners=False`; `"zeros"` padding leaves off-canvas
    regions black, as at inference. Returns the input's rank.
    """
    single = image.ndim == 3
    batched = image[None] if single else image
    b, _, h, w = batched.shape
    theta = self.grid_theta((h, w), out_hw).to(
        device=batched.device, dtype=batched.dtype
    )
    grid = F.affine_grid(
        theta.expand(b, 2, 3), (b, batched.shape[1], *out_hw), align_corners=False
    )
    out = F.grid_sample(
        batched, grid, mode=mode, padding_mode=padding_mode, align_corners=False
    )
    return out[0] if single else out

Augment dataclass

The training-time augmentation, mild by default, as [augment] writes it.

Each field maps onto sample_affine or photometric_jitter.

Source code in src/deeperfly/pose2d/training/augment.py

@dataclass(frozen=True)
class Augment:
    """The training-time augmentation, mild by default, as `[augment]` writes it.

    Each field maps onto [sample_affine][] or [photometric_jitter][].
    """

    degrees: float = 5.0
    scale: tuple[float, float] = (0.9, 1.1)
    aspect: tuple[float, float] = (1.0, 1.0)
    translate: tuple[float, float] = (0.03, 0.03)
    flip_p: float = 0.5
    rig_symmetry: bool = False
    brightness: float = 0.2
    contrast: float = 0.2
    gamma: tuple[float, float] | None = None
    blur_sigma: tuple[float, float] | None = None
    noise_std: float = 0.0
    photometric_p: float = 0.5

    def photometric(self):
        """The torchvision transform this spec's image-only jitter builds."""
        return photometric_jitter(
            brightness=self.brightness,
            contrast=self.contrast,
            gamma=self.gamma,
            blur_sigma=self.blur_sigma,
            noise_std=self.noise_std,
            p=self.photometric_p,
        )

photometric

photometric()

The torchvision transform this spec's image-only jitter builds.

Source code in src/deeperfly/pose2d/training/augment.py

def photometric(self):
    """The torchvision transform this spec's image-only jitter builds."""
    return photometric_jitter(
        brightness=self.brightness,
        contrast=self.contrast,
        gamma=self.gamma,
        blur_sigma=self.blur_sigma,
        noise_std=self.noise_std,
        p=self.photometric_p,
    )

flip_channels

flip_channels(kp, vis, flip_perm)

Reorder (..., K, 2) keypoints and (..., K) weights by flip_perm.

After a mirror a left joint sits where the right one belongs; flip_perm is an involution, the identity on unpaired points.

Source code in src/deeperfly/pose2d/training/augment.py

def flip_channels(kp, vis, flip_perm):
    """Reorder `(..., K, 2)` keypoints and `(..., K)` weights by `flip_perm`.

    After a mirror a left joint sits where the right one belongs; `flip_perm` is an
    involution, the identity on unpaired points.
    """
    perm = torch.as_tensor(flip_perm, device=kp.device, dtype=torch.long)
    return kp.index_select(-2, perm), vis.index_select(-1, perm)

photometric_jitter

photometric_jitter(
    *,
    brightness: float = 0.0,
    contrast: float = 0.0,
    gamma: tuple[float, float] | None = None,
    blur_sigma: tuple[float, float] | None = None,
    blur_kernel: int = 7,
    noise_std: float = 0.0,
    p: float = 0.5,
) -> v2.Transform

A torchvision transforms.v2 pipeline of image-only, coordinate-free jitter.

Each enabled stage fires with probability p: ColorJitter brightness and contrast (safe on gray), a log-uniform gamma, a Gaussian blur (blur_kernel odd) and additive noise on a [0, 1] image. torchvision draws once per call, so apply it per sample, before collation.

Source code in src/deeperfly/pose2d/training/augment.py

def photometric_jitter(
    *,
    brightness: float = 0.0,
    contrast: float = 0.0,
    gamma: tuple[float, float] | None = None,
    blur_sigma: tuple[float, float] | None = None,
    blur_kernel: int = 7,
    noise_std: float = 0.0,
    p: float = 0.5,
) -> v2.Transform:
    """A torchvision `transforms.v2` pipeline of image-only, coordinate-free jitter.

    Each enabled stage fires with probability `p`: `ColorJitter` brightness and
    contrast (safe on gray), a log-uniform `gamma`, a Gaussian blur (`blur_kernel` odd)
    and additive noise on a `[0, 1]` image. torchvision draws once per call, so apply
    it per sample, before collation.
    """
    stages: list[v2.Transform] = []
    if brightness or contrast:
        stages.append(v2.ColorJitter(brightness=brightness, contrast=contrast))
    if gamma is not None:
        stages.append(_RandomGamma(*gamma))
    if blur_sigma is not None:
        stages.append(v2.GaussianBlur(kernel_size=blur_kernel, sigma=blur_sigma))
    if noise_std:
        stages.append(v2.GaussianNoise(sigma=noise_std))
    if not stages:
        return v2.Identity()
    return v2.Compose([v2.RandomApply([s], p=p) for s in stages])

sample_affine

sample_affine(
    hw: tuple[int, int],
    *,
    degrees: float = 0.0,
    scale: tuple[float, float] = (1.0, 1.0),
    aspect: tuple[float, float] = (1.0, 1.0),
    translate: tuple[float, float] = (0.0, 0.0),
    flip_p: float = 0.0,
    generator: Generator | None = None,
    device=None,
    dtype=torch.float32,
) -> tuple[Affine2D, bool]

A random geometric augmentation on an hw canvas: one affine and a flip flag.

An optional flip first, then rotation ([-degrees, degrees]), isotropic scale and an area-preserving, log-uniform aspect jitter about the canvas center, then a translation (a fraction of the canvas). When the flag is set the caller must permute the keypoint channels (flip_channels). Compose the affine with Affine2D.crop or resize to start from video pixels.

Source code in src/deeperfly/pose2d/training/augment.py

def sample_affine(
    hw: tuple[int, int],
    *,
    degrees: float = 0.0,
    scale: tuple[float, float] = (1.0, 1.0),
    aspect: tuple[float, float] = (1.0, 1.0),
    translate: tuple[float, float] = (0.0, 0.0),
    flip_p: float = 0.0,
    generator: torch.Generator | None = None,
    device=None,
    dtype=torch.float32,
) -> tuple[Affine2D, bool]:
    """A random geometric augmentation on an `hw` canvas: one affine and a flip flag.

    An optional flip first, then rotation (`[-degrees, degrees]`), isotropic scale and
    an area-preserving, log-uniform aspect jitter about the canvas center, then a
    translation (a fraction of the canvas). When the flag is set the caller must
    permute the keypoint channels (`flip_channels`). Compose the affine with
    `Affine2D.crop` or `resize` to start from video pixels.
    """
    h, w = hw
    center = ((w - 1) / 2.0, (h - 1) / 2.0)  # the canvas centre in pixel-centre coords
    angle = math.radians(_uniform(-degrees, degrees, generator))
    s = _uniform(scale[0], scale[1], generator)
    a = math.exp(_uniform(math.log(aspect[0]), math.log(aspect[1]), generator))
    tx = _uniform(-translate[0], translate[0], generator) * w
    ty = _uniform(-translate[1], translate[1], generator) * h
    flipped = bool(torch.rand((), generator=generator) < flip_p)

    kw = {"device": device, "dtype": dtype}
    aff = Affine2D.translate(tx, ty, **kw)
    aff = aff @ Affine2D.rotate(angle, center=center, **kw)
    root = math.sqrt(a)
    aff = aff @ Affine2D.scale(s * root, s / root, center=center, **kw)
    if flipped:
        aff = aff @ Affine2D.flip_x(w, **kw)
    return aff, flipped

warp_sample

warp_sample(
    image,
    kp,
    vis,
    affine: Affine2D,
    out_hw: tuple[int, int],
    *,
    flipped: bool = False,
    flip_perm=None,
    mode: str = "bilinear",
    padding_mode: str = "zeros",
)

Apply one geometric augmentation to an (image, kp, vis) sample.

Warps the image and the source-pixel keypoints through affine, permuting the channels by flip_perm when flipped; coordinates stay in out_hw pixels. Returns (image, kp, vis) in the augmented frame.

Source code in src/deeperfly/pose2d/training/augment.py

def warp_sample(
    image,
    kp,
    vis,
    affine: Affine2D,
    out_hw: tuple[int, int],
    *,
    flipped: bool = False,
    flip_perm=None,
    mode: str = "bilinear",
    padding_mode: str = "zeros",
):
    """Apply one geometric augmentation to an `(image, kp, vis)` sample.

    Warps the image and the source-pixel keypoints through `affine`, permuting the
    channels by `flip_perm` when `flipped`; coordinates stay in `out_hw` pixels.
    Returns `(image, kp, vis)` in the augmented frame.
    """
    image = affine.warp_image(image, out_hw, mode=mode, padding_mode=padding_mode)
    kp = affine.apply(kp)
    if flipped:
        if flip_perm is None:
            raise ValueError("flip_perm is required when flipped is True")
        kp, vis = flip_channels(kp, vis, flip_perm)
    return image, kp, vis

The frame cache: a project's labeled frames, decoded once, on disk.

Labels are sparse (a few dozen frames of thousands), so each is decoded once, cropped and resized to the input by the detectors' own anti-aliased resize, and kept as grayscale uint8 in an HDF5 file keyed by the crop and input size.

FrameCache

The labeled frames of one project, on disk.

One HDF5 file per (project, crops, input_hw, videos), its digest in the name, so a changed crop plan or video never reads a stale cache. With input_hw each view's window is stored resized to the input; without it, the source mode stores each video's native frame once and the window mode (with crops) each view's window at video resolution. Nothing is flipped on the way in: a mirrored view's crop is reflected by the sample's window, as at inference. ensure decodes what is missing.

Source code in src/deeperfly/pose2d/training/cache.py

class FrameCache:
    """The labeled frames of one project, on disk.

    One HDF5 file per `(project, crops, input_hw, videos)`, its digest in the name,
    so a changed crop plan or video never reads a stale cache. With `input_hw` each
    view's window is stored resized to the input; without it, the *source* mode stores
    each video's native frame once and the *window* mode (with `crops`) each view's
    window at video resolution. Nothing is flipped on the way in: a mirrored view's crop
    is reflected by the sample's window, as at inference. `ensure` decodes what is
    missing.
    """

    def __init__(
        self,
        project: Project,
        name: str,
        *,
        input_hw: tuple[int, int] | None,
        crops: dict[str, tuple[int, int, int, int]] | None = None,
        cache_dir: str | Path | None = None,
        frame_hw: Mapping[str, tuple[int, int]] | None = None,
    ) -> None:
        """Open (or create) the cache of `project` for `input_hw` and `crops`.

        `name` is the project's name in the training set, the cache's directory
        under `<cache_dir>/frames/` (`<project>/cache` by default); `frame_hw` is
        each view's frame `(height, width)`, probed from the videos when absent.
        """
        self.project = project
        self.name = name
        self.layout = project.layout
        self._frame_hw = None if frame_hw is None else dict(frame_hw)
        self.input_hw = (
            None if input_hw is None else (int(input_hw[0]), int(input_hw[1]))
        )
        self.view_names = tuple(project.views)
        self.crops = crops
        if self.input_hw is not None and crops is None:
            self.crops = full_frame_crops(project, self.frame_hw)
        #: Window mode: the clipped `(x0, y0, x1, y1)` window per view, filled in by
        #: `ensure` and persisted in the file so `window_affine` needs no video access.
        self.windows: dict[str, tuple[int, int, int, int]] = {}
        #: The dataset each view is stored in: its video in the source mode, itself
        #: otherwise.
        self._store = {
            view: (self.layout.video_of(view) if self.source_mode else view)
            for view in self.view_names
        }
        #: `(height, width)` of each stored dataset, from the file or from `ensure`.
        self.sizes: dict[str, tuple[int, int]] = {}
        root = Path(cache_dir) if cache_dir is not None else project.folder / "cache"
        self.path = root / "frames" / name / f"{self.digest()}.h5"
        self._index: dict[int, int] = {}
        if self.path.exists():
            with h5py.File(self.path, "r") as f:
                self._index = {int(t): i for i, t in enumerate(f["frame"][:])}
                if "windows" in f.attrs:
                    self.windows = {
                        k: tuple(int(x) for x in v)
                        for k, v in json.loads(f.attrs["windows"]).items()
                    }
                if "frames" in f:
                    self.sizes = {
                        k: (int(v.shape[1]), int(v.shape[2]))
                        for k, v in f["frames"].items()
                    }

    @property
    def frame_hw(self) -> dict[str, tuple[int, int]]:
        """Each view's frame `(height, width)`, probed once when not given."""
        if self._frame_hw is None:
            self._frame_hw = self.project.frame_sizes()
        return self._frame_hw

    def digest(self) -> str:
        """What the cache is keyed by: input size, per-view crops and video paths.

        A video path is keyed relative to the project folder, both resolved up to the
        file itself, so moving or relinking a project's tree keeps its caches
        (ADR-0124); a video that is a symlink keeps its own name.
        """
        root = os.path.realpath(self.project.folder)
        videos = {
            view: [
                os.path.relpath(os.path.join(os.path.realpath(p.parent), p.name), root)
                for p in self.project.video_paths(self._video_of(view))
            ]
            for view in self.view_names
        }
        key: dict = {
            "version": FRAME_CACHE_VERSION,
            "input_hw": None if self.input_hw is None else list(self.input_hw),
            "crops": None
            if self.crops is None
            else {n: list(self.crops[n]) for n in self.view_names},
            "videos": videos,
        }
        # A view that shares a video is keyed by the dataset it shares and by
        # whether the apparatus mirrors it; a 1:1 rig adds neither key.
        mirrored = [v for v in self.view_names if self.layout.is_mirrored(v)]
        if mirrored:
            key["mirrored"] = mirrored
        if any(store != name for name, store in self._store.items()):
            key["stores"] = dict(self._store)
        key = json.dumps(key, sort_keys=True)
        return hashlib.sha1(key.encode()).hexdigest()[:16]

    def crop_affine(self, view: int) -> Affine2D:
        """Video pixels -> input pixels for one view, in the resizing mode only."""
        if self.crops is None or self.input_hw is None:
            raise ValueError("only a cache resized to an input has a crop affine")
        return Affine2D.crop(
            tuple(float(v) for v in self.crops[self.view_names[view]]), self.input_hw
        )

    @property
    def window_mode(self) -> bool:
        """Whether frames are stored as native-resolution windows."""
        return self.input_hw is None and self.crops is not None

    @property
    def source_mode(self) -> bool:
        """Whether frames are stored whole, one per source (see the class docs)."""
        return self.input_hw is None and self.crops is None

    def window_affine(self, view: int) -> Affine2D:
        """Video pixels -> stored-window pixels (a translation), in window mode."""
        if not self.window_mode:
            raise ValueError("window_affine is for the windowed native mode")
        name = self.view_names[view]
        if name not in self.windows:
            raise ValueError(f"window of {name} not yet known; call ensure first")
        x0, y0, _, _ = self.windows[name]
        return Affine2D.translate(-float(x0), -float(y0))

    def image_affine(self, view: int) -> Affine2D | None:
        """Video pixels -> stored-image pixels for one view; `None` when they are one.

        The crop-and-resize, the window's translation, or, in the source mode, nothing:
        the stored image is the view's video frame.
        """
        if self.input_hw is not None:
            return self.crop_affine(view)
        if self.window_mode:
            return self.window_affine(view)
        return None

    def _stores(self) -> dict[str, str]:
        """Dataset name -> the view whose video it holds, in view order."""
        out: dict[str, str] = {}
        for name in self.view_names:
            out.setdefault(self._store[name], name)
        return out

    def _open(self, store: str):
        """The reader a stored image comes from; nothing is flipped (ADR-0053)."""
        video = store if self.source_mode else self._video_of(store)
        return open_reader(self.project.video_paths(video))

    def _video_of(self, view: str) -> str:
        """The video a view's pixels come from."""
        return self.layout.video_of(view)

    def mirrored(self, view: int) -> bool:
        """Whether the apparatus reflects this view, so its crop is mirrored."""
        return self.layout.is_mirrored(self.view_names[view])

    def stored_hw(self, view: int) -> tuple[int, int]:
        """The `(height, width)` of the image one view is read out of."""
        return self.sizes[self._store[self.view_names[view]]]

    @property
    def frames(self) -> list[int]:
        """The cached frame indices, sorted."""
        return sorted(self._index)

    def ensure(self, frames: Sequence[int]) -> None:
        """Decode every frame of `frames` not yet cached and append it."""
        missing = sorted({int(t) for t in frames} - set(self._index))
        if not missing:
            return
        from ..infer import prepare_images

        native = self.input_hw is None
        stores = self._stores()
        sizes = {}
        for name, view in stores.items():
            if not native:
                sizes[name] = self.input_hw
                continue
            # A stored video's size is that of any view it carries.
            fh, fw = self.frame_hw[view]
            if self.crops is None:
                sizes[name] = (int(fh), int(fw))
                continue
            x0, y0, x1, y1 = (round(float(v)) for v in self.crops[view])
            win = (max(0, x0), max(0, y0), min(int(fw), x1), min(int(fh), y1))
            if win[2] <= win[0] or win[3] <= win[1]:
                raise ValueError(
                    f"window {self.crops[view]} of {view} lies outside its "
                    f"{fw}x{fh} frame"
                )
            self.windows[view] = win
            sizes[name] = (win[3] - win[1], win[2] - win[0])
        self.sizes = dict(sizes)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        with h5py.File(self.path, "a") as f:
            if "frame" not in f:
                f.attrs["format_version"] = FRAME_CACHE_VERSION
                f.attrs["view_names"] = json.dumps(list(self.view_names))
                if self.window_mode:
                    f.attrs["windows"] = json.dumps(self.windows)
                f.create_dataset("frame", shape=(0,), maxshape=(None,), dtype=np.int32)
                for name in stores:
                    h, w = sizes[name]
                    f.create_dataset(
                        f"frames/{name}",
                        shape=(0, h, w),
                        maxshape=(None, h, w),
                        dtype=np.uint8,
                        chunks=(1, h, w),
                    )
            n0 = f["frame"].shape[0]
            f["frame"].resize((n0 + len(missing),))
            f["frame"][n0:] = np.asarray(missing, dtype=np.int32)
            for name, view in stores.items():
                h, w = sizes[name]
                ds = f[f"frames/{name}"]
                ds.resize((n0 + len(missing), h, w))
                with self._open(name if self.source_mode else view) as r:
                    for start in range(0, len(missing), _DECODE_CHUNK):
                        chunk = missing[start : start + _DECODE_CHUNK]
                        raw = r[chunk]  # (T, H, W, C)
                        if native and self.crops is not None:
                            x0, y0, x1, y1 = self.windows[view]
                            gray = np.asarray(raw)[:, y0:y1, x0:x1, 0]
                        elif native:
                            # Monochrome video: the first plane is the luma.
                            gray = np.asarray(raw)[..., 0]
                        else:
                            x0, y0, x1, y1 = self.crops[view]
                            img = prepare_images(
                                raw[:, y0:y1, x0:x1], self.input_hw, "cpu"
                            )
                            gray = (
                                (img[:, 0] * 255.0)
                                .round()
                                .clamp(0, 255)
                                .to(torch.uint8)
                            ).numpy()
                        ds[n0 + start : n0 + start + len(chunk)] = gray
        for i, t in enumerate(missing):
            self._index[t] = n0 + i

    def get(self, view: int, frame: int) -> np.ndarray:
        """The cached `(H, W)` uint8 image one view of one frame is drawn from."""
        with h5py.File(self.path, "r") as f:
            name = self._store[self.view_names[view]]
            return f[f"frames/{name}"][self._index[int(frame)]]

frame_hw property

frame_hw: dict[str, tuple[int, int]]

Each view's frame (height, width), probed once when not given.

frames property

frames: list[int]

The cached frame indices, sorted.

source_mode property

source_mode: bool

Whether frames are stored whole, one per source (see the class docs).

window_mode property

window_mode: bool

Whether frames are stored as native-resolution windows.

__init__

__init__(
    project: Project,
    name: str,
    *,
    input_hw: tuple[int, int] | None,
    crops: dict[str, tuple[int, int, int, int]]
    | None = None,
    cache_dir: str | Path | None = None,
    frame_hw: Mapping[str, tuple[int, int]] | None = None,
) -> None

Open (or create) the cache of project for input_hw and crops.

name is the project's name in the training set, the cache's directory under <cache_dir>/frames/ (<project>/cache by default); frame_hw is each view's frame (height, width), probed from the videos when absent.

Source code in src/deeperfly/pose2d/training/cache.py

def __init__(
    self,
    project: Project,
    name: str,
    *,
    input_hw: tuple[int, int] | None,
    crops: dict[str, tuple[int, int, int, int]] | None = None,
    cache_dir: str | Path | None = None,
    frame_hw: Mapping[str, tuple[int, int]] | None = None,
) -> None:
    """Open (or create) the cache of `project` for `input_hw` and `crops`.

    `name` is the project's name in the training set, the cache's directory
    under `<cache_dir>/frames/` (`<project>/cache` by default); `frame_hw` is
    each view's frame `(height, width)`, probed from the videos when absent.
    """
    self.project = project
    self.name = name
    self.layout = project.layout
    self._frame_hw = None if frame_hw is None else dict(frame_hw)
    self.input_hw = (
        None if input_hw is None else (int(input_hw[0]), int(input_hw[1]))
    )
    self.view_names = tuple(project.views)
    self.crops = crops
    if self.input_hw is not None and crops is None:
        self.crops = full_frame_crops(project, self.frame_hw)
    #: Window mode: the clipped `(x0, y0, x1, y1)` window per view, filled in by
    #: `ensure` and persisted in the file so `window_affine` needs no video access.
    self.windows: dict[str, tuple[int, int, int, int]] = {}
    #: The dataset each view is stored in: its video in the source mode, itself
    #: otherwise.
    self._store = {
        view: (self.layout.video_of(view) if self.source_mode else view)
        for view in self.view_names
    }
    #: `(height, width)` of each stored dataset, from the file or from `ensure`.
    self.sizes: dict[str, tuple[int, int]] = {}
    root = Path(cache_dir) if cache_dir is not None else project.folder / "cache"
    self.path = root / "frames" / name / f"{self.digest()}.h5"
    self._index: dict[int, int] = {}
    if self.path.exists():
        with h5py.File(self.path, "r") as f:
            self._index = {int(t): i for i, t in enumerate(f["frame"][:])}
            if "windows" in f.attrs:
                self.windows = {
                    k: tuple(int(x) for x in v)
                    for k, v in json.loads(f.attrs["windows"]).items()
                }
            if "frames" in f:
                self.sizes = {
                    k: (int(v.shape[1]), int(v.shape[2]))
                    for k, v in f["frames"].items()
                }

crop_affine

crop_affine(view: int) -> Affine2D

Video pixels -> input pixels for one view, in the resizing mode only.

Source code in src/deeperfly/pose2d/training/cache.py

def crop_affine(self, view: int) -> Affine2D:
    """Video pixels -> input pixels for one view, in the resizing mode only."""
    if self.crops is None or self.input_hw is None:
        raise ValueError("only a cache resized to an input has a crop affine")
    return Affine2D.crop(
        tuple(float(v) for v in self.crops[self.view_names[view]]), self.input_hw
    )

digest

digest() -> str

What the cache is keyed by: input size, per-view crops and video paths.

A video path is keyed relative to the project folder, both resolved up to the file itself, so moving or relinking a project's tree keeps its caches (ADR-0124); a video that is a symlink keeps its own name.

Source code in src/deeperfly/pose2d/training/cache.py

def digest(self) -> str:
    """What the cache is keyed by: input size, per-view crops and video paths.

    A video path is keyed relative to the project folder, both resolved up to the
    file itself, so moving or relinking a project's tree keeps its caches
    (ADR-0124); a video that is a symlink keeps its own name.
    """
    root = os.path.realpath(self.project.folder)
    videos = {
        view: [
            os.path.relpath(os.path.join(os.path.realpath(p.parent), p.name), root)
            for p in self.project.video_paths(self._video_of(view))
        ]
        for view in self.view_names
    }
    key: dict = {
        "version": FRAME_CACHE_VERSION,
        "input_hw": None if self.input_hw is None else list(self.input_hw),
        "crops": None
        if self.crops is None
        else {n: list(self.crops[n]) for n in self.view_names},
        "videos": videos,
    }
    # A view that shares a video is keyed by the dataset it shares and by
    # whether the apparatus mirrors it; a 1:1 rig adds neither key.
    mirrored = [v for v in self.view_names if self.layout.is_mirrored(v)]
    if mirrored:
        key["mirrored"] = mirrored
    if any(store != name for name, store in self._store.items()):
        key["stores"] = dict(self._store)
    key = json.dumps(key, sort_keys=True)
    return hashlib.sha1(key.encode()).hexdigest()[:16]

ensure

ensure(frames: Sequence[int]) -> None

Decode every frame of frames not yet cached and append it.

Source code in src/deeperfly/pose2d/training/cache.py

def ensure(self, frames: Sequence[int]) -> None:
    """Decode every frame of `frames` not yet cached and append it."""
    missing = sorted({int(t) for t in frames} - set(self._index))
    if not missing:
        return
    from ..infer import prepare_images

    native = self.input_hw is None
    stores = self._stores()
    sizes = {}
    for name, view in stores.items():
        if not native:
            sizes[name] = self.input_hw
            continue
        # A stored video's size is that of any view it carries.
        fh, fw = self.frame_hw[view]
        if self.crops is None:
            sizes[name] = (int(fh), int(fw))
            continue
        x0, y0, x1, y1 = (round(float(v)) for v in self.crops[view])
        win = (max(0, x0), max(0, y0), min(int(fw), x1), min(int(fh), y1))
        if win[2] <= win[0] or win[3] <= win[1]:
            raise ValueError(
                f"window {self.crops[view]} of {view} lies outside its "
                f"{fw}x{fh} frame"
            )
        self.windows[view] = win
        sizes[name] = (win[3] - win[1], win[2] - win[0])
    self.sizes = dict(sizes)
    self.path.parent.mkdir(parents=True, exist_ok=True)
    with h5py.File(self.path, "a") as f:
        if "frame" not in f:
            f.attrs["format_version"] = FRAME_CACHE_VERSION
            f.attrs["view_names"] = json.dumps(list(self.view_names))
            if self.window_mode:
                f.attrs["windows"] = json.dumps(self.windows)
            f.create_dataset("frame", shape=(0,), maxshape=(None,), dtype=np.int32)
            for name in stores:
                h, w = sizes[name]
                f.create_dataset(
                    f"frames/{name}",
                    shape=(0, h, w),
                    maxshape=(None, h, w),
                    dtype=np.uint8,
                    chunks=(1, h, w),
                )
        n0 = f["frame"].shape[0]
        f["frame"].resize((n0 + len(missing),))
        f["frame"][n0:] = np.asarray(missing, dtype=np.int32)
        for name, view in stores.items():
            h, w = sizes[name]
            ds = f[f"frames/{name}"]
            ds.resize((n0 + len(missing), h, w))
            with self._open(name if self.source_mode else view) as r:
                for start in range(0, len(missing), _DECODE_CHUNK):
                    chunk = missing[start : start + _DECODE_CHUNK]
                    raw = r[chunk]  # (T, H, W, C)
                    if native and self.crops is not None:
                        x0, y0, x1, y1 = self.windows[view]
                        gray = np.asarray(raw)[:, y0:y1, x0:x1, 0]
                    elif native:
                        # Monochrome video: the first plane is the luma.
                        gray = np.asarray(raw)[..., 0]
                    else:
                        x0, y0, x1, y1 = self.crops[view]
                        img = prepare_images(
                            raw[:, y0:y1, x0:x1], self.input_hw, "cpu"
                        )
                        gray = (
                            (img[:, 0] * 255.0)
                            .round()
                            .clamp(0, 255)
                            .to(torch.uint8)
                        ).numpy()
                    ds[n0 + start : n0 + start + len(chunk)] = gray
    for i, t in enumerate(missing):
        self._index[t] = n0 + i

get

get(view: int, frame: int) -> np.ndarray

The cached (H, W) uint8 image one view of one frame is drawn from.

Source code in src/deeperfly/pose2d/training/cache.py

def get(self, view: int, frame: int) -> np.ndarray:
    """The cached `(H, W)` uint8 image one view of one frame is drawn from."""
    with h5py.File(self.path, "r") as f:
        name = self._store[self.view_names[view]]
        return f[f"frames/{name}"][self._index[int(frame)]]

image_affine

image_affine(view: int) -> Affine2D | None

Video pixels -> stored-image pixels for one view; None when they are one.

The crop-and-resize, the window's translation, or, in the source mode, nothing: the stored image is the view's video frame.

Source code in src/deeperfly/pose2d/training/cache.py

def image_affine(self, view: int) -> Affine2D | None:
    """Video pixels -> stored-image pixels for one view; `None` when they are one.

    The crop-and-resize, the window's translation, or, in the source mode, nothing:
    the stored image is the view's video frame.
    """
    if self.input_hw is not None:
        return self.crop_affine(view)
    if self.window_mode:
        return self.window_affine(view)
    return None

mirrored

mirrored(view: int) -> bool

Whether the apparatus reflects this view, so its crop is mirrored.

Source code in src/deeperfly/pose2d/training/cache.py

def mirrored(self, view: int) -> bool:
    """Whether the apparatus reflects this view, so its crop is mirrored."""
    return self.layout.is_mirrored(self.view_names[view])

stored_hw

stored_hw(view: int) -> tuple[int, int]

The (height, width) of the image one view is read out of.

Source code in src/deeperfly/pose2d/training/cache.py

def stored_hw(self, view: int) -> tuple[int, int]:
    """The `(height, width)` of the image one view is read out of."""
    return self.sizes[self._store[self.view_names[view]]]

window_affine

window_affine(view: int) -> Affine2D

Video pixels -> stored-window pixels (a translation), in window mode.

Source code in src/deeperfly/pose2d/training/cache.py

def window_affine(self, view: int) -> Affine2D:
    """Video pixels -> stored-window pixels (a translation), in window mode."""
    if not self.window_mode:
        raise ValueError("window_affine is for the windowed native mode")
    name = self.view_names[view]
    if name not in self.windows:
        raise ValueError(f"window of {name} not yet known; call ensure first")
    x0, y0, _, _ = self.windows[name]
    return Affine2D.translate(-float(x0), -float(y0))

full_frame_crops

full_frame_crops(
    project: Project,
    frame_hw: Mapping[str, tuple[int, int]] | None = None,
) -> dict[str, tuple[int, int, int, int]]

The whole-frame crop window (x0, y0, x1, y1) of every view of a project.

frame_hw is each view's (height, width); without it the videos are probed.

Source code in src/deeperfly/pose2d/training/cache.py

def full_frame_crops(
    project: Project, frame_hw: Mapping[str, tuple[int, int]] | None = None
) -> dict[str, tuple[int, int, int, int]]:
    """The whole-frame crop window `(x0, y0, x1, y1)` of every view of a project.

    `frame_hw` is each view's `(height, width)`; without it the videos are probed.
    """
    sizes = project.frame_sizes() if frame_hw is None else frame_hw
    return {
        view: (0, 0, int(sizes[view][1]), int(sizes[view][0])) for view in project.views
    }

Training samples from projects' labels: per-view samples and view groups.

LabelDataset is one item per (project, frame, view), each with its image (from a FrameCache), input-normalized keypoints, a per-keypoint weight and the geometry to invert it to video pixels. A cell weighs 0 when unlabeled, hidden, absent or past the detector's off-frame margin; a cell inside the margin keeps its weight, so the padded field learns to extrapolate. GroupedLabelDataset carries every view of one moment together for the multiview transformer, with one flip decision per group, since a half-mirrored rig is not a rig.

A training set is several projects (ADR-0085), given as name -> Project in the order the collection names them (project_names); the name is what a split, a report and the frame cache call the project. A sample's project index is the project's position in that mapping.

GroupedLabelDataset

Bases: Dataset

All views of one (project, frame) as (V, ...) tensors: the multiview shape.

One flip decision per group, with the targets permuted as the predictions are, so a flipped group still feeds the geometric loss. With view_drop_p an augmented group keeps a uniform number of views in [view_min, V], the dropped slots invalid, so the model trains at every view count it may be served.

Source code in src/deeperfly/pose2d/training/dataset.py

class GroupedLabelDataset(Dataset):
    """All views of one `(project, frame)` as `(V, ...)` tensors: the multiview shape.

    One flip decision per group, with the targets permuted as the predictions are, so a
    flipped group still feeds the geometric loss. With `view_drop_p` an augmented group
    keeps a uniform number of views in `[view_min, V]`, the dropped slots invalid, so
    the model trains at every view count it may be served.
    """

    def __init__(
        self, base: SampleDataset, *, view_drop_p: float = 0.0, view_min: int = 1
    ) -> None:
        """Group the items of `base` by `(project, frame)`."""
        self.base = base
        self.groups = base.groups()
        self.view_drop_p = float(view_drop_p)
        self.view_min = int(view_min)

    def __len__(self) -> int:
        """The number of `(project, frame)` groups."""
        return len(self.groups)

    def n_views(self) -> list[int]:
        """The view count of every group, for bucketing."""
        return [len(g) for g in self.groups]

    def group_projects(self) -> list[str]:
        """The project name of every group."""
        return [self.base.project_name(self.base.items[g[0]][0]) for g in self.groups]

    def __getitem__(self, g: int) -> dict:
        """Group `g` as stacked per-view tensors, one flip decision for all views."""
        flip = None
        turns: list = [None] * len(self.groups[g])
        if self.base.augment is not None:
            flip = bool(torch.rand(()) < self.base.augment.flip_p)
            symmetries = (
                self.base.symmetries(self.base.items[self.groups[g][0]][0])
                if self.base.augment.rig_symmetry
                else []
            )
            if symmetries:
                flip, signs = symmetries[int(torch.randint(len(symmetries), ()))]
                turns = [signs[self.base.items[i][2]] for i in self.groups[g]]
        views = [
            self.base.sample(i, flip=flip, turn=turn)
            for i, turn in zip(self.groups[g], turns, strict=True)
        ]
        out = {
            key: torch.stack([s[key] for s in views])
            for key in (
                "image",
                "kp",
                "weight",
                "in_frame",
                "affine",
                "gt_video_px",
                "proj",
                "centers",
                "center_valid",
            )
        }
        n = len(views)
        valid = torch.ones(n, dtype=torch.bool)
        if (
            self.base.augment is not None
            and self.view_drop_p > 0
            and n > self.view_min
            and bool(torch.rand(()) < self.view_drop_p)
        ):
            keep = int(torch.randint(self.view_min, n + 1, ()))
            valid[torch.randperm(n)[keep:]] = False
            out["image"][~valid] = 0.0
            out["weight"][~valid] = 0.0
            out["in_frame"][~valid] = False
            out["center_valid"][~valid] = False
        out["view"] = torch.tensor([s["view"] for s in views])
        out["view_valid"] = valid
        out["group_ok"] = bool(views[0]["calibrated"])
        out["flipped"] = bool(views[0]["flipped"])
        out["project"] = views[0]["project"]
        out["frame"] = views[0]["frame"]
        return out

__getitem__

__getitem__(g: int) -> dict

Group g as stacked per-view tensors, one flip decision for all views.

Source code in src/deeperfly/pose2d/training/dataset.py

def __getitem__(self, g: int) -> dict:
    """Group `g` as stacked per-view tensors, one flip decision for all views."""
    flip = None
    turns: list = [None] * len(self.groups[g])
    if self.base.augment is not None:
        flip = bool(torch.rand(()) < self.base.augment.flip_p)
        symmetries = (
            self.base.symmetries(self.base.items[self.groups[g][0]][0])
            if self.base.augment.rig_symmetry
            else []
        )
        if symmetries:
            flip, signs = symmetries[int(torch.randint(len(symmetries), ()))]
            turns = [signs[self.base.items[i][2]] for i in self.groups[g]]
    views = [
        self.base.sample(i, flip=flip, turn=turn)
        for i, turn in zip(self.groups[g], turns, strict=True)
    ]
    out = {
        key: torch.stack([s[key] for s in views])
        for key in (
            "image",
            "kp",
            "weight",
            "in_frame",
            "affine",
            "gt_video_px",
            "proj",
            "centers",
            "center_valid",
        )
    }
    n = len(views)
    valid = torch.ones(n, dtype=torch.bool)
    if (
        self.base.augment is not None
        and self.view_drop_p > 0
        and n > self.view_min
        and bool(torch.rand(()) < self.view_drop_p)
    ):
        keep = int(torch.randint(self.view_min, n + 1, ()))
        valid[torch.randperm(n)[keep:]] = False
        out["image"][~valid] = 0.0
        out["weight"][~valid] = 0.0
        out["in_frame"][~valid] = False
        out["center_valid"][~valid] = False
    out["view"] = torch.tensor([s["view"] for s in views])
    out["view_valid"] = valid
    out["group_ok"] = bool(views[0]["calibrated"])
    out["flipped"] = bool(views[0]["flipped"])
    out["project"] = views[0]["project"]
    out["frame"] = views[0]["frame"]
    return out

__init__

__init__(
    base: SampleDataset,
    *,
    view_drop_p: float = 0.0,
    view_min: int = 1,
) -> None

Group the items of base by (project, frame).

Source code in src/deeperfly/pose2d/training/dataset.py

def __init__(
    self, base: SampleDataset, *, view_drop_p: float = 0.0, view_min: int = 1
) -> None:
    """Group the items of `base` by `(project, frame)`."""
    self.base = base
    self.groups = base.groups()
    self.view_drop_p = float(view_drop_p)
    self.view_min = int(view_min)

__len__

__len__() -> int

The number of (project, frame) groups.

Source code in src/deeperfly/pose2d/training/dataset.py

def __len__(self) -> int:
    """The number of `(project, frame)` groups."""
    return len(self.groups)

group_projects

group_projects() -> list[str]

The project name of every group.

Source code in src/deeperfly/pose2d/training/dataset.py

def group_projects(self) -> list[str]:
    """The project name of every group."""
    return [self.base.project_name(self.base.items[g[0]][0]) for g in self.groups]

n_views

n_views() -> list[int]

The view count of every group, for bucketing.

Source code in src/deeperfly/pose2d/training/dataset.py

def n_views(self) -> list[int]:
    """The view count of every group, for bucketing."""
    return [len(g) for g in self.groups]

LabelDataset

Bases: SampleDataset

One item per (project, frame, view) of the projects' labeled frames.

Parameters:

Name	Type	Description	Default
projects	Mapping[str, Project]	name -> Project whose labels to read, in order.	required
skeleton	Skeleton | None	The skeleton every project labels; the first project's by default.	None
animals	int | None	Instance slots per frame; the most any project holds by default.	None
input_hw	tuple[int, int]	Network input (height, width).	required
margin_px	int | tuple[int, int]	The detector's off-frame margin in input pixels, (y, x) or one value; a label beyond it weighs 0 in that view.	0
augment	Augment | None	Training-time augmentation, or None for evaluation.	None
crops	dict[str, dict[str, tuple[int, int, int, int]]] | None	Per-project, per-view video windows; the whole frame when absent (see label_windows).	None
keypoint_weights		Optional (K,) weights multiplied into every cell's weight.	None
cache_dir	str | Path | None	Where frame caches live (<project>/cache by default).	None
supervised_only	bool	Drop (frame, view) items with no weighted label, as a monocular run wants; grouped training keeps them as context.	False
fill_weight	float	Weight of the fills (ADR-0070); 0 trains and scores GT only, as an evaluation set must.	0.0
off_view_weight	float	Weight of a target outside its own view's image.	1.0
fill_min_views	int	Labeled views a keypoint needs before it is filled.	2
fill_max_residual_px	float	Skip the fill of a keypoint whose 3D reprojects further than this from one of its own labels.	10.0
crop_tightness	tuple[float, float] | None	(lo, hi) log-uniform crop size relative to the animal's box; when set, frames are cached at native resolution and crops is ignored.	None
crop_jitter	float	Training-time center jitter, as a fraction of the crop.	0.1
crop_eval_tightness	float	The crop size used without augmentation.	1.25
view_aspect	dict[str, float] | None	View name -> the aspect its crop box is drawn at.	None
crop_min_points	int	Below this many finite points a view falls back to the largest window of the frame at the input's aspect.	4
crop_full_p	float	Probability of a letterboxed whole-frame crop instead.	0.1
mark_others	bool	Mark the other animals' centers in each sample (mark_others).	False
mark_jitter	float	Training-time error of a marked center, in input pixels.	6.0
frames	Mapping[str, Sequence[int]] | None	project -> frames to restrict a project to; this is how a split holds out frames of a project whose others are trained on.	None
exclude_frames	Mapping[str, Sequence[int]] | None	project -> frames to drop. A frame without labels in either table is a ConfigError: in a split file it is a typo.	None

Source code in src/deeperfly/pose2d/training/dataset.py

class LabelDataset(SampleDataset):
    """One item per `(project, frame, view)` of the projects' labeled frames.

    Args:
        projects: `name -> Project` whose labels to read, in order.
        skeleton: The skeleton every project labels; the first project's by default.
        animals: Instance slots per frame; the most any project holds by default.
        input_hw: Network input `(height, width)`.
        margin_px: The detector's off-frame margin in input pixels, `(y, x)` or one
            value; a label beyond it weighs 0 in that view.
        augment: Training-time augmentation, or `None` for evaluation.
        crops: Per-project, per-view video windows; the whole frame when absent (see
            [label_windows][deeperfly.pose2d.training.dataset.label_windows]).
        keypoint_weights: Optional `(K,)` weights multiplied into every cell's weight.
        cache_dir: Where frame caches live (`<project>/cache` by default).
        supervised_only: Drop `(frame, view)` items with no weighted label, as a
            monocular run wants; grouped training keeps them as context.
        fill_weight: Weight of the fills (ADR-0070); 0 trains and scores GT only, as an
            evaluation set must.
        off_view_weight: Weight of a target outside its own view's image.
        fill_min_views: Labeled views a keypoint needs before it is filled.
        fill_max_residual_px: Skip the fill of a keypoint whose 3D reprojects further
            than this from one of its own labels.
        crop_tightness: `(lo, hi)` log-uniform crop size relative to the animal's box;
            when set, frames are cached at native resolution and `crops` is ignored.
        crop_jitter: Training-time center jitter, as a fraction of the crop.
        crop_eval_tightness: The crop size used without augmentation.
        view_aspect: View name -> the aspect its crop box is drawn at.
        crop_min_points: Below this many finite points a view falls back to the
            largest window of the frame at the input's aspect.
        crop_full_p: Probability of a letterboxed whole-frame crop instead.
        mark_others: Mark the other animals' centers in each sample (`mark_others`).
        mark_jitter: Training-time error of a marked center, in input pixels.
        frames: `project -> frames` to restrict a project to; this is how a split
            holds out frames of a project whose others are trained on.
        exclude_frames: `project -> frames` to drop. A frame without labels in either
            table is a `ConfigError`: in a split file it is a typo.
    """

    def __init__(
        self,
        projects: Mapping[str, Project],
        *,
        skeleton: Skeleton | None = None,
        animals: int | None = None,
        input_hw: tuple[int, int],
        margin_px: int | tuple[int, int] = 0,
        augment: Augment | None = None,
        crops: dict[str, dict[str, tuple[int, int, int, int]]] | None = None,
        keypoint_weights=None,
        cache_dir: str | Path | None = None,
        supervised_only: bool = False,
        fill_weight: float = 0.0,
        fill_max_residual_px: float = 10.0,
        fill_min_views: int = 2,
        off_view_weight: float = 1.0,
        crop_tightness: tuple[float, float] | None = None,
        crop_jitter: float = 0.1,
        crop_eval_tightness: float = 1.25,
        crop_min_points: int = 4,
        crop_full_p: float = 0.1,
        view_aspect: dict[str, float] | None = None,
        mark_others: bool = False,
        mark_jitter: float = 6.0,
        frames: Mapping[str, Sequence[int]] | None = None,
        exclude_frames: Mapping[str, Sequence[int]] | None = None,
    ) -> None:
        """Index the projects' labeled frames, fill their caches and solve the fills."""
        projects = dict(projects)
        skeleton = _skeleton_of(projects, skeleton)
        if animals is None:
            animals = max((p.animals for p in projects.values()), default=1)
        super().__init__(
            skeleton,
            input_hw=input_hw,
            margin_px=margin_px,
            augment=augment,
            keypoint_weights=keypoint_weights,
            crop_tightness=crop_tightness,
            crop_jitter=crop_jitter,
            crop_eval_tightness=crop_eval_tightness,
            crop_min_points=crop_min_points,
            crop_full_p=crop_full_p,
            view_aspect=view_aspect,
            animals=animals,
            mark_others=mark_others,
            mark_jitter=mark_jitter,
        )
        self.projects = projects
        for name, project in projects.items():
            labels = project.label_files().load()
            selected = _select_frames(
                name, labels.labeled_frames(), frames, exclude_frames
            )
            if not selected:
                continue
            view_names = tuple(project.views)
            frame_hw = project.frame_sizes()
            cache = FrameCache(
                project,
                name,
                input_hw=None if crop_tightness is not None else self.input_hw,
                crops=(crops or {}).get(name),
                cache_dir=cache_dir,
                frame_hw=frame_hw,
            )
            cache.ensure(selected)
            slots = sorted(
                {0, *(i for t in selected for i in labels.labeled_instances(t))}
            )
            rig, targets = _project_fill(
                project,
                labels,
                selected,
                slots,
                keypoint_names=self.keypoint_names,
                frame_hw=frame_hw,
                fill_weight=fill_weight,
                max_residual_px=fill_max_residual_px,
                min_views=fill_min_views,
                off_view_weight=off_view_weight,
            )
            proj = None if rig is None else _pinhole_projections(rig, view_names)
            r = len(self.entries)
            self.entries.append(
                _ProjectEntry(name, view_names, cache, labels, proj, targets)
            )
            for t in selected:
                # A frame's animals are its labeled slots; slot 0 is always indexed,
                # so a single-animal project is the item list it has always been.
                for slot in sorted({0, *labels.labeled_instances(t)}):
                    for v in range(len(view_names)):
                        if supervised_only and not targets[(t, slot)][1][v].any():
                            continue
                        self.items.append((r, t, v, slot))

    def _image(self, r: int, t: int, v: int):
        cache = self.entries[r].cache
        return cache.get(v, t), cache.image_affine(v)

    def _target(self, r: int, t: int, v: int, i: int):
        xy, w, ref, _ = self.entries[r].targets[(t, i)]
        return xy[v], w[v], ref[v]

    def _center(self, r: int, t: int, v: int):
        # Every animal of the frame, whichever one this item is cropped around: the
        # center channel locates them all in the image it is given.
        out = np.full((self.animals, 2), np.nan)
        for slot in range(self.animals):
            entry = self.entries[r].targets.get((t, slot))
            if entry is not None:
                out[slot] = entry[3][v]
        return out

    def _base_affine(self, r: int, v: int) -> Affine2D:
        cache = self.entries[r].cache
        if cache.crops is None:
            return cache.crop_affine(v)
        return input_window(
            cache.crops[cache.view_names[v]], self.input_hw, self._mirrored(r, v)
        )

    def _mirrored(self, r: int, v: int) -> bool:
        return self.entries[r].cache.mirrored(v)

__init__

__init__(
    projects: Mapping[str, Project],
    *,
    skeleton: Skeleton | None = None,
    animals: int | None = None,
    input_hw: tuple[int, int],
    margin_px: int | tuple[int, int] = 0,
    augment: Augment | None = None,
    crops: dict[str, dict[str, tuple[int, int, int, int]]]
    | None = None,
    keypoint_weights=None,
    cache_dir: str | Path | None = None,
    supervised_only: bool = False,
    fill_weight: float = 0.0,
    fill_max_residual_px: float = 10.0,
    fill_min_views: int = 2,
    off_view_weight: float = 1.0,
    crop_tightness: tuple[float, float] | None = None,
    crop_jitter: float = 0.1,
    crop_eval_tightness: float = 1.25,
    crop_min_points: int = 4,
    crop_full_p: float = 0.1,
    view_aspect: dict[str, float] | None = None,
    mark_others: bool = False,
    mark_jitter: float = 6.0,
    frames: Mapping[str, Sequence[int]] | None = None,
    exclude_frames: Mapping[str, Sequence[int]]
    | None = None,
) -> None

Index the projects' labeled frames, fill their caches and solve the fills.

Source code in src/deeperfly/pose2d/training/dataset.py

def __init__(
    self,
    projects: Mapping[str, Project],
    *,
    skeleton: Skeleton | None = None,
    animals: int | None = None,
    input_hw: tuple[int, int],
    margin_px: int | tuple[int, int] = 0,
    augment: Augment | None = None,
    crops: dict[str, dict[str, tuple[int, int, int, int]]] | None = None,
    keypoint_weights=None,
    cache_dir: str | Path | None = None,
    supervised_only: bool = False,
    fill_weight: float = 0.0,
    fill_max_residual_px: float = 10.0,
    fill_min_views: int = 2,
    off_view_weight: float = 1.0,
    crop_tightness: tuple[float, float] | None = None,
    crop_jitter: float = 0.1,
    crop_eval_tightness: float = 1.25,
    crop_min_points: int = 4,
    crop_full_p: float = 0.1,
    view_aspect: dict[str, float] | None = None,
    mark_others: bool = False,
    mark_jitter: float = 6.0,
    frames: Mapping[str, Sequence[int]] | None = None,
    exclude_frames: Mapping[str, Sequence[int]] | None = None,
) -> None:
    """Index the projects' labeled frames, fill their caches and solve the fills."""
    projects = dict(projects)
    skeleton = _skeleton_of(projects, skeleton)
    if animals is None:
        animals = max((p.animals for p in projects.values()), default=1)
    super().__init__(
        skeleton,
        input_hw=input_hw,
        margin_px=margin_px,
        augment=augment,
        keypoint_weights=keypoint_weights,
        crop_tightness=crop_tightness,
        crop_jitter=crop_jitter,
        crop_eval_tightness=crop_eval_tightness,
        crop_min_points=crop_min_points,
        crop_full_p=crop_full_p,
        view_aspect=view_aspect,
        animals=animals,
        mark_others=mark_others,
        mark_jitter=mark_jitter,
    )
    self.projects = projects
    for name, project in projects.items():
        labels = project.label_files().load()
        selected = _select_frames(
            name, labels.labeled_frames(), frames, exclude_frames
        )
        if not selected:
            continue
        view_names = tuple(project.views)
        frame_hw = project.frame_sizes()
        cache = FrameCache(
            project,
            name,
            input_hw=None if crop_tightness is not None else self.input_hw,
            crops=(crops or {}).get(name),
            cache_dir=cache_dir,
            frame_hw=frame_hw,
        )
        cache.ensure(selected)
        slots = sorted(
            {0, *(i for t in selected for i in labels.labeled_instances(t))}
        )
        rig, targets = _project_fill(
            project,
            labels,
            selected,
            slots,
            keypoint_names=self.keypoint_names,
            frame_hw=frame_hw,
            fill_weight=fill_weight,
            max_residual_px=fill_max_residual_px,
            min_views=fill_min_views,
            off_view_weight=off_view_weight,
        )
        proj = None if rig is None else _pinhole_projections(rig, view_names)
        r = len(self.entries)
        self.entries.append(
            _ProjectEntry(name, view_names, cache, labels, proj, targets)
        )
        for t in selected:
            # A frame's animals are its labeled slots; slot 0 is always indexed,
            # so a single-animal project is the item list it has always been.
            for slot in sorted({0, *labels.labeled_instances(t)}):
                for v in range(len(view_names)):
                    if supervised_only and not targets[(t, slot)][1][v].any():
                        continue
                    self.items.append((r, t, v, slot))

LocatorDataset

Bases: SampleDataset

One item per (project, frame, video): where every view sees the animal.

The full-frame locator's training shape (ADR-0064): a sample is a whole source frame letterboxed, never cropped, and channel k carries view k's animal center (the center channel's target, ADR-0070), weighing 0 where view k does not read this video. At one animal a channel is one Gaussian; above it, a channel is the center channel's mixture with a peak per animal of the frame (ADR-0127), carried as the item's centers, (C * A, 2) channel-major.

Parameters:

Name	Type	Description	Default
projects	Mapping[str, Project]	name -> Project whose labels to read, in order.	required
views	Sequence[str]	The channel order, one view name per channel, naming every view of every project.	required
input_hw	tuple[int, int]	Network input (height, width).	required
margin_px	int | tuple[int, int]	The detector's off-frame margin in input pixels.	0
augment	Augment | None	Training-time augmentation, or None for the plain resize.	None
cache_dir	str | Path | None	Where frame caches live (<project>/cache by default).	None
fill_max_residual_px	float	Skip a keypoint whose 3D reprojects further than this from its own labels, before the center is taken.	10.0
fill_min_views	int	Labeled views a keypoint needs before it is triangulated.	2
frames	Mapping[str, Sequence[int]] | None	project -> frames to restrict a project to.	None
exclude_frames	Mapping[str, Sequence[int]] | None	project -> frames to drop.	None
animals	int | None	Instance slots per frame; the most any project holds by default.	None

Source code in src/deeperfly/pose2d/training/dataset.py

class LocatorDataset(SampleDataset):
    """One item per `(project, frame, video)`: where every view sees the animal.

    The full-frame locator's training shape (ADR-0064): a sample is a whole source
    frame letterboxed, never cropped, and channel `k` carries view `k`'s animal center
    (the center channel's target, ADR-0070), weighing 0 where view `k` does not read
    this video. At one animal a channel is one Gaussian; above it, a channel is the
    center channel's mixture with a peak per animal of the frame (ADR-0127), carried
    as the item's `centers`, `(C * A, 2)` channel-major.

    Args:
        projects: `name -> Project` whose labels to read, in order.
        views: The channel order, one view name per channel, naming every view of every
            project.
        input_hw: Network input `(height, width)`.
        margin_px: The detector's off-frame margin in input pixels.
        augment: Training-time augmentation, or `None` for the plain resize.
        cache_dir: Where frame caches live (`<project>/cache` by default).
        fill_max_residual_px: Skip a keypoint whose 3D reprojects further than this
            from its own labels, before the center is taken.
        fill_min_views: Labeled views a keypoint needs before it is triangulated.
        frames: `project -> frames` to restrict a project to.
        exclude_frames: `project -> frames` to drop.
        animals: Instance slots per frame; the most any project holds by default.
    """

    def __init__(
        self,
        projects: Mapping[str, Project],
        *,
        views: Sequence[str],
        input_hw: tuple[int, int],
        margin_px: int | tuple[int, int] = 0,
        augment: Augment | None = None,
        cache_dir: str | Path | None = None,
        fill_max_residual_px: float = 10.0,
        fill_min_views: int = 2,
        frames: Mapping[str, Sequence[int]] | None = None,
        exclude_frames: Mapping[str, Sequence[int]] | None = None,
        animals: int | None = None,
    ) -> None:
        """Index the labeled frames of the projects, one item per video."""
        if animals is None:
            animals = max((p.animals for p in projects.values()), default=1)
        #: Animals a channel locates; above one, `centers` carries them all.
        self.located = int(animals)
        super().__init__(
            locator_skeleton(views),
            input_hw=input_hw,
            margin_px=margin_px,
            augment=augment,
            # A locator sample is the whole frame; nothing draws a box around it.
            crop_tightness=None,
            # One center per (channel, animal) above one animal, as `_center` lays
            # them out; at one the channels' own targets carry it.
            animals=1 if self.located == 1 else len(views) * self.located,
        )
        self.projects = dict(projects)
        channel = {n: k for k, n in enumerate(self.keypoint_names)}
        for name, project in self.projects.items():
            labels = project.label_files().load()
            selected = _select_frames(
                name, labels.labeled_frames(), frames, exclude_frames
            )
            if not selected:
                continue
            view_names = tuple(project.views)
            unknown = [c for c in view_names if c not in channel]
            if unknown:
                raise ConfigError(
                    f"{name}: views {unknown} have no locator channel; the channels "
                    f"are {list(self.keypoint_names)}"
                )
            frame_hw = project.frame_sizes()
            cache = FrameCache(
                project,
                name,
                input_hw=None,
                crops=None,
                cache_dir=cache_dir,
                frame_hw=frame_hw,
            )
            cache.ensure(selected)
            layout = project.layout
            sources: list[str] = []
            source_of = []
            for view_name in view_names:
                src = layout.video_of(view_name)
                if src not in sources:
                    sources.append(src)
                source_of.append(sources.index(src))
            rep = tuple(source_of.index(s) for s in range(len(sources)))
            slots = sorted(
                {0, *(i for t in selected for i in labels.labeled_instances(t))}
            )
            slots = [i for i in slots if i < self.located]
            _, filled = _project_fill(
                project,
                labels,
                selected,
                slots,
                keypoint_names=tuple(project.skeleton().point_names),
                frame_hw=frame_hw,
                fill_weight=0.0,
                max_residual_px=fill_max_residual_px,
                min_views=fill_min_views,
            )
            shape = (len(sources), self.located, len(channel))
            targets = {
                int(t): (np.full((*shape, 2), np.nan), np.zeros(shape, np.float32))
                for t in selected
            }
            for (t, slot), (_, _, _, center) in filled.items():
                xy, w = targets[int(t)]
                for v, view_name in enumerate(view_names):
                    if not np.isfinite(center[v]).all():
                        continue
                    k, s = channel[view_name], source_of[v]
                    xy[s, slot, k] = center[v]
                    w[s, slot, k] = 1.0
            r = len(self.entries)
            self.entries.append(
                _LocatorEntry(
                    name,
                    tuple(sources),
                    cache,
                    rep,
                    tuple(cache.sizes[src] for src in sources),
                    targets,
                )
            )
            for t in selected:
                for s in range(len(sources)):
                    # A frame with no solved center locates nothing; it would weigh 0
                    # in every channel and only dilute the batch.
                    if targets[int(t)][1][s].any():
                        self.items.append((r, int(t), s, 0))

    def _image(self, r: int, t: int, v: int):
        entry = self.entries[r]
        # The stored image IS the source frame and the targets are in source pixels,
        # so nothing composes between them.
        return entry.cache.get(entry.rep[v], t), None

    def _target(self, r: int, t: int, v: int, i: int):
        xy, w = self.entries[r].targets[t]
        # No reference points: a locator sample is never cropped around anything.
        return xy[v, 0], w[v, 0], None

    def _center(self, r: int, t: int, v: int):
        if self.located == 1:
            return None
        xy, _ = self.entries[r].targets[t]
        # (A, C, 2) -> (C, A, 2) -> (C * A, 2): channel k's animals at k * A onward.
        return xy[v].transpose(1, 0, 2).reshape(-1, 2)

    def _base_affine(self, r: int, v: int) -> Affine2D:
        hw = self.entries[r].frame_hw[v]
        return Affine2D.crop(padded_window(hw, self.input_hw), self.input_hw)

__init__

__init__(
    projects: Mapping[str, Project],
    *,
    views: Sequence[str],
    input_hw: tuple[int, int],
    margin_px: int | tuple[int, int] = 0,
    augment: Augment | None = None,
    cache_dir: str | Path | None = None,
    fill_max_residual_px: float = 10.0,
    fill_min_views: int = 2,
    frames: Mapping[str, Sequence[int]] | None = None,
    exclude_frames: Mapping[str, Sequence[int]]
    | None = None,
    animals: int | None = None,
) -> None

Index the labeled frames of the projects, one item per video.

Source code in src/deeperfly/pose2d/training/dataset.py

def __init__(
    self,
    projects: Mapping[str, Project],
    *,
    views: Sequence[str],
    input_hw: tuple[int, int],
    margin_px: int | tuple[int, int] = 0,
    augment: Augment | None = None,
    cache_dir: str | Path | None = None,
    fill_max_residual_px: float = 10.0,
    fill_min_views: int = 2,
    frames: Mapping[str, Sequence[int]] | None = None,
    exclude_frames: Mapping[str, Sequence[int]] | None = None,
    animals: int | None = None,
) -> None:
    """Index the labeled frames of the projects, one item per video."""
    if animals is None:
        animals = max((p.animals for p in projects.values()), default=1)
    #: Animals a channel locates; above one, `centers` carries them all.
    self.located = int(animals)
    super().__init__(
        locator_skeleton(views),
        input_hw=input_hw,
        margin_px=margin_px,
        augment=augment,
        # A locator sample is the whole frame; nothing draws a box around it.
        crop_tightness=None,
        # One center per (channel, animal) above one animal, as `_center` lays
        # them out; at one the channels' own targets carry it.
        animals=1 if self.located == 1 else len(views) * self.located,
    )
    self.projects = dict(projects)
    channel = {n: k for k, n in enumerate(self.keypoint_names)}
    for name, project in self.projects.items():
        labels = project.label_files().load()
        selected = _select_frames(
            name, labels.labeled_frames(), frames, exclude_frames
        )
        if not selected:
            continue
        view_names = tuple(project.views)
        unknown = [c for c in view_names if c not in channel]
        if unknown:
            raise ConfigError(
                f"{name}: views {unknown} have no locator channel; the channels "
                f"are {list(self.keypoint_names)}"
            )
        frame_hw = project.frame_sizes()
        cache = FrameCache(
            project,
            name,
            input_hw=None,
            crops=None,
            cache_dir=cache_dir,
            frame_hw=frame_hw,
        )
        cache.ensure(selected)
        layout = project.layout
        sources: list[str] = []
        source_of = []
        for view_name in view_names:
            src = layout.video_of(view_name)
            if src not in sources:
                sources.append(src)
            source_of.append(sources.index(src))
        rep = tuple(source_of.index(s) for s in range(len(sources)))
        slots = sorted(
            {0, *(i for t in selected for i in labels.labeled_instances(t))}
        )
        slots = [i for i in slots if i < self.located]
        _, filled = _project_fill(
            project,
            labels,
            selected,
            slots,
            keypoint_names=tuple(project.skeleton().point_names),
            frame_hw=frame_hw,
            fill_weight=0.0,
            max_residual_px=fill_max_residual_px,
            min_views=fill_min_views,
        )
        shape = (len(sources), self.located, len(channel))
        targets = {
            int(t): (np.full((*shape, 2), np.nan), np.zeros(shape, np.float32))
            for t in selected
        }
        for (t, slot), (_, _, _, center) in filled.items():
            xy, w = targets[int(t)]
            for v, view_name in enumerate(view_names):
                if not np.isfinite(center[v]).all():
                    continue
                k, s = channel[view_name], source_of[v]
                xy[s, slot, k] = center[v]
                w[s, slot, k] = 1.0
        r = len(self.entries)
        self.entries.append(
            _LocatorEntry(
                name,
                tuple(sources),
                cache,
                rep,
                tuple(cache.sizes[src] for src in sources),
                targets,
            )
        )
        for t in selected:
            for s in range(len(sources)):
                # A frame with no solved center locates nothing; it would weigh 0
                # in every channel and only dilute the batch.
                if targets[int(t)][1][s].any():
                    self.items.append((r, int(t), s, 0))

SampleDataset

Bases: Dataset

The per-view sample logic every label source shares.

A source subclasses it, fills self.entries (one per project: objects with name, view_names and a (V, 3, 4) video-pixel proj or None) and self.items ((project, frame, view, instance) tuples), and implements _image, _target and _base_affine. sample turns an item into the training dict.

Source code in src/deeperfly/pose2d/training/dataset.py

class SampleDataset(Dataset):
    """The per-view sample logic every label source shares.

    A source subclasses it, fills `self.entries` (one per project: objects with
    `name`, `view_names` and a `(V, 3, 4)` video-pixel `proj` or `None`) and
    `self.items` (`(project, frame, view, instance)` tuples), and implements
    `_image`, `_target` and `_base_affine`. `sample` turns an item into the
    training dict.
    """

    entries: list
    items: list[tuple[int, int, int, int]]

    def __init__(
        self,
        skeleton,
        *,
        input_hw: tuple[int, int],
        margin_px: int | tuple[int, int],
        augment: Augment | None,
        keypoint_weights=None,
        crop_tightness: tuple[float, float] | None = None,
        crop_jitter: float = 0.1,
        crop_eval_tightness: float = 1.25,
        crop_min_points: int = 4,
        crop_full_p: float = 0.1,
        view_aspect: dict[str, float] | None = None,
        animals: int = 1,
        mark_others: bool = False,
        mark_jitter: float = 6.0,
    ) -> None:
        """Set the sample knobs every source shares; the subclass fills the items.

        `crop_*` shape the training crop drawn around an animal (see `crop_box`),
        `view_aspect` gives a view its own crop aspect, `animals` is how many
        instance slots a frame of this source may hold, and `mark_others` marks the
        other animals' centers in each sample (`mark_others`), off by `mark_jitter`
        pixels when training.
        """
        self.mark_others = bool(mark_others)
        self.mark_jitter = float(mark_jitter)
        #: Instance slots a frame of this source may hold; the width of every sample's
        #: center target, whether or not the slots above 0 are ever filled.
        self.animals = int(animals)
        self.crop_tightness = crop_tightness
        self.crop_jitter = float(crop_jitter)
        self.crop_eval_tightness = float(crop_eval_tightness)
        self.crop_min_points = int(crop_min_points)
        self.crop_full_p = float(crop_full_p)
        self.view_aspect = dict(view_aspect or {})
        self.skeleton = skeleton
        self.keypoint_names = tuple(skeleton.point_names)
        self.input_hw = (int(input_hw[0]), int(input_hw[1]))
        my, mx = (margin_px, margin_px) if np.isscalar(margin_px) else margin_px
        self.margin_px = (int(my), int(mx))
        self.augment = augment
        self._photometric = augment.photometric() if augment is not None else None
        self.flip_perm = torch.as_tensor(skeleton.flip_perm(), dtype=torch.long)
        k = len(self.keypoint_names)
        self.keypoint_weights = (
            torch.ones(k)
            if keypoint_weights is None
            else (torch.as_tensor(np.asarray(keypoint_weights, dtype=np.float32)))
        )
        if tuple(self.keypoint_weights.shape) != (k,):
            raise ValueError(
                f"keypoint_weights must be ({k},), "
                f"got {tuple(self.keypoint_weights.shape)}"
            )
        self.entries = []
        self.items = []

    # -- what a source provides ----------------------------------------------------

    def _image(self, r: int, t: int, v: int) -> tuple[np.ndarray, Affine2D | None]:
        """The `(H, W)` uint8 image of one item, and the video -> image-pixel affine.

        `None` is the identity: the image is the view's own frame.
        """
        raise NotImplementedError

    def _target(
        self, r: int, t: int, v: int, i: int
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray | None]:
        """`(targets, weights, reference)` of one item, for animal `i`.

        `(K, 2)` video-pixel targets (NaN where none), `(K,)` weights, and the `(K, 2)`
        points a crop is drawn around, `None` for a source whose images are already cut.
        """
        raise NotImplementedError

    def _base_affine(self, r: int, v: int) -> Affine2D:
        """Video pixels -> input pixels for one view when no crop is drawn."""
        raise NotImplementedError

    def _mirrored(self, r: int, v: int) -> bool:
        """Whether the apparatus reflects this view, so its crop mirrors (ADR-0053)."""
        return False

    def _center(self, r: int, t: int, v: int) -> np.ndarray | None:
        """`(A, 2)` video-pixel animal centers in this view; `None` without a rig."""
        return None

    # -- the per-sample crop ------------------------------------------------------

    def crop_box(
        self,
        ref: np.ndarray,
        frame_hw: tuple[int, int],
        *,
        train: bool,
        aspect: float | None = None,
    ) -> tuple[float, float, float, float]:
        """A video-pixel window `(x0, y0, x1, y1)` around the animal (ADR-0046).

        The finite `ref` points' bounding box grown to the input's aspect (or `aspect`)
        and scaled by a tightness, log-uniform in `crop_tightness` for training (below 1
        a deliberate bad crop) and `crop_eval_tightness`, centered, for evaluation;
        training centers jitter by `crop_jitter`. Never narrower than half the input,
        never clamped to the frame. With probability `crop_full_p`, or fewer than
        `crop_min_points` points, the sample is the whole frame letterboxed instead.
        """
        h_in, w_in = self.input_hw
        aspect = w_in / h_in if aspect is None else float(aspect)
        pts = ref[np.isfinite(ref).all(axis=-1)]
        if len(pts) < self.crop_min_points:
            return padded_window(frame_hw, self.input_hw)
        if train and self.crop_full_p > 0 and _rand(0.0, 1.0) < self.crop_full_p:
            return padded_window(frame_hw, self.input_hw)
        lo, hi = pts.min(axis=0), pts.max(axis=0)
        cx, cy = (lo + hi) / 2
        w, h = hi - lo
        bw = max(float(w), float(h) * aspect, 1.0)
        if train and self.crop_tightness is not None:
            t_lo, t_hi = self.crop_tightness
            t = math.exp(_rand(math.log(t_lo), math.log(t_hi)))
            j = self.crop_jitter
            cx += _rand(-j, j) * bw
            cy += _rand(-j, j) * bw / aspect
        else:
            t = self.crop_eval_tightness
        bw = max(bw * t, w_in / 2)
        bh = bw / aspect
        cx, cy = float(cx), float(cy)
        return (cx - bw / 2, cy - bh / 2, cx + bw / 2, cy + bh / 2)

    # -- lookups ----------------------------------------------------------------

    def view_name(self, project: int, view: int) -> str:
        """The name of view `view` of the project at index `project`."""
        return self.entries[project].view_names[view]

    def project_name(self, project: int) -> str:
        """The name of the project at index `project`."""
        return self.entries[project].name

    def groups(self) -> list[list[int]]:
        """Item indices by `(project, frame, animal)` group, each in view order."""
        by_frame: dict[tuple[int, int, int], list[int]] = {}
        for index, (r, t, _, i) in enumerate(self.items):
            by_frame.setdefault((r, t, i), []).append(index)
        return list(by_frame.values())

    # -- items ------------------------------------------------------------------

    def __len__(self) -> int:
        """The number of `(project, frame, view, animal)` items."""
        return len(self.items)

    def __getitem__(self, i: int) -> dict:
        """Item `i` with a fresh augmentation draw; see `sample`."""
        return self.sample(i)

    def symmetries(self, r: int) -> list:
        """Project `r`'s `rig_symmetries`, found once; empty without a calibration."""
        cache = self.__dict__.setdefault("_symmetries", {})
        if r not in cache:
            proj = self.entries[r].proj
            cache[r] = [] if proj is None else rig_symmetries(proj)
        return cache[r]

    def sample(
        self,
        i: int,
        *,
        flip: bool | None = None,
        turn: tuple[int, int] | None = None,
    ) -> dict:
        """Item `i`, optionally forcing the mirror decision (for a consistent group).

        `turn` is the `(x, y)` axis flip of this view under a rig symmetry the group
        drew (`rig_symmetries`), with `flip` saying whether it is a mirror.

        A dict of `image (1, H, W)` in `[0, 1]`, `kp (K, 2)` input-normalized (NaN where
        unlabeled), `weight (K,)`, `in_frame (K,)`, `affine (2, 3)` video -> input
        pixels, `gt_video_px (K, 2)` in `kp`'s channel order (mirrored when `flipped`),
        `proj (3, 4)` with `calibrated`, `flipped`, the item's indices, and the
        input-normalized `centers (A, 2)` of every animal this view sees, with
        `center_valid (A,)`.
        """
        r, t, v, n = self.items[i]
        entry = self.entries[r]
        h, w = self.input_hw
        raw, img_affine = self._image(r, t, v)
        img = torch.from_numpy(raw).float().div_(255.0)[None]
        xy, wt, ref = self._target(r, t, v, n)
        gt_video_px = torch.as_tensor(xy, dtype=torch.float32)
        weight = torch.as_tensor(wt, dtype=torch.float32) * self.keypoint_weights
        train = self.augment is not None
        if self.crop_tightness is not None and ref is not None:
            base = input_window(
                self.crop_box(
                    ref,
                    raw.shape,
                    train=train,
                    aspect=self.view_aspect.get(self.view_name(r, v)),
                ),
                self.input_hw,
                self._mirrored(r, v),
            )
        else:
            base = self._base_affine(r, v)
        flipped = False
        affine = base
        if train:
            aug, flipped = sample_affine(
                self.input_hw,
                degrees=self.augment.degrees,
                scale=self.augment.scale,
                aspect=self.augment.aspect,
                translate=self.augment.translate,
                flip_p=(
                    0.0
                    if turn is not None
                    else self.augment.flip_p
                    if flip is None
                    else float(flip)
                ),
            )
            affine = aug @ base
            if turn is not None:
                # One symmetry of the rig for the whole group: this view's axis
                # flips, and the channels swap when it is a mirror.
                affine = aug @ _axis_flip(turn, self.input_hw) @ base
                flipped = bool(flip)
        warp = affine if img_affine is None else affine @ img_affine.inverse()
        img = _warp_antialiased(img, warp, self.input_hw)
        kp = affine.apply(gt_video_px)
        if flipped:
            kp, weight = flip_channels(kp, weight, self.flip_perm)
            gt_video_px = gt_video_px.index_select(0, self.flip_perm)
        if train:
            img = self._photometric(img)
        in_frame = _inside(kp, self.input_hw, (0, 0))
        weight = weight * _inside(kp, self.input_hw, self.margin_px).to(weight.dtype)
        cen = self._center(r, t, v)
        cen = np.full((self.animals, 2), np.nan) if cen is None else np.asarray(cen)
        centers = affine.apply(torch.as_tensor(cen, dtype=torch.float32))
        if self.mark_others and self.animals > 1:
            jitter = self.mark_jitter if train else 0.0
            img = mark_others(img, centers, n, jitter=jitter)
        proj = entry.proj[v] if entry.proj is not None else np.zeros((3, 4))
        return {
            "image": img,
            "kp": kp / torch.tensor([w, h], dtype=kp.dtype),
            "weight": weight,
            "in_frame": in_frame & (weight > 0),
            "affine": affine.matrix.to(torch.float32),
            "gt_video_px": gt_video_px,
            "centers": centers / torch.tensor([w, h], dtype=centers.dtype),
            "center_valid": _inside(centers, self.input_hw, self.margin_px),
            "proj": torch.as_tensor(proj, dtype=torch.float32),
            "calibrated": entry.proj is not None,
            "flipped": flipped,
            "project": r,
            "frame": t,
            "view": v,
            "instance": n,
        }

__getitem__

__getitem__(i: int) -> dict

Item i with a fresh augmentation draw; see sample.

Source code in src/deeperfly/pose2d/training/dataset.py

def __getitem__(self, i: int) -> dict:
    """Item `i` with a fresh augmentation draw; see `sample`."""
    return self.sample(i)

__init__

__init__(
    skeleton,
    *,
    input_hw: tuple[int, int],
    margin_px: int | tuple[int, int],
    augment: Augment | None,
    keypoint_weights=None,
    crop_tightness: tuple[float, float] | None = None,
    crop_jitter: float = 0.1,
    crop_eval_tightness: float = 1.25,
    crop_min_points: int = 4,
    crop_full_p: float = 0.1,
    view_aspect: dict[str, float] | None = None,
    animals: int = 1,
    mark_others: bool = False,
    mark_jitter: float = 6.0,
) -> None

Set the sample knobs every source shares; the subclass fills the items.

crop_* shape the training crop drawn around an animal (see crop_box), view_aspect gives a view its own crop aspect, animals is how many instance slots a frame of this source may hold, and mark_others marks the other animals' centers in each sample (mark_others), off by mark_jitter pixels when training.

Source code in src/deeperfly/pose2d/training/dataset.py

def __init__(
    self,
    skeleton,
    *,
    input_hw: tuple[int, int],
    margin_px: int | tuple[int, int],
    augment: Augment | None,
    keypoint_weights=None,
    crop_tightness: tuple[float, float] | None = None,
    crop_jitter: float = 0.1,
    crop_eval_tightness: float = 1.25,
    crop_min_points: int = 4,
    crop_full_p: float = 0.1,
    view_aspect: dict[str, float] | None = None,
    animals: int = 1,
    mark_others: bool = False,
    mark_jitter: float = 6.0,
) -> None:
    """Set the sample knobs every source shares; the subclass fills the items.

    `crop_*` shape the training crop drawn around an animal (see `crop_box`),
    `view_aspect` gives a view its own crop aspect, `animals` is how many
    instance slots a frame of this source may hold, and `mark_others` marks the
    other animals' centers in each sample (`mark_others`), off by `mark_jitter`
    pixels when training.
    """
    self.mark_others = bool(mark_others)
    self.mark_jitter = float(mark_jitter)
    #: Instance slots a frame of this source may hold; the width of every sample's
    #: center target, whether or not the slots above 0 are ever filled.
    self.animals = int(animals)
    self.crop_tightness = crop_tightness
    self.crop_jitter = float(crop_jitter)
    self.crop_eval_tightness = float(crop_eval_tightness)
    self.crop_min_points = int(crop_min_points)
    self.crop_full_p = float(crop_full_p)
    self.view_aspect = dict(view_aspect or {})
    self.skeleton = skeleton
    self.keypoint_names = tuple(skeleton.point_names)
    self.input_hw = (int(input_hw[0]), int(input_hw[1]))
    my, mx = (margin_px, margin_px) if np.isscalar(margin_px) else margin_px
    self.margin_px = (int(my), int(mx))
    self.augment = augment
    self._photometric = augment.photometric() if augment is not None else None
    self.flip_perm = torch.as_tensor(skeleton.flip_perm(), dtype=torch.long)
    k = len(self.keypoint_names)
    self.keypoint_weights = (
        torch.ones(k)
        if keypoint_weights is None
        else (torch.as_tensor(np.asarray(keypoint_weights, dtype=np.float32)))
    )
    if tuple(self.keypoint_weights.shape) != (k,):
        raise ValueError(
            f"keypoint_weights must be ({k},), "
            f"got {tuple(self.keypoint_weights.shape)}"
        )
    self.entries = []
    self.items = []

__len__

__len__() -> int

The number of (project, frame, view, animal) items.

Source code in src/deeperfly/pose2d/training/dataset.py

def __len__(self) -> int:
    """The number of `(project, frame, view, animal)` items."""
    return len(self.items)

crop_box

crop_box(
    ref: ndarray,
    frame_hw: tuple[int, int],
    *,
    train: bool,
    aspect: float | None = None,
) -> tuple[float, float, float, float]

A video-pixel window (x0, y0, x1, y1) around the animal (ADR-0046).

The finite ref points' bounding box grown to the input's aspect (or aspect) and scaled by a tightness, log-uniform in crop_tightness for training (below 1 a deliberate bad crop) and crop_eval_tightness, centered, for evaluation; training centers jitter by crop_jitter. Never narrower than half the input, never clamped to the frame. With probability crop_full_p, or fewer than crop_min_points points, the sample is the whole frame letterboxed instead.

Source code in src/deeperfly/pose2d/training/dataset.py

def crop_box(
    self,
    ref: np.ndarray,
    frame_hw: tuple[int, int],
    *,
    train: bool,
    aspect: float | None = None,
) -> tuple[float, float, float, float]:
    """A video-pixel window `(x0, y0, x1, y1)` around the animal (ADR-0046).

    The finite `ref` points' bounding box grown to the input's aspect (or `aspect`)
    and scaled by a tightness, log-uniform in `crop_tightness` for training (below 1
    a deliberate bad crop) and `crop_eval_tightness`, centered, for evaluation;
    training centers jitter by `crop_jitter`. Never narrower than half the input,
    never clamped to the frame. With probability `crop_full_p`, or fewer than
    `crop_min_points` points, the sample is the whole frame letterboxed instead.
    """
    h_in, w_in = self.input_hw
    aspect = w_in / h_in if aspect is None else float(aspect)
    pts = ref[np.isfinite(ref).all(axis=-1)]
    if len(pts) < self.crop_min_points:
        return padded_window(frame_hw, self.input_hw)
    if train and self.crop_full_p > 0 and _rand(0.0, 1.0) < self.crop_full_p:
        return padded_window(frame_hw, self.input_hw)
    lo, hi = pts.min(axis=0), pts.max(axis=0)
    cx, cy = (lo + hi) / 2
    w, h = hi - lo
    bw = max(float(w), float(h) * aspect, 1.0)
    if train and self.crop_tightness is not None:
        t_lo, t_hi = self.crop_tightness
        t = math.exp(_rand(math.log(t_lo), math.log(t_hi)))
        j = self.crop_jitter
        cx += _rand(-j, j) * bw
        cy += _rand(-j, j) * bw / aspect
    else:
        t = self.crop_eval_tightness
    bw = max(bw * t, w_in / 2)
    bh = bw / aspect
    cx, cy = float(cx), float(cy)
    return (cx - bw / 2, cy - bh / 2, cx + bw / 2, cy + bh / 2)

groups

groups() -> list[list[int]]

Item indices by (project, frame, animal) group, each in view order.

Source code in src/deeperfly/pose2d/training/dataset.py

def groups(self) -> list[list[int]]:
    """Item indices by `(project, frame, animal)` group, each in view order."""
    by_frame: dict[tuple[int, int, int], list[int]] = {}
    for index, (r, t, _, i) in enumerate(self.items):
        by_frame.setdefault((r, t, i), []).append(index)
    return list(by_frame.values())

project_name

project_name(project: int) -> str

The name of the project at index project.

Source code in src/deeperfly/pose2d/training/dataset.py

def project_name(self, project: int) -> str:
    """The name of the project at index `project`."""
    return self.entries[project].name

sample

sample(
    i: int,
    *,
    flip: bool | None = None,
    turn: tuple[int, int] | None = None,
) -> dict

Item i, optionally forcing the mirror decision (for a consistent group).

turn is the (x, y) axis flip of this view under a rig symmetry the group drew (rig_symmetries), with flip saying whether it is a mirror.

A dict of image (1, H, W) in [0, 1], kp (K, 2) input-normalized (NaN where unlabeled), weight (K,), in_frame (K,), affine (2, 3) video -> input pixels, gt_video_px (K, 2) in kp's channel order (mirrored when flipped), proj (3, 4) with calibrated, flipped, the item's indices, and the input-normalized centers (A, 2) of every animal this view sees, with center_valid (A,).

Source code in src/deeperfly/pose2d/training/dataset.py

def sample(
    self,
    i: int,
    *,
    flip: bool | None = None,
    turn: tuple[int, int] | None = None,
) -> dict:
    """Item `i`, optionally forcing the mirror decision (for a consistent group).

    `turn` is the `(x, y)` axis flip of this view under a rig symmetry the group
    drew (`rig_symmetries`), with `flip` saying whether it is a mirror.

    A dict of `image (1, H, W)` in `[0, 1]`, `kp (K, 2)` input-normalized (NaN where
    unlabeled), `weight (K,)`, `in_frame (K,)`, `affine (2, 3)` video -> input
    pixels, `gt_video_px (K, 2)` in `kp`'s channel order (mirrored when `flipped`),
    `proj (3, 4)` with `calibrated`, `flipped`, the item's indices, and the
    input-normalized `centers (A, 2)` of every animal this view sees, with
    `center_valid (A,)`.
    """
    r, t, v, n = self.items[i]
    entry = self.entries[r]
    h, w = self.input_hw
    raw, img_affine = self._image(r, t, v)
    img = torch.from_numpy(raw).float().div_(255.0)[None]
    xy, wt, ref = self._target(r, t, v, n)
    gt_video_px = torch.as_tensor(xy, dtype=torch.float32)
    weight = torch.as_tensor(wt, dtype=torch.float32) * self.keypoint_weights
    train = self.augment is not None
    if self.crop_tightness is not None and ref is not None:
        base = input_window(
            self.crop_box(
                ref,
                raw.shape,
                train=train,
                aspect=self.view_aspect.get(self.view_name(r, v)),
            ),
            self.input_hw,
            self._mirrored(r, v),
        )
    else:
        base = self._base_affine(r, v)
    flipped = False
    affine = base
    if train:
        aug, flipped = sample_affine(
            self.input_hw,
            degrees=self.augment.degrees,
            scale=self.augment.scale,
            aspect=self.augment.aspect,
            translate=self.augment.translate,
            flip_p=(
                0.0
                if turn is not None
                else self.augment.flip_p
                if flip is None
                else float(flip)
            ),
        )
        affine = aug @ base
        if turn is not None:
            # One symmetry of the rig for the whole group: this view's axis
            # flips, and the channels swap when it is a mirror.
            affine = aug @ _axis_flip(turn, self.input_hw) @ base
            flipped = bool(flip)
    warp = affine if img_affine is None else affine @ img_affine.inverse()
    img = _warp_antialiased(img, warp, self.input_hw)
    kp = affine.apply(gt_video_px)
    if flipped:
        kp, weight = flip_channels(kp, weight, self.flip_perm)
        gt_video_px = gt_video_px.index_select(0, self.flip_perm)
    if train:
        img = self._photometric(img)
    in_frame = _inside(kp, self.input_hw, (0, 0))
    weight = weight * _inside(kp, self.input_hw, self.margin_px).to(weight.dtype)
    cen = self._center(r, t, v)
    cen = np.full((self.animals, 2), np.nan) if cen is None else np.asarray(cen)
    centers = affine.apply(torch.as_tensor(cen, dtype=torch.float32))
    if self.mark_others and self.animals > 1:
        jitter = self.mark_jitter if train else 0.0
        img = mark_others(img, centers, n, jitter=jitter)
    proj = entry.proj[v] if entry.proj is not None else np.zeros((3, 4))
    return {
        "image": img,
        "kp": kp / torch.tensor([w, h], dtype=kp.dtype),
        "weight": weight,
        "in_frame": in_frame & (weight > 0),
        "affine": affine.matrix.to(torch.float32),
        "gt_video_px": gt_video_px,
        "centers": centers / torch.tensor([w, h], dtype=centers.dtype),
        "center_valid": _inside(centers, self.input_hw, self.margin_px),
        "proj": torch.as_tensor(proj, dtype=torch.float32),
        "calibrated": entry.proj is not None,
        "flipped": flipped,
        "project": r,
        "frame": t,
        "view": v,
        "instance": n,
    }

symmetries

symmetries(r: int) -> list

Project r's rig_symmetries, found once; empty without a calibration.

Source code in src/deeperfly/pose2d/training/dataset.py

def symmetries(self, r: int) -> list:
    """Project `r`'s `rig_symmetries`, found once; empty without a calibration."""
    cache = self.__dict__.setdefault("_symmetries", {})
    if r not in cache:
        proj = self.entries[r].proj
        cache[r] = [] if proj is None else rig_symmetries(proj)
    return cache[r]

view_name

view_name(project: int, view: int) -> str

The name of view view of the project at index project.

Source code in src/deeperfly/pose2d/training/dataset.py

def view_name(self, project: int, view: int) -> str:
    """The name of view `view` of the project at index `project`."""
    return self.entries[project].view_names[view]

ViewBucketSampler

Bases: Sampler

Batches of group indices that share a view count, so none needs padding.

Buckets are shuffled (with shuffle), cut into batches and the batches shuffled across buckets. With images_per_batch a bucket of v views takes images_per_batch // v groups a batch, so every step sees about as many images. repeats, one per group, is how often an epoch holds each: its whole part, and once more with the chance of its fraction, drawn anew every epoch (ADR-0136).

Source code in src/deeperfly/pose2d/training/dataset.py

class ViewBucketSampler(Sampler):
    """Batches of group indices that share a view count, so none needs padding.

    Buckets are shuffled (with `shuffle`), cut into batches and the batches shuffled
    across buckets. With `images_per_batch` a bucket of `v` views takes
    `images_per_batch // v` groups a batch, so every step sees about as many images.
    `repeats`, one per group, is how often an epoch holds each: its whole part, and
    once more with the chance of its fraction, drawn anew every epoch (ADR-0136).
    """

    def __init__(
        self,
        n_views: Sequence[int],
        batch_size: int,
        *,
        shuffle: bool = True,
        drop_last: bool = False,
        generator: torch.Generator | None = None,
        images_per_batch: int | None = None,
        repeats: Sequence[float] | None = None,
    ) -> None:
        """Bucket `n_views` (one per group) into batches of `batch_size`."""
        self.n_views = list(n_views)
        self.repeats = None if repeats is None else [float(r) for r in repeats]
        if self.repeats is not None and len(self.repeats) != len(self.n_views):
            raise ValueError(
                f"repeats has {len(self.repeats)} entries for {len(self.n_views)} "
                "groups"
            )
        self.batch_size = int(batch_size)
        self.shuffle = shuffle
        self.drop_last = drop_last
        self.generator = generator
        self.images_per_batch = (
            None if images_per_batch is None else int(images_per_batch)
        )

    def _bucket_batch_size(self, v: int) -> int:
        if self.images_per_batch is None:
            return self.batch_size
        return max(1, self.images_per_batch // max(1, v))

    def _batches(self) -> list[list[int]]:
        buckets: dict[int, list[int]] = {}
        for i, v in enumerate(self.n_views):
            copies = 1
            if self.repeats is not None:
                whole = int(self.repeats[i])
                chance = torch.rand((), generator=self.generator).item()
                copies = whole + (chance < self.repeats[i] - whole)
            buckets.setdefault(v, []).extend([i] * copies)
        batches = []
        for v, members in buckets.items():
            if self.shuffle:
                order = torch.randperm(len(members), generator=self.generator).tolist()
                members = [members[j] for j in order]
            bs = self._bucket_batch_size(v)
            for start in range(0, len(members), bs):
                b = members[start : start + bs]
                if len(b) == bs or not self.drop_last:
                    batches.append(b)
        if self.shuffle:
            order = torch.randperm(len(batches), generator=self.generator).tolist()
            batches = [batches[j] for j in order]
        return batches

    def __iter__(self) -> Iterator[list[int]]:
        """Yield one epoch of equal-view-count batches."""
        return iter(self._batches())

    def __len__(self) -> int:
        """The number of batches per epoch."""
        return len(self._batches())

__init__

__init__(
    n_views: Sequence[int],
    batch_size: int,
    *,
    shuffle: bool = True,
    drop_last: bool = False,
    generator: Generator | None = None,
    images_per_batch: int | None = None,
    repeats: Sequence[float] | None = None,
) -> None

Bucket n_views (one per group) into batches of batch_size.

Source code in src/deeperfly/pose2d/training/dataset.py

def __init__(
    self,
    n_views: Sequence[int],
    batch_size: int,
    *,
    shuffle: bool = True,
    drop_last: bool = False,
    generator: torch.Generator | None = None,
    images_per_batch: int | None = None,
    repeats: Sequence[float] | None = None,
) -> None:
    """Bucket `n_views` (one per group) into batches of `batch_size`."""
    self.n_views = list(n_views)
    self.repeats = None if repeats is None else [float(r) for r in repeats]
    if self.repeats is not None and len(self.repeats) != len(self.n_views):
        raise ValueError(
            f"repeats has {len(self.repeats)} entries for {len(self.n_views)} "
            "groups"
        )
    self.batch_size = int(batch_size)
    self.shuffle = shuffle
    self.drop_last = drop_last
    self.generator = generator
    self.images_per_batch = (
        None if images_per_batch is None else int(images_per_batch)
    )

__iter__

__iter__() -> Iterator[list[int]]

Yield one epoch of equal-view-count batches.

Source code in src/deeperfly/pose2d/training/dataset.py

def __iter__(self) -> Iterator[list[int]]:
    """Yield one epoch of equal-view-count batches."""
    return iter(self._batches())

__len__

__len__() -> int

The number of batches per epoch.

Source code in src/deeperfly/pose2d/training/dataset.py

def __len__(self) -> int:
    """The number of batches per epoch."""
    return len(self._batches())

active_rig

active_rig(
    project: Project,
    frame_hw: Mapping[str, tuple[int, int]] | None = None,
) -> Rig | None

The project's active calibration on its manifest's views, as the GUI's rig.

None when no calibration is active; frame_hw spares a probe of the videos.

Source code in src/deeperfly/pose2d/training/dataset.py

def active_rig(
    project: Project, frame_hw: Mapping[str, tuple[int, int]] | None = None
) -> Rig | None:
    """The project's active calibration on its manifest's views, as the GUI's rig.

    `None` when no calibration is active; `frame_hw` spares a probe of the videos.
    """
    calibration = project.calibration()
    if calibration is None:
        return None
    return project.rig(calibration, frame_hw=frame_hw)

collate_flat

collate_flat(items: Sequence[dict]) -> dict

Stack flat items into a (B, ...) batch; scalars become 1D tensors.

Source code in src/deeperfly/pose2d/training/dataset.py

def collate_flat(items: Sequence[dict]) -> dict:
    """Stack flat items into a `(B, ...)` batch; scalars become 1D tensors."""
    out = {}
    for key, first in items[0].items():
        vals = [s[key] for s in items]
        if torch.is_tensor(first):
            out[key] = torch.stack(vals)
        else:
            out[key] = torch.as_tensor(vals)
    return out

collate_groups

collate_groups(items: Sequence[dict]) -> dict

Stack groups into a (B, Vmax, ...) batch, padding short groups.

A padded slot has view_valid 0 and weight 0, so it enters no loss; the model keeps it out of the real views' attention.

Source code in src/deeperfly/pose2d/training/dataset.py

def collate_groups(items: Sequence[dict]) -> dict:
    """Stack groups into a `(B, Vmax, ...)` batch, padding short groups.

    A padded slot has `view_valid` 0 and `weight` 0, so it enters no loss; the model
    keeps it out of the real views' attention.
    """
    vmax = max(int(s["view_valid"].shape[0]) for s in items)
    out = {}
    for key, first in items[0].items():
        vals = [s[key] for s in items]
        if torch.is_tensor(first) and first.ndim >= 1:
            padded = []
            for x in vals:
                pad = vmax - x.shape[0]
                if pad:
                    x = torch.cat([x, torch.zeros((pad, *x.shape[1:]), dtype=x.dtype)])
                padded.append(x)
            out[key] = torch.stack(padded)
        else:
            out[key] = torch.as_tensor(vals)
    return out

fill_targets

fill_targets(
    rig: Rig | None,
    detections: Detections | None,
    labels,
    frames: Sequence[int],
    *,
    view_boxes: ndarray | None = None,
    fill_weight: float = 1.0,
    max_residual_px: float = 10.0,
    min_views: int = 2,
    off_view_weight: float = 1.0,
    instance: int = 0,
) -> dict[
    int,
    tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray],
]

One animal's (V, K, 2) targets, (V, K) weights, reference points and center.

GT where the operator clicked; elsewhere the stabilized two-click solve through rig reprojected, weighing fill_weight (ADR-0070), where the clicked views, or the detections agreeing with them, place the keypoint in depth (ADR-0130). The reference points are every solved position, whatever its weight, which a per-sample crop is drawn around; the center is the (V, 2) projection of the per-axis median of the solved 3D keypoints, NaN without a rig. instance is the animal filled, stabilized by the slot of the project's merged detections (ADR-0083) its clicks sit on in each frame (ADR-0128); the other animals are plain background. view_boxes are the (V, 4) whole-frame boxes an off_view_weight other than 1 needs.

Source code in src/deeperfly/pose2d/training/dataset.py

def fill_targets(
    rig: Rig | None,
    detections: Detections | None,
    labels,
    frames: Sequence[int],
    *,
    view_boxes: np.ndarray | None = None,
    fill_weight: float = 1.0,
    max_residual_px: float = 10.0,
    min_views: int = 2,
    off_view_weight: float = 1.0,
    instance: int = 0,
) -> dict[int, tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]]:
    """One animal's `(V, K, 2)` targets, `(V, K)` weights, reference points and center.

    GT where the operator clicked; elsewhere the stabilized two-click solve through
    `rig` reprojected, weighing `fill_weight` (ADR-0070), where the clicked views, or
    the detections agreeing with them, place the keypoint in depth (ADR-0130). The
    reference points are every solved position, whatever its weight, which a per-sample
    crop is drawn around; the center is the `(V, 2)` projection of the per-axis median
    of the solved 3D keypoints, NaN without a rig. `instance` is the animal filled,
    stabilized by the slot of the project's merged `detections` (ADR-0083) its clicks
    sit on in each frame (ADR-0128); the other animals are plain background.
    `view_boxes` are the `(V, 4)` whole-frame boxes an `off_view_weight` other than 1
    needs.
    """
    from ...pose3d.triangulation import (
        placed_in_depth,
        triangulate,
        triangulate_stabilized,
    )

    if off_view_weight != 1.0 and view_boxes is None:
        raise ValueError("off_view_weight needs the views' frame boxes")
    out = {}
    for t in frames:
        gt = np.asarray(
            labels.frame_effective(t, instance=instance), dtype=float
        )  # (V, K, 2)
        ok = np.isfinite(gt).all(axis=-1)
        xy = gt.astype(np.float32)
        w = ok.astype(np.float32)
        solved = np.full((gt.shape[1], 3), np.nan)
        if rig is not None:
            det = None
            if detections is not None:
                det = instance_frame(detections, labels, t, instance)[0]
            if det is not None:
                det[:, labels.absent_frame(t, instance=instance)] = np.nan
            for k in np.flatnonzero(ok.sum(axis=0) >= min_views):
                g = gt[:, k]
                x0 = triangulate(rig, g)
                if not np.isfinite(x0).all():
                    continue
                x = x0
                stab = None if det is None else det[:, k]
                if stab is not None and np.isfinite(stab).all(axis=-1).any():
                    x = triangulate_stabilized(rig, x0, g, stab)
                p = np.asarray(rig.project(x[None, :]))[:, 0]  # (V, 2)
                resid = np.linalg.norm(p[ok[:, k]] - g[ok[:, k]], axis=-1)
                if not np.isfinite(resid).all() or resid.max() > max_residual_px:
                    continue
                # Clicks on one line's two ends (ventral and dorsal) leave the depth to
                # the detections; with none at an angle agreeing, it is a guess the
                # fill would draw into every other view (ADR-0130).
                agree = ok[:, k].copy()
                if stab is not None:
                    gap = np.linalg.norm(p - stab, axis=-1)
                    agree |= np.isfinite(gap) & (gap <= max_residual_px)
                if not placed_in_depth(rig, x, agree):
                    continue
                solved[k] = x
                fill = ~ok[:, k] & np.isfinite(p).all(axis=-1)
                xy[fill, k] = p[fill]
                w[fill, k] = fill_weight
        if off_view_weight != 1.0:
            w[_outside_view(xy, view_boxes)] *= float(off_view_weight)
        w[labels.hidden_frame(t, instance=instance)] = 0.0
        ref = xy.copy()
        xy[w == 0] = np.nan
        center = np.full((gt.shape[0], 2), np.nan)
        if np.isfinite(solved).all(axis=-1).any():
            middle = np.nanmedian(solved, axis=0)
            center = np.asarray(rig.project(middle[None, :]))[:, 0]
        out[int(t)] = (xy, w, ref, center)
    return out

label_windows

label_windows(
    projects: Mapping[str, Project],
    *,
    aspect: float = 2.0,
    pad: float = 0.15,
    quantiles: tuple[float, float] = (0.5, 99.5),
    min_points: int = 50,
    min_hw: tuple[int, int] = (256, 512),
) -> dict[str, dict[str, tuple[int, int, int, int]]]

Per-project, per-view video windows centered on the labels, at a fixed aspect.

The corpus mixes rigs (side views 512x1984 and axial ones 1008x1600 on the ball-tracking recordings), so a whole-frame resize squashes the fly 3.9x one way and 2x the other. Per view, the labels' box (the quantiles, padded by pad) is grown to aspect, centered and fitted inside the frame, never smaller than min_hw and never stretched; a view with fewer than min_points labels gets the largest centered window.

Source code in src/deeperfly/pose2d/training/dataset.py

def label_windows(
    projects: Mapping[str, Project],
    *,
    aspect: float = 2.0,
    pad: float = 0.15,
    quantiles: tuple[float, float] = (0.5, 99.5),
    min_points: int = 50,
    min_hw: tuple[int, int] = (256, 512),
) -> dict[str, dict[str, tuple[int, int, int, int]]]:
    """Per-project, per-view video windows centered on the labels, at a fixed aspect.

    The corpus mixes rigs (side views 512x1984 and axial ones 1008x1600 on the
    ball-tracking recordings), so a whole-frame resize squashes the fly 3.9x one way and
    2x the other. Per view, the labels' box (the `quantiles`, padded by `pad`) is grown
    to `aspect`, centered and fitted inside the frame, never smaller than `min_hw` and
    never stretched; a view with fewer than `min_points` labels gets the largest
    centered window.
    """
    out: dict[str, dict[str, tuple[int, int, int, int]]] = {}
    for name, project in projects.items():
        labels = project.label_files().load()
        frames = labels.labeled_frames()
        frame_hw = project.frame_sizes()
        windows = {}
        for v, view_name in enumerate(project.views):
            fh, fw = frame_hw[view_name]

            seen = [
                labels.frame_effective(t, instance=i)[v]
                for t in frames
                for i in labels.labeled_instances(t)
            ]
            pts = np.concatenate(seen, axis=0) if seen else np.zeros((0, 2))
            pts = pts[np.isfinite(pts).all(axis=-1)]
            if len(pts) < min_points:
                cx, cy, w, h = fw / 2, fh / 2, float(fw), float(fh)
            else:
                lo = np.percentile(pts, quantiles[0], axis=0)
                hi = np.percentile(pts, quantiles[1], axis=0)
                size = hi - lo
                lo, hi = lo - pad * size, hi + pad * size
                cx, cy = (lo + hi) / 2
                w, h = hi - lo
            if w / h < aspect:
                w = h * aspect
            else:
                h = w / aspect
            w, h = max(w, min_hw[1]), max(h, min_hw[0])
            if w > fw:
                w, h = float(fw), fw / aspect
            if h > fh:
                w, h = fh * aspect, float(fh)
            x0 = min(max(cx - w / 2, 0.0), fw - w)
            y0 = min(max(cy - h / 2, 0.0), fh - h)
            windows[view_name] = (round(x0), round(y0), round(x0 + w), round(y0 + h))
        out[name] = windows
    return out

locator_skeleton

locator_skeleton(views: Sequence[str]) -> Skeleton

The point set a full-frame locator trains against: one point per view.

Channel k is view k's picture of the animal, so the pose path reads the views as keypoints. No edges and no mirror pairs: a locator learns the rig's fixed layout, which a mirrored frame never shows.

Source code in src/deeperfly/pose2d/training/dataset.py

def locator_skeleton(views: Sequence[str]) -> Skeleton:
    """The point set a full-frame locator trains against: one point per view.

    Channel `k` is view `k`'s picture of the animal, so the pose path reads the views as
    keypoints. No edges and no mirror pairs: a locator learns the rig's fixed layout,
    which a mirrored frame never shows.
    """
    return Skeleton(
        name="locator",
        point_names=tuple(views),
        edges=np.empty((0, 2), dtype=np.int64),
    )

rig_symmetries

rig_symmetries(
    proj: ndarray, tol_deg: float = 15.0
) -> list[tuple[bool, tuple[tuple[int, int], ...]]]

The world isometries that carry a rig's views onto each other by axis flips.

A mirror or a turn of the world that maps every view onto some view of the rig, up to a flip of that view's image axes and within tol_deg, is a picture the rig could have taken: the views stay consistent with each other, which a flip of every view about its own x axis is not on a rig whose cameras are turned against each other. A model that does not know which view is which never sees the views change places, so a symmetry is one axis flip per view.

Returns:

Type	Description
list[tuple[bool, tuple[tuple[int, int], ...]]]	(improper, signs) per symmetry, the identity first: improper is a mirror
list[tuple[bool, tuple[tuple[int, int], ...]]]	(left and right keypoints change places), signs[v] the (x, y) axis signs
list[tuple[bool, tuple[tuple[int, int], ...]]]	view v's picture goes through.

Source code in src/deeperfly/pose2d/training/dataset.py

def rig_symmetries(
    proj: np.ndarray, tol_deg: float = 15.0
) -> list[tuple[bool, tuple[tuple[int, int], ...]]]:
    """The world isometries that carry a rig's views onto each other by axis flips.

    A mirror or a turn of the world that maps every view onto some view of the rig,
    up to a flip of that view's image axes and within `tol_deg`, is a picture the rig
    could have taken: the views stay consistent with each other, which a
    flip of every view about its own x axis is not on a rig whose cameras are turned
    against each other. A model that does not know which view is which never sees
    the views change places, so a symmetry is one axis flip per view.

    Returns:
        `(improper, signs)` per symmetry, the identity first: `improper` is a mirror
        (left and right keypoints change places), `signs[v]` the `(x, y)` axis signs
        view `v`'s picture goes through.
    """
    frames = _view_frames(np.asarray(proj, float))
    n = len(frames)
    found: dict[tuple, tuple[bool, float]] = {}
    for w in range(n):
        for s in _AXIS_SIGNS:
            world = frames[0].T @ np.diag([*s, 1.0]) @ frames[w]
            signs: list = [None] * n
            worst = 0.0
            for v in range(n):
                best = (np.inf, -1, s)
                for w2 in range(n):
                    for s2 in _AXIS_SIGNS:
                        q = np.diag([*s2, 1.0]) @ frames[w2] @ (frames[v] @ world).T
                        if np.linalg.det(q) < 0:
                            continue
                        cos = np.clip((np.trace(q) - 1) / 2, -1, 1)
                        ang = float(np.degrees(np.arccos(cos)))
                        if ang < best[0]:
                            best = (ang, w2, s2)
                if best[0] > tol_deg or signs[best[1]] is not None:
                    break
                # Old view `best[1]` fills slot `v`, flipped by `best[2]`.
                signs[best[1]] = best[2]
                worst = max(worst, best[0])
            else:
                key = tuple(signs)
                if key not in found or worst < found[key][1]:
                    found[key] = (bool(np.linalg.det(world) < 0), worst)
    order = sorted(found, key=lambda k: (any(s != (1, 1) for s in k), found[k][0], k))
    return [(found[k][0], k) for k in order]

Label sources outside a project: the df3d corpus and a Lightning Pose export.

Pretraining stages, read into a project's SampleDataset shape with no calibration: no fill, no geometric loss. Df3dDataset reads the DeepFly3D annotations (Gunel et al. 2019, cleaned in the previous training repo), keypoint columns mapped onto the skeleton by name. LightningPoseDataset reads a multiview Lightning Pose export, its window affine taking the labels back to video pixels.

Df3dDataset

Bases: SampleDataset

The df3d corpus as (experiment, frame, camera) items grouped by frame.

Parameters:

Name	Type	Description	Default
root	str | Path	The corpus root (images/<experiment>/annotations.csv below it).	required
skeleton		The detector's skeleton; corpus columns are matched to it by name.	required
input_hw	tuple[int, int]	Network input (height, width).	required
margin_px	int | tuple[int, int]	The detector's off-frame margin in input pixels.	0
augment	Augment | None	Training-time augmentation, or None for the plain resize.	None
experiments	Sequence[str] | None	Which experiments to read (all by default).	None
drop_thorax_coxa_sources	Sequence[str]	Annotation sources whose *_thorax_coxa labels weigh 0; the corpus's sources.csv flags gizem_2022 as "annotated too high".	('gizem_2022',)

Source code in src/deeperfly/pose2d/training/sources.py

class Df3dDataset(SampleDataset):
    """The df3d corpus as `(experiment, frame, camera)` items grouped by frame.

    Args:
        root: The corpus root (`images/<experiment>/annotations.csv` below it).
        skeleton: The detector's skeleton; corpus columns are matched to it by name.
        input_hw: Network input `(height, width)`.
        margin_px: The detector's off-frame margin in input pixels.
        augment: Training-time augmentation, or `None` for the plain resize.
        experiments: Which experiments to read (all by default).
        drop_thorax_coxa_sources: Annotation sources whose `*_thorax_coxa` labels weigh
            0; the corpus's `sources.csv` flags `gizem_2022` as "annotated too high".
    """

    def __init__(
        self,
        root: str | Path,
        skeleton,
        *,
        input_hw: tuple[int, int],
        margin_px: int | tuple[int, int] = 0,
        augment: Augment | None = None,
        experiments: Sequence[str] | None = None,
        drop_thorax_coxa_sources: Sequence[str] = ("gizem_2022",),
        crop_tightness: tuple[float, float] | None = None,
        crop_jitter: float = 0.1,
        crop_eval_tightness: float = 1.25,
        crop_min_points: int = 4,
        crop_full_p: float = 0.1,
    ) -> None:
        """Index every annotated `(experiment, frame, camera)` of the corpus."""
        super().__init__(
            skeleton,
            input_hw=input_hw,
            margin_px=margin_px,
            augment=augment,
            crop_tightness=crop_tightness,
            crop_jitter=crop_jitter,
            crop_eval_tightness=crop_eval_tightness,
            crop_min_points=crop_min_points,
            crop_full_p=crop_full_p,
        )
        self.root = Path(root).expanduser()
        names = (
            list(experiments)
            if experiments is not None
            else sorted(p.name for p in (self.root / "images").iterdir())
        )
        k = len(self.keypoint_names)
        index = {n: i for i, n in enumerate(self.keypoint_names)}
        thc = {i for n, i in index.items() if n.endswith("_thorax_coxa")}
        self._paths: list[Path] = []
        self._size: dict[int, tuple[int, int]] = {}  # recording -> (h, w) source size
        self._xy: list[np.ndarray] = []
        self._w: list[np.ndarray] = []
        for name in names:
            path = self.root / "images" / name / "annotations.csv"
            with open(path, newline="") as f:
                reader = csv.DictReader(f)
                cols = [
                    (c[:-2], index[_rename(c[:-2])])
                    for c in reader.fieldnames
                    if c.endswith("_x") and _rename(c[:-2]) in index
                ]
                rows = list(reader)
            if not rows:
                continue
            r = len(self.entries)
            # The corpus names its seven views `camera_<n>`; that spelling is its own.
            view_names = tuple(f"camera_{c}" for c in range(7))
            self.entries.append(_Source(name, view_names))
            for row in rows:
                w_img, h_img = int(row["img_width"]), int(row["img_height"])
                xy = np.full((k, 2), np.nan, dtype=np.float32)
                for col, i in cols:
                    sx, sy = row[f"{col}_x"], row[f"{col}_y"]
                    if sx and sy:
                        xy[i] = (float(sx) * w_img, float(sy) * h_img)
                wt = np.isfinite(xy).all(axis=-1).astype(np.float32)
                if row["source"] in drop_thorax_coxa_sources:
                    wt[list(thc)] = 0.0
                    xy[list(thc)] = np.nan
                self.items.append((r, int(row["frame"]), int(row["camera"]), 0))
                self._paths.append(path.parent / row["image"])
                self._size.setdefault(r, (h_img, w_img))
                self._xy.append(xy)
                self._w.append(wt)
        self._lookup = {key[:3]: i for i, key in enumerate(self.items)}

    def _image(self, r: int, t: int, v: int):
        path = self._paths[self._lookup[(r, t, v)]]
        raw = np.array(Image.open(path).convert("L"))
        if raw.shape != self._size[r]:
            raise ValueError(
                f"{path} is {raw.shape}; its experiment's images are {self._size[r]}"
            )
        return raw, None

    def _target(self, r: int, t: int, v: int, i: int = 0):
        index = self._lookup[(r, t, v)]
        return self._xy[index], self._w[index], self._xy[index]

    def _base_affine(self, r: int, v: int) -> Affine2D:
        # Without a per-sample crop: the largest window at the input's aspect, centred,
        # so the 800x480 images are cut to 800x400 rather than stretched.
        return Affine2D.crop(aspect_window(self._size[r], self.input_hw), self.input_hw)

__init__

__init__(
    root: str | Path,
    skeleton,
    *,
    input_hw: tuple[int, int],
    margin_px: int | tuple[int, int] = 0,
    augment: Augment | None = None,
    experiments: Sequence[str] | None = None,
    drop_thorax_coxa_sources: Sequence[str] = (
        "gizem_2022",
    ),
    crop_tightness: tuple[float, float] | None = None,
    crop_jitter: float = 0.1,
    crop_eval_tightness: float = 1.25,
    crop_min_points: int = 4,
    crop_full_p: float = 0.1,
) -> None

Index every annotated (experiment, frame, camera) of the corpus.

Source code in src/deeperfly/pose2d/training/sources.py

def __init__(
    self,
    root: str | Path,
    skeleton,
    *,
    input_hw: tuple[int, int],
    margin_px: int | tuple[int, int] = 0,
    augment: Augment | None = None,
    experiments: Sequence[str] | None = None,
    drop_thorax_coxa_sources: Sequence[str] = ("gizem_2022",),
    crop_tightness: tuple[float, float] | None = None,
    crop_jitter: float = 0.1,
    crop_eval_tightness: float = 1.25,
    crop_min_points: int = 4,
    crop_full_p: float = 0.1,
) -> None:
    """Index every annotated `(experiment, frame, camera)` of the corpus."""
    super().__init__(
        skeleton,
        input_hw=input_hw,
        margin_px=margin_px,
        augment=augment,
        crop_tightness=crop_tightness,
        crop_jitter=crop_jitter,
        crop_eval_tightness=crop_eval_tightness,
        crop_min_points=crop_min_points,
        crop_full_p=crop_full_p,
    )
    self.root = Path(root).expanduser()
    names = (
        list(experiments)
        if experiments is not None
        else sorted(p.name for p in (self.root / "images").iterdir())
    )
    k = len(self.keypoint_names)
    index = {n: i for i, n in enumerate(self.keypoint_names)}
    thc = {i for n, i in index.items() if n.endswith("_thorax_coxa")}
    self._paths: list[Path] = []
    self._size: dict[int, tuple[int, int]] = {}  # recording -> (h, w) source size
    self._xy: list[np.ndarray] = []
    self._w: list[np.ndarray] = []
    for name in names:
        path = self.root / "images" / name / "annotations.csv"
        with open(path, newline="") as f:
            reader = csv.DictReader(f)
            cols = [
                (c[:-2], index[_rename(c[:-2])])
                for c in reader.fieldnames
                if c.endswith("_x") and _rename(c[:-2]) in index
            ]
            rows = list(reader)
        if not rows:
            continue
        r = len(self.entries)
        # The corpus names its seven views `camera_<n>`; that spelling is its own.
        view_names = tuple(f"camera_{c}" for c in range(7))
        self.entries.append(_Source(name, view_names))
        for row in rows:
            w_img, h_img = int(row["img_width"]), int(row["img_height"])
            xy = np.full((k, 2), np.nan, dtype=np.float32)
            for col, i in cols:
                sx, sy = row[f"{col}_x"], row[f"{col}_y"]
                if sx and sy:
                    xy[i] = (float(sx) * w_img, float(sy) * h_img)
            wt = np.isfinite(xy).all(axis=-1).astype(np.float32)
            if row["source"] in drop_thorax_coxa_sources:
                wt[list(thc)] = 0.0
                xy[list(thc)] = np.nan
            self.items.append((r, int(row["frame"]), int(row["camera"]), 0))
            self._paths.append(path.parent / row["image"])
            self._size.setdefault(r, (h_img, w_img))
            self._xy.append(xy)
            self._w.append(wt)
    self._lookup = {key[:3]: i for i, key in enumerate(self.items)}

LightningPoseDataset

Bases: SampleDataset

A multiview Lightning Pose project as (recording, moment, view) items.

Parameters:

Name	Type	Description	Default
root	str | Path	The project directory (<view>.csv, bboxes_<view>.csv, labeled-data/).	required
skeleton		The detector's skeleton; the CSV bodyparts are matched by name.	required
input_hw	tuple[int, int]	Network input (height, width), the exported images' size.	required
margin_px	int | tuple[int, int]	The detector's off-frame margin in input pixels.	0
augment	Augment | None	Training-time augmentation, or None for the plain resize.	None
views	Sequence[str]	The view names, in the project's order (config.yaml's view_names).	('rh', 'rm', 'rf', 'f', 'lf', 'lm', 'lh', 'h')
exclude	Sequence[str]	Recording names to leave out (a scored animal, say).	()
min_visible	float	The visible value a cell needs to weigh 1 (the export writes 2 for a real label and 0 for none).	1.0

Source code in src/deeperfly/pose2d/training/sources.py

class LightningPoseDataset(SampleDataset):
    """A multiview Lightning Pose project as `(recording, moment, view)` items.

    Args:
        root: The project directory (`<view>.csv`, `bboxes_<view>.csv`,
            `labeled-data/`).
        skeleton: The detector's skeleton; the CSV `bodyparts` are matched by name.
        input_hw: Network input `(height, width)`, the exported images' size.
        margin_px: The detector's off-frame margin in input pixels.
        augment: Training-time augmentation, or `None` for the plain resize.
        views: The view names, in the project's order (`config.yaml`'s `view_names`).
        exclude: Recording names to leave out (a scored animal, say).
        min_visible: The `visible` value a cell needs to weigh 1 (the export writes 2
            for a real label and 0 for none).
    """

    def __init__(
        self,
        root: str | Path,
        skeleton,
        *,
        input_hw: tuple[int, int],
        margin_px: int | tuple[int, int] = 0,
        augment: Augment | None = None,
        views: Sequence[str] = ("rh", "rm", "rf", "f", "lf", "lm", "lh", "h"),
        exclude: Sequence[str] = (),
        min_visible: float = 1.0,
    ) -> None:
        """Index every `(recording, moment, view)` with an image and a bounding box."""
        super().__init__(
            skeleton, input_hw=input_hw, margin_px=margin_px, augment=augment
        )
        self.root = Path(root).expanduser()
        views = tuple(views)
        k = len(self.keypoint_names)
        index = {n: i for i, n in enumerate(self.keypoint_names)}
        placeholders: dict[str, set[str]] = {}
        manifest = self.root / "export_manifest.json"
        if manifest.exists():
            for part in json.loads(manifest.read_text()).get("partial", []):
                placeholders[part["recording"]] = set(part.get("placeholder_views", []))
        # (recording, image file name) -> per-view (path, xy, weight, box)
        table: dict[tuple[str, str], dict[int, tuple]] = {}
        to_video_px: dict[tuple, Affine2D] = {}
        for v, view in enumerate(views):
            boxes = {}
            with open(self.root / f"bboxes_{view}.csv", newline="") as f:
                for row in csv.DictReader(f):
                    boxes[row[""]] = tuple(
                        float(row[key]) for key in ("x", "y", "h", "w")
                    )
            with open(self.root / f"{view}.csv", newline="") as f:
                rows = list(csv.reader(f))
            parts, coords = rows[1][1:], rows[2][1:]
            slot = [None] * len(parts)
            for j, (part, coord) in enumerate(zip(parts, coords, strict=True)):
                if part in index:
                    slot[j] = (index[part], coord)
            for row in rows[3:]:
                rel = row[0]
                rec_dir, fname = rel.split("/")[1], rel.split("/")[-1]
                rec = rec_dir[: -(len(view) + 1)]
                if rec in exclude or view in placeholders.get(rec, ()):
                    continue
                xy = np.full((k, 2), np.nan, dtype=np.float32)
                vis = np.zeros(k, dtype=np.float32)
                for j, val in enumerate(row[1:]):
                    if slot[j] is None or val == "":
                        continue
                    i, coord = slot[j]
                    if coord == "x":
                        xy[i, 0] = float(val)
                    elif coord == "y":
                        xy[i, 1] = float(val)
                    elif coord == "visible":
                        vis[i] = float(val)
                wt = (np.isfinite(xy).all(axis=-1) & (vis >= min_visible)).astype(
                    np.float32
                )
                xy[wt == 0] = np.nan
                box = boxes[rel]
                if box not in to_video_px:
                    x0, y0, bh, bw = box
                    to_video_px[box] = Affine2D.crop(
                        (x0, y0, x0 + bw, y0 + bh), self.input_hw
                    ).inverse()
                xy = to_video_px[box].apply(torch.from_numpy(xy)).numpy()
                table.setdefault((rec, fname), {})[v] = (
                    self.root / rel,
                    xy,
                    wt,
                    box,
                )
        self._data: list[tuple] = []
        rec_index: dict[str, int] = {}
        for (rec, fname), per_view in sorted(table.items()):
            if rec not in rec_index:
                rec_index[rec] = len(self.entries)
                self.entries.append(_Source(rec, views))
            r = rec_index[rec]
            t = int("".join(ch for ch in fname if ch.isdigit()) or len(self._data))
            for v in sorted(per_view):
                self.items.append((r, t, v, 0))
                self._data.append(per_view[v])
        self._lookup = {key[:3]: i for i, key in enumerate(self.items)}
        self._boxes = {
            (r, v): self._data[i][3] for i, (r, _, v, _) in enumerate(self.items)
        }

    def _image(self, r: int, t: int, v: int):
        path = self._data[self._lookup[(r, t, v)]][0]
        img = np.array(Image.open(path).convert("L"))
        if img.shape != tuple(self.input_hw):
            raise ValueError(
                f"{path} is {img.shape}, not the network input {self.input_hw}"
            )
        return img, self._base_affine(r, v)

    def _target(self, r: int, t: int, v: int, i: int = 0):
        _, xy, wt, _ = self._data[self._lookup[(r, t, v)]]
        return xy, wt, None  # the images are already cut: no per-sample crop

    def _base_affine(self, r: int, v: int) -> Affine2D:
        x, y, h, w = self._boxes[(r, v)]
        return Affine2D.crop((x, y, x + w, y + h), self.input_hw)

__init__

__init__(
    root: str | Path,
    skeleton,
    *,
    input_hw: tuple[int, int],
    margin_px: int | tuple[int, int] = 0,
    augment: Augment | None = None,
    views: Sequence[str] = (
        "rh",
        "rm",
        "rf",
        "f",
        "lf",
        "lm",
        "lh",
        "h",
    ),
    exclude: Sequence[str] = (),
    min_visible: float = 1.0,
) -> None

Index every (recording, moment, view) with an image and a bounding box.

Source code in src/deeperfly/pose2d/training/sources.py

def __init__(
    self,
    root: str | Path,
    skeleton,
    *,
    input_hw: tuple[int, int],
    margin_px: int | tuple[int, int] = 0,
    augment: Augment | None = None,
    views: Sequence[str] = ("rh", "rm", "rf", "f", "lf", "lm", "lh", "h"),
    exclude: Sequence[str] = (),
    min_visible: float = 1.0,
) -> None:
    """Index every `(recording, moment, view)` with an image and a bounding box."""
    super().__init__(
        skeleton, input_hw=input_hw, margin_px=margin_px, augment=augment
    )
    self.root = Path(root).expanduser()
    views = tuple(views)
    k = len(self.keypoint_names)
    index = {n: i for i, n in enumerate(self.keypoint_names)}
    placeholders: dict[str, set[str]] = {}
    manifest = self.root / "export_manifest.json"
    if manifest.exists():
        for part in json.loads(manifest.read_text()).get("partial", []):
            placeholders[part["recording"]] = set(part.get("placeholder_views", []))
    # (recording, image file name) -> per-view (path, xy, weight, box)
    table: dict[tuple[str, str], dict[int, tuple]] = {}
    to_video_px: dict[tuple, Affine2D] = {}
    for v, view in enumerate(views):
        boxes = {}
        with open(self.root / f"bboxes_{view}.csv", newline="") as f:
            for row in csv.DictReader(f):
                boxes[row[""]] = tuple(
                    float(row[key]) for key in ("x", "y", "h", "w")
                )
        with open(self.root / f"{view}.csv", newline="") as f:
            rows = list(csv.reader(f))
        parts, coords = rows[1][1:], rows[2][1:]
        slot = [None] * len(parts)
        for j, (part, coord) in enumerate(zip(parts, coords, strict=True)):
            if part in index:
                slot[j] = (index[part], coord)
        for row in rows[3:]:
            rel = row[0]
            rec_dir, fname = rel.split("/")[1], rel.split("/")[-1]
            rec = rec_dir[: -(len(view) + 1)]
            if rec in exclude or view in placeholders.get(rec, ()):
                continue
            xy = np.full((k, 2), np.nan, dtype=np.float32)
            vis = np.zeros(k, dtype=np.float32)
            for j, val in enumerate(row[1:]):
                if slot[j] is None or val == "":
                    continue
                i, coord = slot[j]
                if coord == "x":
                    xy[i, 0] = float(val)
                elif coord == "y":
                    xy[i, 1] = float(val)
                elif coord == "visible":
                    vis[i] = float(val)
            wt = (np.isfinite(xy).all(axis=-1) & (vis >= min_visible)).astype(
                np.float32
            )
            xy[wt == 0] = np.nan
            box = boxes[rel]
            if box not in to_video_px:
                x0, y0, bh, bw = box
                to_video_px[box] = Affine2D.crop(
                    (x0, y0, x0 + bw, y0 + bh), self.input_hw
                ).inverse()
            xy = to_video_px[box].apply(torch.from_numpy(xy)).numpy()
            table.setdefault((rec, fname), {})[v] = (
                self.root / rel,
                xy,
                wt,
                box,
            )
    self._data: list[tuple] = []
    rec_index: dict[str, int] = {}
    for (rec, fname), per_view in sorted(table.items()):
        if rec not in rec_index:
            rec_index[rec] = len(self.entries)
            self.entries.append(_Source(rec, views))
        r = rec_index[rec]
        t = int("".join(ch for ch in fname if ch.isdigit()) or len(self._data))
        for v in sorted(per_view):
            self.items.append((r, t, v, 0))
            self._data.append(per_view[v])
    self._lookup = {key[:3]: i for i, key in enumerate(self.items)}
    self._boxes = {
        (r, v): self._data[i][3] for i, (r, _, v, _) in enumerate(self.items)
    }

Unlabeled moments as training samples: a teacher's consensus (ADR-0071).

A teacher's 2D over every view of a sampled moment is triangulated robustly and the consensus reprojected into every view as the target, the cells it corrected flagged. UnlabeledDataset serves them through the labeled pipeline's crops, augmentation, margin rule and grouped batches, prism mirrors included, from a native FrameCache.

Without a teacher artifact the teacher is each project's merged detections (ADR-0083), and in both cases the windows are the static crops those detections were read through.

UnlabeledDataset

Bases: SampleDataset

One item per (project, frame, view) of teacher-consensus targets.

Built by build_unlabeled from projects whose targets are solved. The crop is drawn around the consensus points, its tightness capped below the labeled range since the stored window is only 1.5x the animal; no whole-frame letterbox.

Source code in src/deeperfly/pose2d/training/unlabeled.py

class UnlabeledDataset(SampleDataset):
    """One item per `(project, frame, view)` of teacher-consensus targets.

    Built by [build_unlabeled][deeperfly.pose2d.training.unlabeled.build_unlabeled] from
    projects whose targets are solved. The crop is drawn around the consensus points,
    its tightness capped below the labeled range since the stored window is only 1.5x
    the animal; no whole-frame letterbox.
    """

    def __init__(
        self,
        skeleton,
        entries: Sequence[_UnlabeledEntry],
        *,
        input_hw: tuple[int, int],
        margin_px: int | tuple[int, int],
        augment,
        crop_tightness: tuple[float, float] = (0.75, 1.6),
        crop_jitter: float = 0.1,
        crop_eval_tightness: float = 1.25,
        crop_min_points: int = 4,
    ) -> None:
        """Index every view of every solved moment of `entries`."""
        super().__init__(
            skeleton,
            input_hw=input_hw,
            margin_px=margin_px,
            augment=augment,
            crop_tightness=crop_tightness,
            crop_jitter=crop_jitter,
            crop_eval_tightness=crop_eval_tightness,
            crop_min_points=crop_min_points,
            crop_full_p=0.0,
        )
        self.entries = list(entries)
        for r, entry in enumerate(self.entries):
            for t in sorted(entry.targets):
                for v in range(len(entry.view_names)):
                    self.items.append((r, t, v, 0))

    def _image(self, r: int, t: int, v: int):
        cache = self.entries[r].cache
        return cache.get(v, t), cache.image_affine(v)

    def _target(self, r: int, t: int, v: int, i: int = 0):
        xy, w, ref = self.entries[r].targets[t]
        return xy[v], w[v], ref[v]

    def _mirrored(self, r: int, v: int) -> bool:
        return self.entries[r].cache.mirrored(v)

    def _base_affine(self, r: int, v: int) -> Affine2D:
        cache = self.entries[r].cache
        h, w = cache.stored_hw(v)
        box = cache.windows[cache.view_names[v]] if cache.window_mode else (0, 0, w, h)
        return input_window(box, self.input_hw, self._mirrored(r, v))

    def n_moments(self) -> int:
        """How many `(project, frame)` moments carry targets."""
        return sum(len(entry.targets) for entry in self.entries)

__init__

__init__(
    skeleton,
    entries: Sequence[_UnlabeledEntry],
    *,
    input_hw: tuple[int, int],
    margin_px: int | tuple[int, int],
    augment,
    crop_tightness: tuple[float, float] = (0.75, 1.6),
    crop_jitter: float = 0.1,
    crop_eval_tightness: float = 1.25,
    crop_min_points: int = 4,
) -> None

Index every view of every solved moment of entries.

Source code in src/deeperfly/pose2d/training/unlabeled.py

def __init__(
    self,
    skeleton,
    entries: Sequence[_UnlabeledEntry],
    *,
    input_hw: tuple[int, int],
    margin_px: int | tuple[int, int],
    augment,
    crop_tightness: tuple[float, float] = (0.75, 1.6),
    crop_jitter: float = 0.1,
    crop_eval_tightness: float = 1.25,
    crop_min_points: int = 4,
) -> None:
    """Index every view of every solved moment of `entries`."""
    super().__init__(
        skeleton,
        input_hw=input_hw,
        margin_px=margin_px,
        augment=augment,
        crop_tightness=crop_tightness,
        crop_jitter=crop_jitter,
        crop_eval_tightness=crop_eval_tightness,
        crop_min_points=crop_min_points,
        crop_full_p=0.0,
    )
    self.entries = list(entries)
    for r, entry in enumerate(self.entries):
        for t in sorted(entry.targets):
            for v in range(len(entry.view_names)):
                self.items.append((r, t, v, 0))

n_moments

n_moments() -> int

How many (project, frame) moments carry targets.

Source code in src/deeperfly/pose2d/training/unlabeled.py

def n_moments(self) -> int:
    """How many `(project, frame)` moments carry targets."""
    return sum(len(entry.targets) for entry in self.entries)

build_unlabeled

build_unlabeled(
    projects: Mapping[str, Project],
    *,
    skeleton: Skeleton | None = None,
    teacher: str | Path | None,
    n_per_project: int = 200,
    agree_px: float,
    min_views: int,
    conf_min: float,
    input_hw: tuple[int, int],
    margin_px,
    augment,
    geom_weight: float = 1.0,
    skip_frames: Mapping[str, Sequence[int]] | None = None,
    guard_frames: int = 0,
    crop_tightness: tuple[float, float] = (0.75, 1.6),
    cache_dir: str | Path | None = None,
    device: str | None = None,
    log=print,
) -> UnlabeledDataset

Sample, cache, teach, solve the consensus: the dataset over projects.

Per project (name -> Project), n_per_project frames spread over the video (the labeled and skip_frames ones skipped, guard_frames either side) are cached behind the discovered windows, taught and solved through the active calibration, and the targets stored beside the frame cache. teacher=None takes the points from each project's merged detections, sampling only detected frames (ADR-0071, ADR-0083). A project without an active calibration is skipped. cache_dir is <project>/cache/unlabeled by default.

Source code in src/deeperfly/pose2d/training/unlabeled.py

def build_unlabeled(
    projects: Mapping[str, Project],
    *,
    skeleton: Skeleton | None = None,
    teacher: str | Path | None,
    n_per_project: int = 200,
    agree_px: float,
    min_views: int,
    conf_min: float,
    input_hw: tuple[int, int],
    margin_px,
    augment,
    geom_weight: float = 1.0,
    skip_frames: Mapping[str, Sequence[int]] | None = None,
    guard_frames: int = 0,
    crop_tightness: tuple[float, float] = (0.75, 1.6),
    cache_dir: str | Path | None = None,
    device: str | None = None,
    log=print,
) -> UnlabeledDataset:
    """Sample, cache, teach, solve the consensus: the dataset over `projects`.

    Per project (`name -> Project`), `n_per_project` frames spread over the video
    (the labeled and `skip_frames` ones skipped, `guard_frames` either side) are
    cached behind the discovered windows, taught and solved through the active
    calibration, and the targets stored beside the frame cache. `teacher=None` takes
    the points from each project's merged detections, sampling only detected frames
    (ADR-0071, ADR-0083). A project without an active calibration is skipped.
    `cache_dir` is `<project>/cache/unlabeled` by default.
    """
    from ..models import ModelSpec, _load_mvt

    projects = dict(projects)
    skeleton = _skeleton_of(projects, skeleton)
    from_detections = teacher is None
    teacher = None if from_detections else Path(teacher).expanduser()
    model = None
    entries: list[_UnlabeledEntry] = []
    for name, project in projects.items():
        root = (
            Path(cache_dir)
            if cache_dir is not None
            else project.folder / "cache" / "unlabeled"
        )
        views = tuple(project.views)
        labels = project.label_files().load()
        frame_hw = project.frame_sizes()
        rig = active_rig(project, frame_hw)
        if rig is None:
            log(f"unlabeled {name}: no active calibration to solve through, skipped")
            continue
        skip = list(labels.labeled_frames()) + list((skip_frames or {}).get(name, ()))
        detections = Detections.open(
            project,
            view_names=views,
            keypoint_names=skeleton.point_names,
            n_frames=labels.n_frames,
        )
        try:
            if from_detections:
                frames = _detected_frames(detections, skip, n_per_project, guard_frames)
                if not frames:
                    log(f"unlabeled {name}: no detected frame to draw from, skipped")
                    continue
            else:
                frames = sample_frames(
                    labels.n_frames, skip, n_per_project, guard=guard_frames
                )
            windows = discovered_windows(
                project,
                input_hw,
                detections=detections,
                whole_view=from_detections,
                frame_hw=frame_hw,
            )
            # A window that is the whole frame is the frame: store it once per video
            # rather than once per view, or a rig that reads three views out of one
            # video keeps three copies of every pixel and decodes it three times.
            whole = windows == full_frame_crops(project, frame_hw)
            cache = FrameCache(
                project,
                name,
                input_hw=None,
                crops=None if whole else windows,
                cache_dir=root,
                frame_hw=frame_hw,
            )
            cache.ensure(frames)
            settings = {
                "agree_px": agree_px,
                "min_views": min_views,
                "conf_min": conf_min,
                "geom_weight": geom_weight,
            }
            digest = _teacher_digest(
                teacher,
                cache,
                frames,
                None if teacher else detections,
                rig=rig,
                **settings,
            )
            tpath = root / "teacher" / name / f"{digest}.h5"
            if tpath.exists():
                solved = _solved_targets(tpath)
            else:
                if from_detections:
                    points, conf = detections.read(frames)
                    txy = np.moveaxis(points, 1, 0)
                    tconf = np.moveaxis(conf, 1, 0)
                else:
                    if model is None:
                        model = _load_mvt(ModelSpec("teacher", "mvt", str(teacher)))
                        model.to(
                            device or ("cuda" if torch.cuda.is_available() else "cpu")
                        )
                    txy, tconf = run_teacher(
                        model, cache, frames, flip_perm=skeleton.flip_perm()
                    )
                solved = _solve_targets(
                    tpath, rig, txy, tconf, frames, teacher, settings
                )
        finally:
            detections.close()
        fr, xy, w, ref, corrected = solved
        targets = {}
        for i, t in enumerate(fr):
            if int((w[i] > 0).any(axis=0).sum()) >= MIN_CONSENSUS_KEYPOINTS:
                targets[int(t)] = (xy[i], w[i], ref[i])
        n_cells = w.size
        log(
            f"unlabeled {name}: {len(targets)}/{len(fr)} moments kept, consensus on "
            f"{100 * float((w > 0).mean()):.0f}% of cells, geometry corrected "
            f"{100 * float(corrected.sum() / max(n_cells, 1)):.1f}%"
        )
        entries.append(
            _UnlabeledEntry(
                name, views, cache, _pinhole_projections(rig, views), targets
            )
        )
    if model is not None:
        del model
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    return UnlabeledDataset(
        skeleton,
        entries,
        input_hw=input_hw,
        margin_px=margin_px,
        augment=augment,
        crop_tightness=crop_tightness,
    )

consensus_targets

consensus_targets(
    rig,
    xy: ndarray,
    conf: ndarray,
    *,
    agree_px: float,
    min_views: int,
    conf_min: float,
    geom_weight: float = 1.0,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]

Teacher points (V, K, 2) and confidences -> the consensus targets (ADR-0071).

Returns (targets, weights, reference, corrected): the consensus reprojected into every view, 1 where a consensus exists, the target or else the teacher's point for the crop, and the cells whose own detection lay more than agree_px away. geom_weight blends the target toward the teacher's own pixel.

Source code in src/deeperfly/pose2d/training/unlabeled.py

def consensus_targets(
    rig,
    xy: np.ndarray,
    conf: np.ndarray,
    *,
    agree_px: float,
    min_views: int,
    conf_min: float,
    geom_weight: float = 1.0,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Teacher points `(V, K, 2)` and confidences -> the consensus targets (ADR-0071).

    Returns `(targets, weights, reference, corrected)`: the consensus reprojected into
    every view, 1 where a consensus exists, the target or else the teacher's point for
    the crop, and the cells whose own detection lay more than `agree_px` away.
    `geom_weight` blends the target toward the teacher's own pixel.
    """
    pts = np.asarray(xy, dtype=float).copy()
    pts[np.asarray(conf) < conf_min] = np.nan
    pts3d, inliers = triangulate_ransac(
        rig, pts, threshold=float(agree_px), min_views=int(min_views)
    )
    ok = np.isfinite(pts3d).all(axis=-1)  # (K,)
    proj = np.full(pts.shape, np.nan)
    if ok.any():
        proj[:, ok] = np.asarray(rig.project(pts3d[ok]))
    weights = np.repeat(ok[None, :].astype(np.float32), pts.shape[0], axis=0)
    has = np.isfinite(proj).all(axis=-1)
    weights[~has] = 0.0
    # Only where the teacher has a pixel to blend with: a view it left NaN keeps the
    # reprojection, so `geom_weight = 1` is the old rule cell for cell.
    a = float(geom_weight)
    raw = np.asarray(xy, dtype=float)
    blend = np.where(np.isfinite(raw), (1.0 - a) * raw + a * proj, proj)
    targets = np.where(has[..., None], blend, np.nan).astype(np.float32)
    reference = np.where(has[..., None], proj, xy).astype(np.float32)
    corrected = has & ~inliers & np.isfinite(xy).all(axis=-1)
    return targets, weights, reference, corrected

detected_crops

detected_crops(
    detections: Detections, n: int = CROP_SAMPLES
) -> dict[str, tuple[int, int, int, int]] | None

The one static crop per view the detections were read through, or None.

Compared over n covered frames spread over the project; None when nothing is covered, a run stored no crops, or the crop moves (a per-frame window, or runs that discovered different ones).

Source code in src/deeperfly/pose2d/training/unlabeled.py

def detected_crops(
    detections: Detections, n: int = CROP_SAMPLES
) -> dict[str, tuple[int, int, int, int]] | None:
    """The one static crop per view the detections were read through, or `None`.

    Compared over `n` covered frames spread over the project; `None` when nothing is
    covered, a run stored no crops, or the crop moves (a per-frame window, or runs
    that discovered different ones).
    """
    covered = detections.covered_frames()
    if covered.size == 0:
        return None
    take = np.unique(np.round(np.linspace(0, covered.size - 1, num=n)).astype(int))
    boxes = [detections.crops(int(t)) for t in covered[take]]
    first = boxes[0]
    if first is None or not np.isfinite(first).all():
        return None
    if any(b is None or not np.array_equal(b, first) for b in boxes[1:]):
        return None
    out = {}
    for i, view in enumerate(detections.view_names):
        x0, y0, x1, y1 = (round(float(v)) for v in first[i])
        out[view] = (x0, y0, x1, y1)
    return out

discovered_windows

discovered_windows(
    project: Project,
    input_hw: tuple[int, int],
    *,
    detections: Detections | None = None,
    whole_view: bool = False,
    frame_hw: Mapping[str, tuple[int, int]] | None = None,
) -> dict[str, tuple[int, int, int, int]]

Per-view (x0, y0, x1, y1) video windows around the animal.

The static crop the detections were read through (detected_crops); failing that, each view whole with whole_view, for a stream whose points come from the detections, or else the largest centered window at the input's aspect. frame_hw is each view's frame (height, width), probed when absent.

Raises:

Type	Description
ValueError	Without a static crop or whole_view, when that centered window covers less than half a view: a prism strip, whose animal is outside it too often (ADR-0071).

Source code in src/deeperfly/pose2d/training/unlabeled.py

def discovered_windows(
    project: Project,
    input_hw: tuple[int, int],
    *,
    detections: Detections | None = None,
    whole_view: bool = False,
    frame_hw: Mapping[str, tuple[int, int]] | None = None,
) -> dict[str, tuple[int, int, int, int]]:
    """Per-view `(x0, y0, x1, y1)` video windows around the animal.

    The static crop the `detections` were read through
    ([detected_crops][deeperfly.pose2d.training.unlabeled.detected_crops]); failing
    that, each view whole with `whole_view`, for a stream whose points come from the
    detections, or else the largest centered window at the input's aspect.
    `frame_hw` is each view's frame `(height, width)`, probed when absent.

    Raises:
        ValueError: Without a static crop or `whole_view`, when that centered window
            covers less than half a view: a prism strip, whose animal is outside it
            too often (ADR-0071).
    """
    crops = None if detections is None else detected_crops(detections)
    if crops is not None:
        return crops
    sizes = project.frame_sizes() if frame_hw is None else frame_hw
    if whole_view:
        return full_frame_crops(project, sizes)
    out = {}
    for c in project.views:
        h, w = sizes[c]
        box = tuple(round(float(v)) for v in aspect_window((h, w), input_hw))
        covered = (box[2] - box[0]) * (box[3] - box[1]) / max(h * w, 1)
        if covered < 0.5:
            raise ValueError(
                f"{project.folder}/{c}: no static crop in the project's detections, "
                f"and the centered {input_hw[1]}:{input_hw[0]} window covers "
                f"{100 * covered:.0f}% of the {w}x{h} view -- too little of it to "
                "find the animal behind. Run pose2d on the project first (its "
                "static crops are the windows), or pass whole_view with teacher "
                "points that came from elsewhere."
            )
        out[c] = box
    return out

run_teacher

run_teacher(
    model,
    cache: FrameCache,
    frames: Sequence[int],
    *,
    flip_perm,
    flip_tta: bool = True,
    moments_per_batch: int = 4,
) -> tuple[np.ndarray, np.ndarray]

The teacher's (T, V, K, 2) video-pixel points and (T, V, K) confidences.

Every view's stored window resized to the model input, as a run's [pose2d] does, the whole group at once; with flip_tta the mirrored group is averaged in. The points are on the pixel-center grid whatever the teacher was trained on (ADR-0078).

Source code in src/deeperfly/pose2d/training/unlabeled.py

@torch.no_grad()
def run_teacher(
    model,
    cache: FrameCache,
    frames: Sequence[int],
    *,
    flip_perm,
    flip_tta: bool = True,
    moments_per_batch: int = 4,
) -> tuple[np.ndarray, np.ndarray]:
    """The teacher's `(T, V, K, 2)` video-pixel points and `(T, V, K)` confidences.

    Every view's stored window resized to the model input, as a run's `[pose2d]` does,
    the whole group at once; with `flip_tta` the mirrored group is averaged in. The
    points are on the pixel-center grid whatever the teacher was trained on (ADR-0078).
    """
    from ..infer import decode_points, joint_channels, prepare_images
    from ..models import grid_shift

    model.eval()
    dev = next(model.parameters()).device
    h, w = model.input_hw
    views = cache.view_names
    n_views = len(views)
    perm = torch.as_tensor(np.asarray(flip_perm), device=dev, dtype=torch.long)
    inv = []
    for v in range(n_views):
        # The student's own map, mirror included: a prism view's window comes out
        # reflected (ADR-0053), so the teacher sees the picture it trained on.
        to_input = input_window(cache.windows[views[v]], (h, w), cache.mirrored(v))
        inv.append(to_input.inverse())

    def one(v: int, t: int):
        img = prepare_images(cache.get(v, t)[None], (h, w), dev)[0]
        # Mirroring the window is a flip of the input, which is what `to_input`
        # above composes and what the student's own samples are warped through.
        return img.flip(-1) if cache.mirrored(v) else img

    xy_out, conf_out = [], []
    for start in range(0, len(frames), moments_per_batch):
        chunk = frames[start : start + moments_per_batch]
        imgs = torch.stack(
            [torch.stack([one(v, t) for v in range(n_views)]) for t in chunk]
        )  # (B, V, 1, h, w)
        b = imgs.shape[0]
        c, k = model.num_classes, model.n_points
        # A center channel is the last of each view's block, so the field is folded per
        # view and the joint channels decoded alone.
        field = model(imgs)
        field = field.reshape(b * n_views, c, *field.shape[-2:])
        pred, conf = decode_points(model, joint_channels(model, field))
        pred = pred.float().reshape(b, n_views, k, 2)
        conf = conf.float().reshape(b, n_views, k)
        if flip_tta:
            mirrored = model(imgs.flip(-1))
            mirrored = mirrored.reshape(b * n_views, c, *mirrored.shape[-2:])
            pm, _ = decode_points(model, joint_channels(model, mirrored))
            pm = pm.float().reshape(b, n_views, k, 2).index_select(2, perm).clone()
            pm[..., 0] = (w - 1 - pm[..., 0] * w) / w
            pred = 0.5 * (pred + pm)
        px = pred * pred.new_tensor([w, h])
        for i in range(b):
            xy = torch.stack(
                [inv[v].to(device=dev).apply(px[i, v]) for v in range(n_views)]
            )
            xy_out.append((xy + grid_shift(model)).cpu().numpy())
            conf_out.append(conf[i].cpu().numpy())
    return np.asarray(xy_out, dtype=np.float32), np.asarray(conf_out, dtype=np.float32)

sample_frames

sample_frames(
    n_frames: int,
    labeled: Sequence[int],
    n: int,
    *,
    guard: int = 0,
) -> list[int]

n frame indices spread evenly over the video, skipping the labeled ones.

guard widens each skipped frame into a band either side, for a fold that holds out frames of a project the stream still draws from.

Source code in src/deeperfly/pose2d/training/unlabeled.py

def sample_frames(
    n_frames: int, labeled: Sequence[int], n: int, *, guard: int = 0
) -> list[int]:
    """`n` frame indices spread evenly over the video, skipping the labeled ones.

    `guard` widens each skipped frame into a band either side, for a fold that holds
    out frames of a project the stream still draws from.
    """
    if n <= 0 or n_frames <= 0:
        return []
    pick = np.unique(np.round(np.linspace(0, n_frames - 1, num=n)).astype(int))
    skip = np.asarray(sorted({int(t) for t in labeled}), dtype=np.int64)
    if skip.size == 0:
        return [int(t) for t in pick]
    if guard <= 0:
        return [int(t) for t in pick if int(t) not in set(skip.tolist())]
    near = np.abs(pick[:, None] - skip[None, :]).min(axis=1)
    return [int(t) for t, d in zip(pick, near, strict=True) if d > guard]

Batch plumbing for the loop and the evaluator: view flattening, the inverse affine.

A batch is flat ((B, 1, H, W) images, (B, K, ...) labels) or grouped ((B, V, 1, H, W), (B, V, K, ...) and a view_valid (B, V) mask); forward_flat hands every consumer the flat layout.

flatten_views

flatten_views(
    batch: dict,
) -> tuple[dict, tuple[int, int] | None]

A grouped batch in the flat (B * V, ...) layout; a flat batch as is.

Per-group fields repeat per view; padded slots stay in (weight 0), and the second return, (B, V), folds the rows back into groups.

Source code in src/deeperfly/pose2d/training/batch.py

def flatten_views(batch: dict) -> tuple[dict, tuple[int, int] | None]:
    """A grouped batch in the flat `(B * V, ...)` layout; a flat batch as is.

    Per-group fields repeat per view; padded slots stay in (weight 0), and the second
    return, `(B, V)`, folds the rows back into groups.
    """
    if batch["image"].ndim == 4:
        out = dict(batch)
        out.setdefault(
            "view_valid",
            torch.ones(
                batch["image"].shape[0], dtype=torch.bool, device=batch["image"].device
            ),
        )
        return out, None
    b, v = batch["image"].shape[:2]
    out = {}
    for key, val in batch.items():
        if key in _PER_VIEW or key == "view_valid":
            out[key] = val.reshape(b * v, *val.shape[2:])
        elif torch.is_tensor(val):
            out[key] = val.repeat_interleave(v, dim=0)
        else:
            out[key] = val
    return out, (b, v)

forward_flat

forward_flat(
    model, batch: dict, *, mask_ratio=None
) -> tuple[object, dict, tuple[int, int] | None]

Run model on batch and return (outputs, flat_batch, group_shape).

The outputs are flat whatever the model: a joint-views model is fed the grouped images, refuses a flat batch and takes mask_ratio; any other sees the flat images.

Source code in src/deeperfly/pose2d/training/batch.py

def forward_flat(
    model, batch: dict, *, mask_ratio=None
) -> tuple[object, dict, tuple[int, int] | None]:
    """Run `model` on `batch` and return `(outputs, flat_batch, group_shape)`.

    The outputs are flat whatever the model: a joint-views model is fed the grouped
    images, refuses a flat batch and takes `mask_ratio`; any other sees the flat images.
    """
    flat, shape = flatten_views(batch)
    if getattr(model, "joint_views", False):
        if shape is None:
            raise ValueError(
                "a joint-views model needs grouped (B, V, 1, H, W) batches"
            )
        kw = {} if mask_ratio is None else {"mask_ratio": mask_ratio}
        out = model(batch["image"], view_valid=batch.get("view_valid"), **kw)
        b, v = shape
        k = out.shape[1] // v
        return out.reshape(b * v, k, *out.shape[2:]), flat, shape
    return model(flat["image"]), flat, shape

to_device

to_device(batch: dict, device) -> dict

Every tensor of batch moved to device (non-blocking); the rest as is.

Source code in src/deeperfly/pose2d/training/batch.py

def to_device(batch: dict, device) -> dict:
    """Every tensor of `batch` moved to `device` (non-blocking); the rest as is."""
    return {
        k: v.to(device, non_blocking=True) if torch.is_tensor(v) else v
        for k, v in batch.items()
    }

to_video_px

to_video_px(xy_norm, affine, input_hw)

Input-normalized (N, K, 2) -> video pixels through the inverse of affine.

affine (N, 2, 3) is the sample's video -> input map, flip included.

Source code in src/deeperfly/pose2d/training/batch.py

def to_video_px(xy_norm, affine, input_hw):
    """Input-normalized `(N, K, 2)` -> video pixels through the inverse of `affine`.

    `affine (N, 2, 3)` is the sample's video -> input map, flip included.
    """
    h, w = input_hw
    px = xy_norm * xy_norm.new_tensor([w, h])
    n = affine.shape[0]
    full = torch.cat(
        [
            affine.to(px.dtype),
            px.new_tensor([0.0, 0.0, 1.0]).expand(n, 1, 3),
        ],
        dim=1,
    )
    inv = torch.linalg.inv(full)[:, :2]  # (N, 2, 3)
    return px @ inv[:, :, :2].transpose(1, 2) + inv[:, None, :, 2]

The training loop for every detector family, and the per-head supervised loss.

supervised_loss matches the model's output_kind: the heatmap head's Gaussian MSE, SimCC's per-axis KL, or a probability field's cross-entropy and coordinate term, read through a local expectation (ADR-0072). batch_loss adds the pair-triangulation loss to a calibrated grouped batch, unlabeled_loss is the distillation term (ADR-0071), and fit is a minimal AdamW and cosine loop with bf16 autocast, an optional EMA, the unlabeled stream and a periodic evaluation.

batch_loss

batch_loss(
    model,
    batch: dict,
    *,
    sigma: float = 2.0,
    fg_weight: float = 30.0,
    field_weight: float = 0.0,
    center_weight: float = 0.25,
    coord_weight: float = 0.0,
    coord_radius: int = 3,
    coord_huber: float = 4.0,
    tri_weight: float = 0.0,
    tri_huber: float = 0.0,
    tri_max_world: float = 0.0,
    train_decode: str = "local",
)

The training loss of one batch, flat or grouped: (loss, aux).

The supervised term runs over every flat row (a padded view weighs 0). With tri_weight > 0 and a calibrated grouped batch the pair-triangulation loss is added per group_ok group, read through training_decode(kind=train_decode).

Source code in src/deeperfly/pose2d/training/loop.py

def batch_loss(
    model,
    batch: dict,
    *,
    sigma: float = 2.0,
    fg_weight: float = 30.0,
    field_weight: float = 0.0,
    center_weight: float = 0.25,
    coord_weight: float = 0.0,
    coord_radius: int = 3,
    coord_huber: float = 4.0,
    tri_weight: float = 0.0,
    tri_huber: float = 0.0,
    tri_max_world: float = 0.0,
    train_decode: str = "local",
):
    """The training loss of one batch, flat or grouped: `(loss, aux)`.

    The supervised term runs over every flat row (a padded view weighs 0). With
    `tri_weight > 0` and a calibrated grouped batch the pair-triangulation loss is added
    per `group_ok` group, read through `training_decode(kind=train_decode)`.
    """
    weight_key = "weight" if "weight" in batch else "vis"
    out, flat, shape = forward_flat(model, batch)
    loss, aux = supervised_loss(
        model,
        out,
        flat["kp"],
        flat[weight_key],
        sigma=sigma,
        fg_weight=fg_weight,
        field_weight=field_weight,
        center_weight=center_weight,
        coord_weight=coord_weight,
        coord_radius=coord_radius,
        coord_huber=coord_huber,
        centers=flat.get("centers"),
        center_valid=flat.get("center_valid"),
    )
    if tri_weight > 0 and shape is not None and "proj" in flat:
        b, v = shape
        pred_norm = training_decode(model, out, kind=train_decode, radius=coord_radius)
        pred_video_px = to_video_px(pred_norm, flat["affine"], model.input_hw)
        w = flat[weight_key].reshape(b, v, -1)
        group_ok = flat["group_ok"].reshape(b, v)[:, 0] if "group_ok" in flat else None
        tri, tri_aux = triangulation_loss(
            pred_video_px.reshape(b, v, -1, 2),
            flat["gt_video_px"].reshape(b, v, -1, 2),
            w,
            flat["proj"].reshape(b, v, 3, 4),
            view_valid=flat["view_valid"].reshape(b, v),
            group_ok=group_ok,
            huber=tri_huber,
            max_world=tri_max_world,
        )
        loss = loss + tri_weight * tri
        aux = {**aux, **tri_aux}
    return loss, aux

center_cross_entropy

center_cross_entropy(
    model, field, centers_norm, valid, *, sigma: float = 2.0
)

Cross-entropy between the center channel and the frame's animal centers.

Every valid center of (N, A, 2) renders the joints' Gaussian, and their renormalized sum is the target: one field with a peak per animal, as the top-k readout reads it back. A row with no valid center weighs 0.

Source code in src/deeperfly/pose2d/training/loop.py

def center_cross_entropy(model, field, centers_norm, valid, *, sigma: float = 2.0):
    """Cross-entropy between the center channel and the frame's animal centers.

    Every valid center of `(N, A, 2)` renders the joints' Gaussian, and their
    renormalized sum is the target: one field with a peak per animal, as the top-k
    readout reads it back. A row with no valid center weighs 0.
    """
    n, h, w = field.shape
    geom = _field_geometry(model, field[:, None])
    with torch.no_grad():
        cells = geom.to_field_cells(centers_norm)
        target = heatmap.render_gaussian_targets(
            cells, valid.to(cells.dtype), hw=(h, w), sigma=sigma
        ).sum(dim=1)
        target = target / target.sum(dim=(-1, -2), keepdim=True).clamp_min(1e-12)
    if getattr(model, "field_is_prob", False):
        logp = torch.log(field.float().clamp_min(1e-12))
    else:
        logp = torch.log_softmax(field.float().flatten(1), dim=-1).view(n, h, w)
    ce = -(target.to(logp.dtype) * logp).sum(dim=(-1, -2))  # (N,)
    seen = valid.any(dim=-1).to(ce.dtype)
    loss = (ce * seen).sum() / seen.sum().clamp_min(1.0)
    return loss, {"center_ce": loss.detach()}

field_cross_entropy

field_cross_entropy(
    model, field, kp_norm, weight, *, sigma: float = 2.0
)

Cross-entropy between a probability field and a normalized Gaussian target.

The cells are placed by the geometry the readout inverts; raw logits go through a spatial log-softmax.

Source code in src/deeperfly/pose2d/training/loop.py

def field_cross_entropy(model, field, kp_norm, weight, *, sigma: float = 2.0):
    """Cross-entropy between a probability field and a normalized Gaussian target.

    The cells are placed by the geometry the readout inverts; raw logits go through a
    spatial log-softmax.
    """
    n, k, h, w = field.shape
    geom = _field_geometry(model, field)
    with torch.no_grad():
        cells = geom.to_field_cells(kp_norm)
        target = heatmap.render_gaussian_targets(
            cells, torch.ones_like(weight), hw=(h, w), sigma=sigma
        )
        target = target / target.sum(dim=(-1, -2), keepdim=True).clamp_min(1e-12)
    if getattr(model, "field_is_prob", False):
        logp = torch.log(field.float().clamp_min(1e-12))
    else:
        logp = torch.log_softmax(field.float().flatten(2), dim=-1).view(n, k, h, w)
    ce = -(target.to(logp.dtype) * logp).sum(dim=(-1, -2))  # (N, K)
    w_sum = weight.sum()
    loss = (ce * weight.to(ce.dtype)).sum() / w_sum.clamp_min(1.0)
    return loss, {"field_ce": loss.detach()}

fit

fit(
    model,
    batches,
    *,
    steps: int,
    base_lr: float = 0.0002,
    weight_decay: float = 0.05,
    betas: tuple[float, float] = (0.9, 0.95),
    warmup_steps: int = 0,
    backbone_lr_mult: float = 1.0,
    sigma: float = 2.0,
    fg_weight: float = 30.0,
    field_weight: float = 0.0,
    center_weight: float = 0.25,
    coord_weight: float = 0.0,
    coord_radius: int = 3,
    coord_huber: float = 4.0,
    tri_weight: float = 0.0,
    tri_huber: float = 0.0,
    tri_max_world: float = 0.0,
    train_decode: str = "local",
    unlabeled_batches=None,
    unlabeled_weight: float = 0.0,
    mask_ratios: tuple[float, ...] = (),
    masked_weight: float = 1.0,
    unmasked_weight: float = 0.5,
    grad_clip: float = 0.0,
    amp: bool = False,
    ema_decay: float = 0.0,
    eval_every: int = 0,
    eval_fn: Callable[[Module], dict] | None = None,
    select_best: str | None = None,
    checkpoint_fn: Callable[[int, Module], None]
    | None = None,
    log_every: int = 0,
    log: Callable[[str], None] = print,
) -> dict

Train model on batches for steps AdamW steps.

batches and unlabeled_batches are cycled; each step adds unlabeled_weight times unlabeled_loss. The LR is a linear warmup then a cosine decay. amp runs the forward in bf16 on CUDA, the geometric loss in float32. With ema_decay > 0 the EMA (ramped from 0, as timm does) is what eval_fn sees and what model holds at the end, the raw weights under history["raw_state_dict"]. eval_fn and checkpoint_fn run every eval_every steps and at the end.

Returns:

Type	Description
dict	{"loss": [...], "aux": [...], "eval": [(step, report)]}.

Source code in src/deeperfly/pose2d/training/loop.py

def fit(
    model,
    batches,
    *,
    steps: int,
    base_lr: float = 2e-4,
    weight_decay: float = 0.05,
    betas: tuple[float, float] = (0.9, 0.95),
    warmup_steps: int = 0,
    backbone_lr_mult: float = 1.0,
    sigma: float = 2.0,
    fg_weight: float = 30.0,
    field_weight: float = 0.0,
    center_weight: float = 0.25,
    coord_weight: float = 0.0,
    coord_radius: int = 3,
    coord_huber: float = 4.0,
    tri_weight: float = 0.0,
    tri_huber: float = 0.0,
    tri_max_world: float = 0.0,
    train_decode: str = "local",
    unlabeled_batches=None,
    unlabeled_weight: float = 0.0,
    mask_ratios: tuple[float, ...] = (),
    masked_weight: float = 1.0,
    unmasked_weight: float = 0.5,
    grad_clip: float = 0.0,
    amp: bool = False,
    ema_decay: float = 0.0,
    eval_every: int = 0,
    eval_fn: Callable[[torch.nn.Module], dict] | None = None,
    select_best: str | None = None,
    checkpoint_fn: Callable[[int, torch.nn.Module], None] | None = None,
    log_every: int = 0,
    log: Callable[[str], None] = print,
) -> dict:
    """Train `model` on `batches` for `steps` AdamW steps.

    `batches` and `unlabeled_batches` are cycled; each step adds `unlabeled_weight`
    times `unlabeled_loss`. The LR is a linear warmup then a cosine decay. `amp` runs
    the forward in bf16 on CUDA, the geometric loss in float32. With `ema_decay > 0` the
    EMA (ramped from 0, as timm does) is what `eval_fn` sees and what `model` holds at
    the end, the raw weights under `history["raw_state_dict"]`. `eval_fn` and
    `checkpoint_fn` run every `eval_every` steps and at the end.

    Returns:
        `{"loss": [...], "aux": [...], "eval": [(step, report)]}`.
    """
    device = next(model.parameters()).device
    model.train()
    groups, base_lrs = build_param_groups(
        model, base_lr, backbone_lr_mult=backbone_lr_mult
    )
    opt = torch.optim.AdamW(groups, betas=betas, weight_decay=weight_decay)
    ema = None
    if ema_decay > 0:
        ema = copy.deepcopy(model).eval()
        for p in ema.parameters():
            p.requires_grad_(False)
    use_amp = bool(amp) and device.type == "cuda"
    stream = _cycle(batches)
    ustream = (
        _cycle(unlabeled_batches)
        if unlabeled_batches is not None and unlabeled_weight > 0
        else None
    )
    history: dict = {"loss": [], "aux": [], "eval": []}
    best: tuple[int, float, dict] | None = None
    t0 = time.time()
    for step in range(steps):
        batch = to_device(next(stream), device)
        apply_lr(opt, base_lrs, cosine_warmup_factor(step, steps, warmup_steps))
        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=use_amp):
            loss, aux = batch_loss(
                model,
                batch,
                sigma=sigma,
                fg_weight=fg_weight,
                field_weight=field_weight,
                center_weight=center_weight,
                coord_weight=coord_weight,
                coord_radius=coord_radius,
                coord_huber=coord_huber,
                tri_weight=tri_weight,
                tri_huber=tri_huber,
                tri_max_world=tri_max_world,
                train_decode=train_decode,
            )
            if ustream is not None:
                ul, uaux = unlabeled_loss(
                    model,
                    to_device(next(ustream), device),
                    sigma=sigma,
                    mask_ratios=mask_ratios,
                    masked_weight=masked_weight,
                    unmasked_weight=unmasked_weight,
                )
                loss = loss + unlabeled_weight * ul
                aux = {**aux, **uaux}
        if not torch.isfinite(loss):
            raise FloatingPointError(
                f"non-finite loss at step {step + 1}: "
                + ", ".join(f"{k}={float(v):.4g}" for k, v in aux.items())
            )
        opt.zero_grad(set_to_none=True)
        loss.backward()
        if grad_clip > 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
        opt.step()
        if ema is not None:
            _ema_update(ema, model, min(ema_decay, (1 + step) / (10 + step)))
        history["loss"].append(float(loss.detach()))
        history["aux"].append({k: float(v) for k, v in aux.items()})
        if log_every and (step + 1) % log_every == 0:
            recent = history["loss"][-log_every:]
            log(
                f"step {step + 1}/{steps} loss {sum(recent) / len(recent):.4g} "
                f"lr {opt.param_groups[0]['lr']:.2e} {time.time() - t0:.0f}s"
            )
        last = step + 1 == steps
        if (eval_every and (step + 1) % eval_every == 0) or last:
            current = ema if ema is not None else model
            if eval_fn is not None:
                report = eval_fn(current)
                history["eval"].append((step + 1, report))
                if select_best is not None:
                    score, larger = metric_at(report, select_best)
                    beats = best is None or (
                        score > best[1] if larger else score < best[1]
                    )
                    if beats:
                        best = (step + 1, score, _cpu_state(current))
                        log(f"best {select_best} {score:.4g} at step {step + 1}")
                model.train()
            if checkpoint_fn is not None and not last:
                checkpoint_fn(step + 1, current)
            if device.type == "cuda":
                # The evaluation's allocations fragment the caching allocator; a late
                # training step then failed to find a contiguous block once.
                torch.cuda.empty_cache()
    if ema is not None:
        history["raw_state_dict"] = {
            k: v.detach().cpu() for k, v in model.state_dict().items()
        }
        model.load_state_dict(ema.state_dict())
    if best is not None:
        history["best_step"], history["best_score"], history["best_state_dict"] = best
    return history

local_expectation

local_expectation(field, model, *, radius: int = 3)

Differentiable sub-cell peak: the field's expectation around its argmax.

(N, K, h, w) -> (N, K, 2) (x, y) in cells. The detached argmax only places the 2 * radius + 1 window, so mass on another limb cannot pull the expectation.

Source code in src/deeperfly/pose2d/training/loop.py

def local_expectation(field, model, *, radius: int = 3):
    """Differentiable sub-cell peak: the field's expectation around its argmax.

    `(N, K, h, w)` -> `(N, K, 2)` `(x, y)` in cells. The detached argmax only places the
    `2 * radius + 1` window, so mass on another limb cannot pull the expectation.
    """
    prob = _field_prob(model, field)
    h, w = prob.shape[-2:]
    with torch.no_grad():
        idx = prob.flatten(2).argmax(-1)
        py, px = idx // w, idx % w
    ys = torch.arange(h, device=prob.device).view(1, 1, h, 1)
    xs = torch.arange(w, device=prob.device).view(1, 1, 1, w)
    win = ((ys - py[..., None, None]).abs() <= radius) & (
        (xs - px[..., None, None]).abs() <= radius
    )
    p = prob * win
    p = p / p.sum(dim=(-1, -2), keepdim=True).clamp_min(1e-12)
    ex = (p * xs).sum(dim=(-1, -2))
    ey = (p * ys).sum(dim=(-1, -2))
    return torch.stack([ex, ey], dim=-1)

supervised_loss

supervised_loss(
    model,
    out,
    kp_norm,
    weight,
    *,
    sigma: float = 2.0,
    fg_weight: float = 30.0,
    simcc_sigma: float = 6.0,
    field_weight: float = 0.0,
    center_weight: float = 0.25,
    coord_weight: float = 0.0,
    coord_radius: int = 3,
    coord_huber: float = 4.0,
    centers=None,
    center_valid=None,
)

The per-head supervised loss of one flat forward output, by output_kind.

kp_norm is (N, K, 2) input-normalized ground truth (NaN where unlabeled) and weight (N, K). For a probability field field_weight weighs the Gaussian-target cross-entropy and coord_weight a Huber loss (coord_huber) on the local expectation of coord_radius cells (ADR-0072); a center channel adds center_weight times center_cross_entropy on centers and center_valid (ADR-0063). A locator for several animals is that cross-entropy per channel alone (ADR-0127). Returns (loss, aux).

Source code in src/deeperfly/pose2d/training/loop.py

def supervised_loss(
    model,
    out,
    kp_norm,
    weight,
    *,
    sigma: float = 2.0,
    fg_weight: float = 30.0,
    simcc_sigma: float = 6.0,
    field_weight: float = 0.0,
    center_weight: float = 0.25,
    coord_weight: float = 0.0,
    coord_radius: int = 3,
    coord_huber: float = 4.0,
    centers=None,
    center_valid=None,
):
    """The per-head supervised loss of one flat forward output, by `output_kind`.

    `kp_norm` is `(N, K, 2)` input-normalized ground truth (NaN where unlabeled) and
    `weight` `(N, K)`. For a probability field `field_weight` weighs the Gaussian-target
    cross-entropy and `coord_weight` a Huber loss (`coord_huber`) on the local
    expectation of `coord_radius` cells (ADR-0072); a center channel adds
    `center_weight` times `center_cross_entropy` on `centers` and `center_valid`
    (ADR-0063). A locator for several animals is that cross-entropy per channel alone
    (ADR-0127). Returns `(loss, aux)`.
    """
    kind = model.output_kind
    center_field = None
    if getattr(model, "center_channel", False):
        whole = out[-1] if isinstance(out, list) else out
        center_field = whole[:, model.n_points]
    out = joint_channels(model, out)
    located = int(getattr(model, "animals", 1))
    if getattr(model, "task", "pose") == "locate" and located > 1:
        # A locator for several animals: each channel is the center channel's mixture,
        # a peak per animal, with the item's `centers` laid out channel-major
        # (ADR-0127). It replaces the joint terms, which hold one point a channel.
        field = out[-1] if isinstance(out, list) else out
        n, c, h, w = field.shape
        loss, aux = center_cross_entropy(
            model,
            field.reshape(n * c, h, w),
            centers.reshape(n * c, located, 2),
            center_valid.reshape(n * c, located),
            sigma=sigma,
        )
        return loss, {"locate_ce": aux["center_ce"]}
    if kind == "heatmap":
        field = out[-1] if isinstance(out, list) else out
        cells = model.geom.to_field_cells(kp_norm)
        loss, aux = heatmap.masked_heatmap_loss(
            field, cells, weight, sigma=sigma, fg_weight=fg_weight
        )
    elif kind == "simcc":
        x_logits, y_logits = out
        tx, ty = simcc.simcc_targets(
            torch.nan_to_num(kp_norm, nan=0.0),
            input_hw=model.input_hw,
            split_ratio=model.split_ratio,
            sigma=simcc_sigma,
            margin=getattr(model, "hm_margin_px", 0),
        )
        loss, aux = simcc.simcc_kl_loss(x_logits, y_logits, tx, ty, weight)
    elif kind == "heatmap_softargmax":
        field = out[-1] if isinstance(out, list) else out
        loss, aux = field.sum() * 0.0, {}
        if coord_weight > 0:
            geom = _field_geometry(model, field)
            pred_cells = local_expectation(field, model, radius=coord_radius)
            gt_cells = geom.to_field_cells(kp_norm)
            cl, caux = coordinate_loss(
                pred_cells * geom.stride,
                gt_cells * geom.stride,
                weight,
                huber=coord_huber,
            )
            loss = coord_weight * cl
            aux = {"coord_px": caux["coord_loss"]}
        if field_weight > 0:
            ce, ce_aux = field_cross_entropy(model, field, kp_norm, weight, sigma=sigma)
            loss = loss + field_weight * ce
            aux = {**aux, **ce_aux}
    else:
        raise NotImplementedError(f"no supervised loss for output_kind {kind!r}")
    if center_field is not None and centers is not None and center_weight > 0:
        ce, ce_aux = center_cross_entropy(
            model, center_field, centers, center_valid, sigma=sigma
        )
        loss = loss + center_weight * ce
        aux = {**aux, **ce_aux}
    return loss, aux

training_decode

training_decode(
    model, out, *, kind: str = "local", radius: int = 3
)

The differentiable coordinate readout the geometric loss trains through.

(N, K, 2) input-normalized: "local" is local_expectation through the field's geometry, "softargmax" the first chain's upsampled readout, kept reproducible (ADR-0072); SimCC decodes through its own soft-argmax.

Raises:

Type	Description
ValueError	For a hard-argmax head, which has no gradient, or an unknown kind.

Source code in src/deeperfly/pose2d/training/loop.py

def training_decode(model, out, *, kind: str = "local", radius: int = 3):
    """The differentiable coordinate readout the geometric loss trains through.

    `(N, K, 2)` input-normalized: `"local"` is `local_expectation` through the field's
    geometry, `"softargmax"` the first chain's upsampled readout, kept reproducible
    (ADR-0072); SimCC decodes through its own soft-argmax.

    Raises:
        ValueError: For a hard-argmax head, which has no gradient, or an unknown `kind`.
    """
    out_kind = model.output_kind
    out = joint_channels(model, out)
    if out_kind == "heatmap":
        raise ValueError(
            "the triangulation loss needs a differentiable decode; use the "
            "softargmax readout or the simcc head"
        )
    if out_kind == "simcc":
        from ..infer import decode_points

        return decode_points(model, out)[0]
    field = out[-1] if isinstance(out, list) else out
    if kind == "local":
        geom = _field_geometry(model, field)
        return geom.to_input_normalized(local_expectation(field, model, radius=radius))
    if kind == "softargmax":
        return heatmap.soft_argmax(
            field,
            input_hw=model.input_hw,
            downsample_factor=model.downsample_factor,
            margin=model.hm_margin_px,
        )[0]
    raise ValueError(f"unknown train_decode {kind!r}; use 'local' or 'softargmax'")

unlabeled_loss

unlabeled_loss(
    model,
    batch: dict,
    *,
    sigma: float = 2.0,
    mask_ratios: tuple[float, ...] = (),
    masked_weight: float = 1.0,
    unmasked_weight: float = 0.5,
)

The distillation term on one grouped batch of teacher targets: (loss, aux).

Each view draws a mask ratio from mask_ratios, one per group left whole; the field cross-entropy weighs masked_weight on masked views and unmasked_weight on whole ones (ADR-0071). Probability-field heads only.

Source code in src/deeperfly/pose2d/training/loop.py

def unlabeled_loss(
    model,
    batch: dict,
    *,
    sigma: float = 2.0,
    mask_ratios: tuple[float, ...] = (),
    masked_weight: float = 1.0,
    unmasked_weight: float = 0.5,
):
    """The distillation term on one grouped batch of teacher targets: `(loss, aux)`.

    Each view draws a mask ratio from `mask_ratios`, one per group left whole; the field
    cross-entropy weighs `masked_weight` on masked views and `unmasked_weight` on whole
    ones (ADR-0071). Probability-field heads only.
    """
    if model.output_kind != "heatmap_softargmax":
        raise ValueError("the unlabeled term needs a probability-field head (mvt)")
    b, v = batch["image"].shape[:2]
    dev = batch["image"].device
    ratios = torch.zeros(b, v, device=dev)
    if mask_ratios:
        choices = torch.as_tensor(list(mask_ratios), device=dev, dtype=torch.float32)
        ratios = choices[torch.randint(len(choices), (b, v), device=dev)]
        ratios[torch.arange(b, device=dev), torch.randint(v, (b,), device=dev)] = 0.0
    flat_ratios = ratios.reshape(-1)
    out, flat, _ = forward_flat(model, batch, mask_ratio=flat_ratios)
    field = joint_channels(model, out[-1] if isinstance(out, list) else out)
    per_view = torch.where(
        flat_ratios > 0,
        torch.full_like(flat_ratios, float(masked_weight)),
        torch.full_like(flat_ratios, float(unmasked_weight)),
    )
    weight = flat["weight"] * per_view[:, None]
    ce, aux = field_cross_entropy(model, field, flat["kp"], weight, sigma=sigma)
    return ce, {
        "unl_ce": aux["field_ce"],
        "unl_masked": (flat_ratios > 0).float().mean().detach(),
    }

Keypoint accuracy metrics over decoded coordinates, and the relations they split on.

The caller passes per-point errors in any unit; breakdown reduces them under any grouping, since a pooled number hides what matters (two claw channels carried 44% of the squared error on one animal), and relation names the grouping that tracks the hard cells: ipsilateral, contralateral, midline, or an axial view by name.

breakdown

breakdown(
    errors, keys: Sequence[Hashable], *, pck_thresh: float
) -> dict

error_summary per distinct key (zipped 1:1 with errors), in key order.

Source code in src/deeperfly/pose2d/training/metrics.py

def breakdown(errors, keys: Sequence[Hashable], *, pck_thresh: float) -> dict:
    """`error_summary` per distinct key (zipped 1:1 with `errors`), in key order."""
    e = np.asarray(errors, dtype=np.float64).ravel()
    keys = np.asarray(list(keys), dtype=object)
    if keys.shape != e.shape:
        raise ValueError(f"keys {keys.shape} must be 1:1 with errors {e.shape}")
    out = {}
    for k in sorted(set(keys.tolist()), key=str):
        out[k] = error_summary(e[keys == k], pck_thresh=pck_thresh)
    return out

error_summary

error_summary(errors, *, pck_thresh: float) -> dict

Per-point errors -> {median, mean, p90, pck, n}; an empty input gives zeros.

Source code in src/deeperfly/pose2d/training/metrics.py

def error_summary(errors, *, pck_thresh: float) -> dict:
    """Per-point errors -> `{median, mean, p90, pck, n}`; an empty input gives zeros."""
    e = np.asarray(errors, dtype=np.float64).ravel()
    n = int(e.size)
    if n == 0:
        return {"median": 0.0, "mean": 0.0, "p90": 0.0, "pck": 0.0, "n": 0}
    return {
        "median": float(np.median(e)),
        "mean": float(e.mean()),
        "p90": float(np.percentile(e, 90)),
        "pck": float((e < pck_thresh).mean()),
        "n": n,
    }

keypoint_side

keypoint_side(skeleton, keypoint: int) -> str

"left" or "right" for a paired point named l* or r*, else "midline".

Source code in src/deeperfly/pose2d/training/metrics.py

def keypoint_side(skeleton, keypoint: int) -> str:
    """`"left"` or `"right"` for a paired point named `l*` or `r*`, else `"midline"`."""
    if skeleton.partner(keypoint) is None:
        return "midline"
    return {"l": "left", "r": "right"}.get(
        skeleton.point_names[keypoint][:1].lower(), "midline"
    )

metric_at

metric_at(report: dict, path: str) -> tuple[float, bool]

One number out of an evaluation report, and whether larger is better.

path is dotted ("overall.pck", "per_view.east.median") and must end in one of error_summary's metric names, which fix the direction.

Raises:

Type	Description
KeyError	If the path names nothing in this report, or ends in no metric name.

Source code in src/deeperfly/pose2d/training/metrics.py

def metric_at(report: dict, path: str) -> tuple[float, bool]:
    """One number out of an evaluation report, and whether larger is better.

    `path` is dotted (`"overall.pck"`, `"per_view.east.median"`) and must end in one of
    `error_summary`'s metric names, which fix the direction.

    Raises:
        KeyError: If the path names nothing in this report, or ends in no metric name.
    """
    metric = path.rsplit(".", 1)[-1]
    if metric not in LARGER_IS_BETTER:
        raise KeyError(f"{path!r}: {metric!r} is not one of {sorted(LARGER_IS_BETTER)}")
    node: object = report
    for key in path.split("."):
        if not isinstance(node, dict) or key not in node:
            raise KeyError(f"{path!r}: no {key!r} in this report")
        node = node[key]
    return float(node), LARGER_IS_BETTER[metric]  # type: ignore[arg-type]

relation

relation(
    view_name: str,
    keypoint: int,
    skeleton,
    *,
    view_sides: Callable[[str], str | None]
    | dict[str, str | None]
    | None = None,
) -> str

The relation of a keypoint to the view that sees it.

"ipsilateral", "contralateral" or "midline" for a side view, the view's own name for an axial one; view_sides overrides view_side's naming convention.

Source code in src/deeperfly/pose2d/training/metrics.py

def relation(
    view_name: str,
    keypoint: int,
    skeleton,
    *,
    view_sides: Callable[[str], str | None] | dict[str, str | None] | None = None,
) -> str:
    """The relation of a keypoint to the view that sees it.

    `"ipsilateral"`, `"contralateral"` or `"midline"` for a side view, the view's own
    name for an axial one; `view_sides` overrides `view_side`'s naming convention.
    """
    if view_sides is None:
        side = view_side(view_name)
    elif callable(view_sides):
        side = view_sides(view_name)
    else:
        side = view_sides.get(view_name)
    if side is None:
        return view_name
    kp_side = keypoint_side(skeleton, keypoint)
    if kp_side == "midline":
        return "midline"
    return "ipsilateral" if kp_side == side else "contralateral"

view_side

view_side(name: str) -> str | None

A view's side from its name's leading l or r; None for an axial view.

Source code in src/deeperfly/pose2d/training/metrics.py

def view_side(name: str) -> str | None:
    """A view's side from its name's leading `l` or `r`; `None` for an axial view."""
    first = name[:1].lower()
    return {"l": "left", "r": "right"}.get(first)

Evaluate a detector over a batch stream: decode, error per cell, break it down.

The inference decode, each sample's affine inverted so the error is in the video pixels the label was clicked in, reduced per keypoint, view, relation, project and in-frame against off-frame. Raw 2D detections are scored, never a reprojection.

evaluate

evaluate(
    model,
    batches,
    *,
    dataset,
    pck_px: float = 10.0,
    view_sides=None,
    return_cells: bool = False,
    flip_tta: bool = False,
) -> dict

Score model on unaugmented batches from dataset; errors in video pixels.

return_cells adds the per-cell arrays under "cells", for a caller's own statistics (bootstrap CIs over frames or animals); flip_tta averages in the mirrored pass. Returns overall, in_frame, off_frame and model_px summaries and per_keypoint, per_view, per_relation, per_project and off_frame_per_view breakdowns.

Source code in src/deeperfly/pose2d/training/eval.py

@torch.no_grad()
def evaluate(
    model,
    batches,
    *,
    dataset,
    pck_px: float = 10.0,
    view_sides=None,
    return_cells: bool = False,
    flip_tta: bool = False,
) -> dict:
    """Score `model` on unaugmented `batches` from `dataset`; errors in video pixels.

    `return_cells` adds the per-cell arrays under `"cells"`, for a caller's own
    statistics (bootstrap CIs over frames or animals); `flip_tta` averages in the
    mirrored pass. Returns `overall`, `in_frame`, `off_frame` and `model_px` summaries
    and `per_keypoint`, `per_view`, `per_relation`, `per_project` and
    `off_frame_per_view` breakdowns.
    """
    from ..infer import decode_points, joint_channels

    model.eval()
    device = next(model.parameters()).device
    h, w = dataset.input_hw
    cols = {
        k: []
        for k in ("err", "err_model", "kp", "view", "project", "frame", "in_frame")
    }
    perm = torch.as_tensor(dataset.skeleton.flip_perm(), device=device)
    for batch in batches:
        batch = to_device(batch, device)
        out, flat, _ = forward_flat(model, batch)
        pred_norm, _ = decode_points(model, joint_channels(model, out))
        pred_norm = pred_norm.float()
        if flip_tta:
            pred_norm = flip_average(model, batch, pred_norm, perm)
        shift = grid_shift(model)
        if shift:
            # A detector trained on corner-origin labels, moved onto the pixel-center
            # grid through each sample's window (ADR-0078): a video-pixel step `d`
            # is `A d` in the input, so the model-space error sees it too.
            lin = flat["affine"].float()[:, :, :2]
            step = lin @ lin.new_full((2,), shift)
            pred_norm = pred_norm + (step / pred_norm.new_tensor([w, h]))[:, None, :]
        weight = flat["weight"].float()
        valid = flat["view_valid"].bool()
        pred_video_px = to_video_px(pred_norm, flat["affine"], (h, w))
        err = (pred_video_px - flat["gt_video_px"].float()).norm(dim=-1)  # (N, K)
        err_model = (
            (pred_norm - flat["kp"].float()) * pred_norm.new_tensor([w, h])
        ).norm(dim=-1)
        mask = (weight > 0) & valid[:, None]
        n, k = mask.shape
        kp_idx = torch.arange(k, device=device).expand(n, k)
        view_idx = flat["view"].view(n, 1).expand(n, k)
        project_idx = (
            torch.as_tensor(flat["project"], device=device).view(n, 1).expand(n, k)
        )
        frame = flat["frame"] if "frame" in flat else torch.full((n,), -1)
        frame_idx = torch.as_tensor(frame, device=device).view(n, 1).expand(n, k)
        cols["err"].append(err[mask].cpu().numpy())
        cols["err_model"].append(err_model[mask].cpu().numpy())
        cols["kp"].append(kp_idx[mask].cpu().numpy())
        cols["view"].append(view_idx[mask].cpu().numpy())
        cols["project"].append(project_idx[mask].cpu().numpy())
        cols["frame"].append(frame_idx[mask].cpu().numpy())
        cols["in_frame"].append(flat["in_frame"][mask].cpu().numpy())
    cells = {
        key: (np.concatenate(v) if v else np.empty(0, dtype=np.float64))
        for key, v in cols.items()
    }
    cells["in_frame"] = cells["in_frame"].astype(bool)
    cells["kp"] = cells["kp"].astype(int)
    cells["view"] = cells["view"].astype(int)
    cells["project"] = cells["project"].astype(int)
    cells["frame"] = cells["frame"].astype(int)

    err = cells["err"]
    view_names = [
        dataset.view_name(int(r), int(v))
        for r, v in zip(cells["project"], cells["view"], strict=True)
    ]
    project_names = [dataset.project_name(int(r)) for r in cells["project"]]
    kp_names = [dataset.keypoint_names[int(i)] for i in cells["kp"]]
    relations = [
        relation(vn, int(i), dataset.skeleton, view_sides=view_sides)
        for vn, i in zip(view_names, cells["kp"], strict=True)
    ]
    off = ~cells["in_frame"]
    report = {
        "units": "video_px",
        "pck_px": float(pck_px),
        "overall": error_summary(err, pck_thresh=pck_px),
        "in_frame": error_summary(err[~off], pck_thresh=pck_px),
        "off_frame": error_summary(err[off], pck_thresh=pck_px),
        "model_px": error_summary(cells["err_model"], pck_thresh=pck_px),
        "per_keypoint": breakdown(err, kp_names, pck_thresh=pck_px),
        "per_view": breakdown(err, view_names, pck_thresh=pck_px),
        "per_relation": breakdown(err, relations, pck_thresh=pck_px),
        "per_project": breakdown(err, project_names, pck_thresh=pck_px),
        "off_frame_per_view": breakdown(
            err[off],
            [vn for vn, o in zip(view_names, off, strict=True) if o],
            pck_thresh=pck_px,
        ),
    }
    if return_cells:
        cells["view_name"] = np.asarray(view_names, dtype=object)
        cells["project"] = np.asarray(project_names, dtype=object)
        cells["relation"] = np.asarray(relations, dtype=object)
        report["cells"] = cells
    return report

flip_average

flip_average(model, batch: dict, pred_norm, flip_perm)

Average pred_norm with the unmirrored prediction on the mirrored batch.

On coordinates, not fields, since the cell grid is not symmetric about the frame center (ADR-0072); channels permuted by flip_perm.

Source code in src/deeperfly/pose2d/training/eval.py

def flip_average(model, batch: dict, pred_norm, flip_perm):
    """Average `pred_norm` with the unmirrored prediction on the mirrored batch.

    On coordinates, not fields, since the cell grid is not symmetric about the frame
    center (ADR-0072); channels permuted by `flip_perm`.
    """
    from ..infer import decode_points, joint_channels

    mirrored = dict(batch)
    mirrored["image"] = batch["image"].flip(-1)
    out_m, _, _ = forward_flat(model, mirrored)
    pm, _ = decode_points(model, joint_channels(model, out_m))
    pm = pm.float().index_select(1, flip_perm.to(pm.device)).clone()
    w = model.input_hw[1]
    pm[..., 0] = (w - 1 - pm[..., 0] * w) / w
    return 0.5 * (pred_norm + pm)

A training from its config: data, model, loop, held-out evaluation, artifacts.

What deeperfly train drives. A training over projects (ADR-0085) takes them from [data] projects or from the paths the command is given, names them as a collection does, and requires one skeleton of them all. It writes into its output directory the config it ran (with the projects it was given, when they replaced the document's), the per-step history, the held-out report (per keypoint, view, relation and project, in video pixels, beside each project's folder) and the weights: a deeperfly-dense38-1 checkpoint for the monocular family, a deeperfly-detector-1 artifact for the others, which models loads for inference.

boost_repeats

boost_repeats(
    owners: Sequence[str],
    projects: Sequence[str],
    share: float,
) -> tuple[list[float], str]

How often an epoch holds each group under [data.boost], and a log line.

owners names each group's project. The groups of the projects projects selects are repeated by one factor, the one that makes them share of the epoch, and never less than once (a selection already past its share is left alone); every other group is held once.

Raises:

Type	Description
ConfigError	As match_names, when a pattern selects no project with a labeled moment in the training set.

Source code in src/deeperfly/pose2d/training/run.py

def boost_repeats(
    owners: Sequence[str], projects: Sequence[str], share: float
) -> tuple[list[float], str]:
    """How often an epoch holds each group under `[data.boost]`, and a log line.

    `owners` names each group's project. The groups of the projects `projects` selects
    are repeated by one factor, the one that makes them `share` of the epoch, and
    never less than once (a selection already past its share is left alone); every
    other group is held once.

    Raises:
        ConfigError: As [match_names][deeperfly.document.match_names], when a pattern
            selects no project with a labeled moment in the training set.
    """
    boosted = set(match_names(projects, sorted(set(owners)), "[data.boost] projects"))
    n_boost = sum(owner in boosted for owner in owners)
    n_rest = len(owners) - n_boost
    factor = max(1.0, share / (1.0 - share) * n_rest / n_boost)
    epoch = factor * n_boost + n_rest
    line = (
        f"boost: {n_boost} labeled animals of {len(boosted)} projects drawn "
        f"x{factor:.1f}, {factor * n_boost / epoch:.0%} of an epoch of {epoch:.0f}"
    )
    return [factor if owner in boosted else 1.0 for owner in owners], line

build_model

build_model(cfg: TrainModel, n_keypoints: int)

Assemble the detector [model] names, its family keys checked at load.

Source code in src/deeperfly/pose2d/training/run.py

def build_model(cfg: TrainModel, n_keypoints: int):
    """Assemble the detector `[model]` names, its family keys checked at load."""
    from ..detect import DETECTORS

    kw = {"input_hw": cfg.input_hw, "pretrained": cfg.pretrained, **cfg.extra}
    if cfg.backbone is not None:
        kw["backbone"] = cfg.backbone
    return DETECTORS[cfg.family](n_keypoints, **kw)

dataset_margin

dataset_margin(model) -> tuple[int, int]

The (y, x) off-frame margin in input pixels the model's field actually has.

Source code in src/deeperfly/pose2d/training/run.py

def dataset_margin(model) -> tuple[int, int]:
    """The `(y, x)` off-frame margin in input pixels the model's field actually has."""
    m = getattr(model, "hm_margin_px", 0)
    return (int(m), int(m)) if np.isscalar(m) else (int(m[0]), int(m[1]))

load_artifact

load_artifact(path: str | Path, **kwargs)

Rebuild a deeperfly-detector-1 artifact's module on the CPU, in eval mode.

The artifact's [model] table goes through the same assembler (never with pretrained) and the state dict loads strictly; kwargs override its keys.

Source code in src/deeperfly/pose2d/training/run.py

def load_artifact(path: str | Path, **kwargs):
    """Rebuild a `deeperfly-detector-1` artifact's module on the CPU, in eval mode.

    The artifact's `[model]` table goes through the same assembler (never with
    `pretrained`) and the state dict loads strictly; `kwargs` override its keys.
    """
    path = artifact_file(path)
    blob = torch.load(path, map_location="cpu", weights_only=False)
    if not isinstance(blob, dict) or blob.get("format") != ARTIFACT_FORMAT:
        raise SystemExit(
            f"{path} is not a {ARTIFACT_FORMAT!r} artifact; train one with "
            '`deeperfly train` (family = "mvt").'
        )
    cfg = parse(TrainModel, {**blob["model"], "pretrained": False, **kwargs}, "model")
    module = build_model(cfg, int(blob["n_keypoints"]))
    module.load_state_dict(blob["state_dict"], strict=True)
    # `n_keypoints` is the channel count; an artifact written before the center channel
    # records no flag and every channel is a joint.
    module.center_channel = bool(blob.get("center_channel", False))
    # What the run trained it to do: "pose" (joints) or "locate" (a channel per
    # view). An artifact written before the locator existed says nothing and poses.
    module.task = str(blob.get("task", "pose"))
    # Animals a locator's channel holds a peak for (ADR-0127); one when unrecorded.
    module.animals = int(blob.get("animals", 1))
    # Whether its inputs mark the other animals' centers (ADR-0148).
    module.mark_others = bool(blob.get("mark_others", False))
    module.norm_mean = float(blob.get("norm_mean", module.norm_mean))
    module.norm_std = float(blob.get("norm_std", module.norm_std))
    module.point_names = list(blob.get("point_names", []))
    module.pixel_grid = artifact_grid(blob, path)
    return module.eval()

load_artifact_state

load_artifact_state(
    model, path: str | Path, *, log=print
) -> None

Warm-start model from a run's artifact, non-strictly.

Mismatches are reported, and a start where under half of the tensors came from the file is refused: that is a different model. A head wider than the parent's takes the parent's rows (_widened); any other mismatched tensor is dropped whole.

Source code in src/deeperfly/pose2d/training/run.py

def load_artifact_state(model, path: str | Path, *, log=print) -> None:
    """Warm-start `model` from a run's artifact, non-strictly.

    Mismatches are reported, and a start where under half of the tensors came from the
    file is refused: that is a different model. A head wider than the parent's takes
    the parent's rows (`_widened`); any other mismatched tensor is dropped whole.
    """
    blob = torch.load(artifact_file(path), map_location="cpu", weights_only=False)
    state = (
        blob["state_dict"] if isinstance(blob, dict) and "state_dict" in blob else blob
    )
    own = model.state_dict()
    n_classes = int(getattr(model, "num_classes", 0))
    usable, widened = {}, []
    for k, v in state.items():
        if k not in own:
            continue
        if tuple(v.shape) == tuple(own[k].shape):
            usable[k] = v
        elif k.startswith("head."):
            wide = _widened(own[k], v, n_classes)
            if wide is not None:
                usable[k] = wide
                widened.append(k)
    missing = sorted(set(own) - set(usable))
    skipped = sorted(set(state) - set(usable))
    if len(usable) < len(own) / 2:
        raise ValueError(
            f"init_from {path}: only {len(usable)} of {len(own)} tensors match; "
            "refusing to call that a warm start"
        )
    model.load_state_dict(usable, strict=False)
    log(
        f"init_from {path}: {len(usable)} tensors loaded, {len(missing)} left at init"
        + (f" ({missing[:4]}...)" if missing else "")
        + (f", {len(skipped)} in the file unused" if skipped else "")
        + (f", {len(widened)} widened to {n_classes} channels" if widened else "")
    )

locator_animals

locator_animals(projects: Mapping[str, Project]) -> int

How many animals a locator trained on projects puts a peak for per channel.

Source code in src/deeperfly/pose2d/training/run.py

def locator_animals(projects: Mapping[str, Project]) -> int:
    """How many animals a locator trained on `projects` puts a peak for per channel."""
    return max((p.animals for p in projects.values()), default=1)

locator_views

locator_views(projects: Mapping[str, Project]) -> list[str]

Every view of the projects, by name order, first seen first: the channels.

Read off every project, labeled or not, so an artifact stays readable on a project the training never saw.

Source code in src/deeperfly/pose2d/training/run.py

def locator_views(projects: Mapping[str, Project]) -> list[str]:
    """Every view of the projects, by name order, first seen first: the channels.

    Read off every project, labeled or not, so an artifact stays readable on a project
    the training never saw.
    """
    out: list[str] = []
    for name in sorted(projects):
        for view in projects[name].views:
            if view not in out:
                out.append(view)
    return out

model_table

model_table(cfg: TrainModel) -> dict

The [model] table as the artifact records it (no warm-start reference).

Source code in src/deeperfly/pose2d/training/run.py

def model_table(cfg: TrainModel) -> dict:
    """The `[model]` table as the artifact records it (no warm-start reference)."""
    table = {
        "family": cfg.family,
        "pretrained": cfg.pretrained,
        "input_hw": list(cfg.input_hw),
        **cfg.extra,
    }
    if cfg.backbone is not None:
        table["backbone"] = cfg.backbone
    return table

out_dir

out_dir(cfg: TrainConfig, out: Path | None = None) -> Path

Where a training writes: out, else [train] out, else beside its config.

Beside the config is <config dir>/<stem>/ (ADR-0085).

Source code in src/deeperfly/pose2d/training/run.py

def out_dir(cfg: TrainConfig, out: Path | None = None) -> Path:
    """Where a training writes: `out`, else `[train] out`, else beside its config.

    Beside the config is `<config dir>/<stem>/` (ADR-0085).
    """
    if out is not None:
        return Path(out).expanduser()
    if cfg.train.out is not None:
        return cfg.train.out
    if cfg.path is None:
        raise ValueError("an in-memory config has no default output directory")
    return cfg.path.with_suffix("")

save_artifact

save_artifact(
    model, path: Path, cfg_model: TrainModel, point_names
) -> None

Write the trained weights: the shipped checkpoint format when one exists.

Source code in src/deeperfly/pose2d/training/run.py

def save_artifact(model, path: Path, cfg_model: TrainModel, point_names) -> None:
    """Write the trained weights: the shipped checkpoint format when one exists."""
    if getattr(model, "family", None) == "monocular":
        from ..detect import export_monocular

        export_monocular(model, path, point_names=point_names)
        return
    blob = {
        "format": ARTIFACT_FORMAT,
        "family": model.family,
        "model": model_table(cfg_model),
        "state_dict": model.state_dict(),
        "n_keypoints": int(model.num_classes),
        "center_channel": bool(getattr(model, "center_channel", False)),
        "task": str(getattr(model, "task", "pose")),
        "input_size": list(model.input_hw),
        "point_names": list(point_names),
        "norm_mean": float(model.norm_mean),
        "norm_std": float(model.norm_std),
        # Every label it trained on is at a pixel center (ADR-0055, ADR-0078).
        "pixel_grid": "center",
    }
    # A locator trained on several animals says so; nothing else carries the key.
    located = int(getattr(model, "animals", 1))
    if located > 1:
        blob["animals"] = located
    # So does a detector whose inputs marked the other animals (ADR-0148): inference
    # has to draw the same mark.
    if getattr(model, "mark_others", False):
        blob["mark_others"] = True
    torch.save(blob, path)

shared_skeleton

shared_skeleton(
    projects: Mapping[str, Project],
) -> Skeleton

The skeleton every project labels (ADR-0085).

Raises:

Type	Description
ConfigError	If two projects' skeletons differ, naming them.

Source code in src/deeperfly/pose2d/training/run.py

def shared_skeleton(projects: Mapping[str, Project]) -> Skeleton:
    """The skeleton every project labels (ADR-0085).

    Raises:
        ConfigError: If two projects' skeletons differ, naming them.
    """
    first: tuple[str, Skeleton] | None = None
    for name, project in projects.items():
        skeleton = project.skeleton()
        if first is None:
            first = (name, skeleton)
        elif skeleton.digest != first[1].digest:
            raise ConfigError(
                f"[data].projects: {first[0]} labels {first[1].label} and {name} "
                f"labels {skeleton.label}; a detector trains on one skeleton"
            )
    if first is None:
        raise ConfigError("[data].projects: no project to train on")
    return first[1]

train

train(
    config: TrainConfig,
    out_dir: str | Path,
    *,
    projects: Mapping[str, Project] | None = None,
    log=print,
) -> dict

Run one training and return its report (also written to out_dir).

projects (name -> Project, as training_projects gives them) replace the document's [data] projects, and the copy of the document in out_dir then lists their folders in its place.

Source code in src/deeperfly/pose2d/training/run.py

def train(
    config: TrainConfig,
    out_dir: str | Path,
    *,
    projects: Mapping[str, Project] | None = None,
    log=print,
) -> dict:
    """Run one training and return its report (also written to `out_dir`).

    `projects` (`name -> Project`, as
    [training_projects][deeperfly.pose2d.training.run.training_projects] gives them)
    replace the document's `[data] projects`, and the copy of the document in
    `out_dir` then lists their folders in its place.
    """
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    data, loop = config.data, config.train

    torch.manual_seed(loop.seed)
    np.random.seed(loop.seed)
    # Read by the CUDA caching allocator at its first allocation: expandable segments
    # avoid the fragmentation that ended a long run with a spurious out-of-memory.
    os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

    named: dict[str, Project] = {}
    if data.kind == "project":
        named = dict(projects) if projects is not None else training_projects(config)
        shared = shared_skeleton(named)
        if data.task == "locate":
            skeleton = locator_skeleton(locator_views(named))
        else:
            skeleton = shared
    elif projects is not None:
        raise ConfigError(f"[data].kind = {data.kind!r} reads a corpus, not projects")
    else:
        shared = None
        skeleton = load_skeleton(data.skeleton or "fly38")
    model = build_model(config.model, skeleton.n_points + int(data.center_channel))
    model.center_channel = bool(data.center_channel)
    model.task = data.task
    model.mark_others = data.mark_others and data.task == "pose"
    if data.task == "locate":
        model.animals = locator_animals(named)
    if data.task == "locate" and not getattr(model, "has_heatmaps", False):
        raise ConfigError(
            '[data].task = "locate": a locator channel is a peak in a 2D field, so '
            f"the {config.model.family!r} head cannot carry one; use a heatmap family"
        )
    if model.center_channel and model.output_kind == "simcc":
        raise ConfigError(
            "[data].center_channel: the center is a peak in a 2D field, and the SimCC "
            "head emits none; train the locator with a heatmap family"
        )
    if config.model.init_from is not None:
        load_artifact_state(model, config.model.init_from, log=log)
    device = config.compute.device or ("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    model.grad_checkpoint = bool(loop.grad_checkpoint)
    joint = bool(getattr(model, "joint_views", False))
    grouped = joint if loop.grouped is None else loop.grouped
    if joint and not grouped:
        raise ConfigError(
            "[train].grouped: the multiview transformer trains on grouped batches only"
        )

    skeleton, train_ds, holdout_ds, holdout, holdout_frames = _build_datasets(
        data,
        named,
        shared,
        input_hw=config.model.input_hw,
        margin=dataset_margin(model),
        augment=_augment(config.augment),
        grouped=grouped,
    )
    if skeleton.n_points != model.n_points:
        raise ValueError(
            f"the skeleton has {skeleton.n_points} points but the model emits "
            f"{model.n_points} joint channels"
        )
    n_train = len({(r, t) for r, t, _, _ in train_ds.items})
    log(
        f"train: {len(train_ds.entries)} projects, {n_train} frames, "
        f"{len(train_ds)} views; holdout: {holdout}"
        + (f", frames {holdout_frames}" if holdout_frames else "")
        + (f" ({len(holdout_ds)} views)" if holdout_ds is not None else "")
    )

    workers = config.compute.workers
    bs = loop.batch_size
    ipb = loop.images_per_batch or None
    unl_loader = None
    if config.unlabeled is not None:
        unl = config.unlabeled
        if data.kind != "project" or not grouped:
            raise ConfigError(
                '[unlabeled]: the distillation stream needs kind = "project" and a '
                "grouped (multiview) model"
            )
        from .unlabeled import build_unlabeled

        names = sorted(named)
        # Only a fold that holds out a whole project leaves it out of the stream; a
        # frame-level fold skips the held frames, `guard_frames` either side (ADR-0071).
        skip_frames = {} if unl.transductive else dict(holdout_frames)
        left_out = set() if unl.transductive else set(holdout)
        if unl.exclude:
            left_out |= set(match_names(unl.exclude, names, "[unlabeled].exclude"))
        kept = [
            n
            for n in names
            if n not in left_out and named[n].manifest.active_calibration is not None
        ]
        unl_ds = build_unlabeled(
            _subset(named, kept),
            skeleton=shared,
            teacher=unl.teacher,
            n_per_project=unl.n_per_project,
            agree_px=unl.agree_px,
            min_views=unl.min_views,
            conf_min=unl.conf_min,
            input_hw=config.model.input_hw,
            margin_px=dataset_margin(model),
            augment=_augment(config.augment),
            geom_weight=unl.geom_weight,
            skip_frames=skip_frames,
            guard_frames=unl.guard_frames,
            crop_tightness=unl.crop_tightness,
            cache_dir=unl.cache,
            device=device,
            log=log,
        )
        log(
            f"unlabeled: {len(unl_ds.entries)} projects, {unl_ds.n_moments()} "
            f"moments, {len(unl_ds)} views; transductive: {unl.transductive}"
        )
        unl_loader = _loader(
            GroupedLabelDataset(
                unl_ds, view_drop_p=loop.view_drop_p, view_min=loop.view_min
            ),
            grouped=True,
            batch_size=bs,
            images_per_batch=ipb,
            shuffle=True,
            workers=workers,
            seed=loop.seed + 1,
        )
    train_groups = (
        GroupedLabelDataset(
            train_ds, view_drop_p=loop.view_drop_p, view_min=loop.view_min
        )
        if grouped
        else train_ds
    )
    repeats = None
    if data.boost.projects:
        if not grouped:
            raise ConfigError(
                "[data.boost] repeats the groups a multiview model is batched by; "
                "this model is not one"
            )
        repeats, line = boost_repeats(
            train_groups.group_projects(), data.boost.projects, data.boost.share
        )
        log(line)
    train_loader = _loader(
        train_groups,
        grouped=grouped,
        batch_size=bs,
        images_per_batch=ipb,
        shuffle=True,
        workers=workers,
        seed=loop.seed,
        repeats=repeats,
    )
    eval_fn = None
    view_sides = _view_sides(data.view_sides)
    if holdout_ds is not None and len(holdout_ds):
        holdout_loader = _loader(
            GroupedLabelDataset(holdout_ds) if grouped else holdout_ds,
            grouped=grouped,
            batch_size=bs,
            images_per_batch=ipb,
            shuffle=False,
            workers=workers,
            seed=loop.seed,
        )

        def eval_fn(module):
            report = evaluate(
                module,
                holdout_loader,
                dataset=holdout_ds,
                pck_px=loop.eval.pck_px,
                view_sides=view_sides,
                flip_tta=loop.eval.tta,
            )
            o = report["overall"]
            log(
                f"holdout: mean {o['mean']:.2f} px, p90 {o['p90']:.2f} px, "
                f"pck@{loop.eval.pck_px:g} {o['pck']:.3f} over {o['n']} cells"
            )
            return report

    if loop.eval.select_best and eval_fn is None:
        raise ValueError(
            f"select_best = {loop.eval.select_best!r} needs a holdout to score, and "
            "this training config holds nothing out"
        )

    # The config is copied, never re-emitted, so its comments and layout survive;
    # projects the command gave in place of the document's are written into it.
    if config.path is not None:
        text = config.recorded_text()
        if projects is not None:
            text = _recorded_document(text, [p.folder for p in named.values()])
        (out / "config.toml").write_text(text)

    def checkpoint_fn(step, module):
        save_artifact(
            module, out / "checkpoint.pth", config.model, skeleton.point_names
        )
        log(f"checkpoint at step {step} -> {out / 'checkpoint.pth'}")

    history = fit(
        model,
        train_loader,
        steps=loop.steps,
        base_lr=loop.optimizer.lr,
        weight_decay=loop.optimizer.weight_decay,
        warmup_steps=loop.optimizer.warmup_steps,
        backbone_lr_mult=loop.optimizer.backbone_lr_mult,
        sigma=loop.loss.sigma,
        fg_weight=loop.loss.fg_weight,
        field_weight=loop.loss.field_weight,
        center_weight=loop.loss.center_weight,
        coord_weight=loop.loss.coord_weight,
        coord_radius=loop.loss.coord_radius,
        coord_huber=loop.loss.coord_huber,
        tri_weight=loop.loss.tri_weight,
        tri_huber=loop.loss.tri_huber,
        tri_max_world=loop.loss.tri_max_world,
        train_decode=loop.loss.decode,
        unlabeled_batches=unl_loader,
        unlabeled_weight=0.0 if config.unlabeled is None else config.unlabeled.weight,
        mask_ratios=() if config.unlabeled is None else config.unlabeled.mask_ratios,
        masked_weight=1.0
        if config.unlabeled is None
        else config.unlabeled.masked_weight,
        unmasked_weight=(
            0.5 if config.unlabeled is None else config.unlabeled.unmasked_weight
        ),
        grad_clip=loop.grad_clip,
        amp=loop.amp,
        ema_decay=loop.ema_decay,
        eval_every=loop.eval.every,
        eval_fn=eval_fn,
        select_best=loop.eval.select_best,
        checkpoint_fn=checkpoint_fn,
        log_every=loop.log_every,
        log=log,
    )
    raw = history.pop("raw_state_dict", None)
    if raw is not None:
        torch.save(raw, out / "model_raw.pth")
    # `select_best` makes the artifact the best evaluated weights, since a fixed step
    # budget overtrains a small label set; the final weights stay beside it.
    best = history.pop("best_state_dict", None)
    if best is not None:
        save_artifact(
            model, out / "model_final.pth", config.model, skeleton.point_names
        )
        model.load_state_dict(best)
        log(
            f"selected step {history['best_step']} by {loop.eval.select_best} "
            f"({history['best_score']:.4g})"
        )
    save_artifact(model, out / "model.pth", config.model, skeleton.point_names)
    (out / "checkpoint.pth").unlink(missing_ok=True)
    report = {
        # Each project by its name in the training and its folder.
        "projects": {name: str(p.folder) for name, p in named.items()},
        "train_projects": [r.name for r in train_ds.entries],
        "holdout": holdout,
        "holdout_frames": holdout_frames,
        "n_train_items": len(train_ds),
        "augment": config.augment.model_dump() if config.augment else None,
        "final_loss": history["loss"][-1] if history["loss"] else None,
        "holdout_report": history["eval"][-1][1] if history["eval"] else None,
        "select_best": loop.eval.select_best,
        "best_step": history.get("best_step"),
        "best_score": history.get("best_score"),
    }
    (out / "history.json").write_text(json.dumps(history, indent=1))
    (out / "report.json").write_text(json.dumps(report, indent=1))
    return report

training_projects

training_projects(
    cfg: TrainConfig,
    paths: Sequence[str | Path] | None = None,
) -> dict[str, Project]

The projects a training takes, by their names in the collection.

paths replace [data] projects and are opened as a collection is (each a project or a directory of them); the document's patterns expand against its directory, a match that is no project skipped.

Raises:

Type	Description
ConfigError	If a pattern matches nothing, a path is no project, or no project is left.

Source code in src/deeperfly/pose2d/training/run.py

def training_projects(
    cfg: TrainConfig, paths: Sequence[str | Path] | None = None
) -> dict[str, Project]:
    """The projects a training takes, by their names in the collection.

    `paths` replace `[data] projects` and are opened as a collection is (each a
    project or a directory of them); the document's patterns expand against its
    directory, a match that is no project skipped.

    Raises:
        ConfigError: If a pattern matches nothing, a path is no project, or no
            project is left.
    """
    try:
        if paths:
            projects = list(Collection.open(list(paths)).projects)
        else:
            root = cfg.root if cfg.root is not None else Path.cwd()
            projects = expand_projects(cfg.data.projects, root)
    except (FileNotFoundError, ValueError) as exc:
        raise ConfigError(str(exc) if paths else f"[data].projects: {exc}") from exc
    if not projects:
        raise ConfigError(
            "[data].projects: no project to train on; list them in the document or "
            "give their paths to `deeperfly train`"
        )
    names = project_names([p.folder for p in projects])
    return dict(zip(names, projects, strict=True))

The one loader: a detector's class, weights and input contract, on a device.

A ModelSpec is the torch-free description and load_model turns it into a LoadedModel; weights is an already resolved path. load_models and load_locator are what a prediction run calls.

LoadedModel

A loaded detector plus its input contract and fused forward and decode.

Wraps the torch module behind the torch-free seam; the resize and normalization are the model's own.

Source code in src/deeperfly/pose2d/models.py

class LoadedModel:
    """A loaded detector plus its input contract and fused forward and decode.

    Wraps the torch module behind the torch-free seam; the resize and normalization
    are the model's own.
    """

    def __init__(self, spec: ModelSpec, module):
        """Wrap a loaded detector `module` with its `spec`."""
        self.spec = spec
        self.module = module

    @property
    def input_size(self) -> tuple[int, int]:
        """The network input `(height, width)`."""
        return self.spec.input_size

    @property
    def n_out_channels(self) -> int:
        """Number of output heatmap channels."""
        return self.spec.n_out_channels

    @property
    def joint_views(self) -> bool:
        """Whether the `V` axis is coupled: views computed together (the multiview net).

        `False` for a per-view detector, whose `V` axis is spare batch.
        """
        return bool(getattr(self.module, "joint_views", False))

    @property
    def padded_field(self) -> bool:
        """Whether a peak may land outside the model input: an off-frame joint."""
        return bool(getattr(self.module, "padded_field", False))

    @property
    def accepts_gray(self) -> bool:
        """Whether the model takes one-channel frames (both shipped classes do)."""
        return bool(getattr(self.module, "accepts_gray", False))

    def _impl(self):
        """The inference surface running every detector (prepare / predict / decode)."""
        from . import infer

        return infer

    def predict_points(self, inputs, *, method: str = "weighted", radius: int = 2):
        """Fused forward and decode -> input-normalized peaks and confidence.

        A peak may fall outside `[0, 1]` on a padded field: an off-frame joint.
        """
        return self._impl().predict_points(
            self.module, inputs, method=method, radius=radius
        )

    def predict_heatmaps(self, inputs):
        """Final heatmaps as host NumPy, for the candidate path above this layer."""
        return self._impl().predict_heatmaps(self.module, inputs)

    def set_precision(self, precision: str) -> None:
        """Set the forward precision (`float32` / `float16` / `bfloat16`)."""
        from .runtime import set_precision

        set_precision(self.module, precision)

    def device(self) -> str:
        """The device the model's parameters live on."""
        from . import runtime

        return runtime.detector_device(self.module)

accepts_gray property

accepts_gray: bool

Whether the model takes one-channel frames (both shipped classes do).

input_size property

input_size: tuple[int, int]

The network input (height, width).

joint_views property

joint_views: bool

Whether the V axis is coupled: views computed together (the multiview net).

False for a per-view detector, whose V axis is spare batch.

n_out_channels property

n_out_channels: int

Number of output heatmap channels.

padded_field property

padded_field: bool

Whether a peak may land outside the model input: an off-frame joint.

__init__

__init__(spec: ModelSpec, module)

Wrap a loaded detector module with its spec.

Source code in src/deeperfly/pose2d/models.py

def __init__(self, spec: ModelSpec, module):
    """Wrap a loaded detector `module` with its `spec`."""
    self.spec = spec
    self.module = module

device

device() -> str

The device the model's parameters live on.

Source code in src/deeperfly/pose2d/models.py

def device(self) -> str:
    """The device the model's parameters live on."""
    from . import runtime

    return runtime.detector_device(self.module)

predict_heatmaps

predict_heatmaps(inputs)

Final heatmaps as host NumPy, for the candidate path above this layer.

Source code in src/deeperfly/pose2d/models.py

def predict_heatmaps(self, inputs):
    """Final heatmaps as host NumPy, for the candidate path above this layer."""
    return self._impl().predict_heatmaps(self.module, inputs)

predict_points

predict_points(
    inputs, *, method: str = "weighted", radius: int = 2
)

Fused forward and decode -> input-normalized peaks and confidence.

A peak may fall outside [0, 1] on a padded field: an off-frame joint.

Source code in src/deeperfly/pose2d/models.py

def predict_points(self, inputs, *, method: str = "weighted", radius: int = 2):
    """Fused forward and decode -> input-normalized peaks and confidence.

    A peak may fall outside `[0, 1]` on a padded field: an off-frame joint.
    """
    return self._impl().predict_points(
        self.module, inputs, method=method, radius=radius
    )

set_precision

set_precision(precision: str) -> None

Set the forward precision (float32 / float16 / bfloat16).

Source code in src/deeperfly/pose2d/models.py

def set_precision(self, precision: str) -> None:
    """Set the forward precision (`float32` / `float16` / `bfloat16`)."""
    from .runtime import set_precision

    set_precision(self.module, precision)

ModelSpec dataclass

A torch-free description of one detector model.

Attributes:

Name	Type	Description
name	str	The model's name, referenced by a pathway's model key.
cls	str	The registry key selecting the model class (e.g. "mvt").
weights	str | None	A resolved checkpoint path. None is refused at load.
input_size	tuple[int, int]	The network input (height, width).
mean	float	Scalar subtracted from the [0, 1] image after the resize.
n_out_channels	int | None	Output heatmap channels (validated against the weights).
precision	str | None	Per-model forward precision, or None to inherit the default.
kwargs	dict	Extra class-specific construction kwargs.

Source code in src/deeperfly/pose2d/models.py

@dataclass(frozen=True)
class ModelSpec:
    """A torch-free description of one detector model.

    Attributes:
        name: The model's name, referenced by a pathway's `model` key.
        cls: The registry key selecting the model class (e.g. `"mvt"`).
        weights: A resolved checkpoint path. `None` is refused at load.
        input_size: The network input `(height, width)`.
        mean: Scalar subtracted from the `[0, 1]` image after the resize.
        n_out_channels: Output heatmap channels (validated against the weights).
        precision: Per-model forward precision, or `None` to inherit the default.
        kwargs: Extra class-specific construction kwargs.
    """

    name: str
    cls: str
    weights: str | None
    input_size: tuple[int, int] = DEFAULT_INPUT_SIZE
    mean: float = 0.0
    n_out_channels: int | None = None
    precision: str | None = None
    kwargs: dict = field(default_factory=dict)

artifact_grid

artifact_grid(blob: dict, path) -> str

The pixel_grid a detector artifact was trained on.

An artifact that records none predates the field, when every label was on the corner grid, so it reads as corner (ADR-0078); that also serves the released detectors on shared storage, which nobody rewrites.

Raises:

Type	Description
SystemExit	If the artifact records a grid this version does not know.

Source code in src/deeperfly/pose2d/models.py

def artifact_grid(blob: dict, path) -> str:
    """The `pixel_grid` a detector artifact was trained on.

    An artifact that records none predates the field, when every label was on the
    corner grid, so it reads as `corner` (ADR-0078); that also serves the released
    detectors on shared storage, which nobody rewrites.

    Raises:
        SystemExit: If the artifact records a grid this version does not know.
    """
    grid = blob.get("pixel_grid", "corner")
    if grid not in GRID_SHIFT:
        raise SystemExit(
            f"{path} records pixel grid {grid!r}; known: {sorted(GRID_SHIFT)}"
        )
    return str(grid)

first_model

first_model(models)

The module (the first member of an ensemble) whose geometry the run follows.

Source code in src/deeperfly/pose2d/models.py

def first_model(models):
    """The module (the first member of an ensemble) whose geometry the run follows."""
    return models[0] if isinstance(models, (list, tuple)) else models

grid_shift

grid_shift(models) -> float

The video-pixel shift that puts models' output on the pixel-center grid.

The mean over an ensemble's members, which is exact for a mean combine; a module built for training and never loaded from an artifact is on the pixel-center grid.

Source code in src/deeperfly/pose2d/models.py

def grid_shift(models) -> float:
    """The video-pixel shift that puts `models`' output on the pixel-center grid.

    The mean over an ensemble's members, which is exact for a mean combine; a module
    built for training and never loaded from an artifact is on the pixel-center grid.
    """
    members = models if isinstance(models, (list, tuple)) else [models]
    shifts = [GRID_SHIFT[getattr(m, "pixel_grid", "center")] for m in members]
    return sum(shifts) / len(shifts)

load_locator

load_locator(
    weights: str | Path,
    *,
    device: str | None = None,
    precision: str = DEFAULT_PRECISION,
) -> tuple[object, torch.device]

Load a full-frame locator artifact onto a device (ADR-0064).

Returns (model, device).

Raises:

Type	Description
ValueError	If the artifact is a pose detector rather than a locator.

Source code in src/deeperfly/pose2d/models.py

def load_locator(
    weights: str | Path,
    *,
    device: str | None = None,
    precision: str = DEFAULT_PRECISION,
) -> tuple[object, torch.device]:
    """Load a full-frame locator artifact onto a device (ADR-0064).

    Returns `(model, device)`.

    Raises:
        ValueError: If the artifact is a pose detector rather than a locator.
    """
    dev = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
    module = _load_artifact(weights, dev, precision)
    if getattr(module, "task", "pose") != "locate":
        raise ValueError(
            f"{weights} is a pose detector, not a locator; train one with "
            '`[data] task = "locate"` (see deeperfly.config.train)'
        )
    return module, dev

load_model

load_model(spec: ModelSpec) -> LoadedModel

Build a LoadedModel from a ModelSpec.

Raises:

Type	Description
SystemExit	On an unknown class, or one with no loadable checkpoint yet.
ValueError	If the module's channel count or trained input size disagrees with the spec.

Source code in src/deeperfly/pose2d/models.py

def load_model(spec: ModelSpec) -> LoadedModel:
    """Build a [LoadedModel][deeperfly.pose2d.models.LoadedModel] from a `ModelSpec`.

    Raises:
        SystemExit: On an unknown class, or one with no loadable checkpoint yet.
        ValueError: If the module's channel count or trained input size disagrees with
            the spec.
    """
    loader = MODEL_CLASSES.get(spec.cls)
    if loader is None:
        raise SystemExit(
            f"unknown model class {spec.cls!r} for model {spec.name!r}; "
            f"known classes: {sorted(MODEL_CLASSES)}"
        )
    module = loader(spec)
    num_classes = getattr(module, "num_classes", None)
    if (
        num_classes is not None
        and spec.n_out_channels is not None
        and int(num_classes) != int(spec.n_out_channels)
    ):
        raise ValueError(
            f"model {spec.name!r} declares n_out_channels={spec.n_out_channels} "
            f"but its weights emit {num_classes} channels"
        )
    input_hw = getattr(module, "input_hw", None)
    if input_hw is not None and tuple(int(v) for v in input_hw) != tuple(
        spec.input_size
    ):
        raise ValueError(
            f"model {spec.name!r} declares input_size={list(spec.input_size)} but its "
            f"weights were trained at {list(input_hw)}. Drop 'input_size' from its "
            "[[pose2d.models]] table to take the checkpoint's."
        )
    return LoadedModel(spec, module)

load_models

load_models(
    weights: str | Path | Sequence[str | Path],
    *,
    device: str | None = None,
    precision: str = DEFAULT_PRECISION,
) -> tuple[object, torch.device]

Load one multiview artifact, or an ensemble of them, onto a device.

precision defaults to "auto" here, not to the float32 reference, which is about 3x slower (ADR-0069). Returns (model, device): the module, or the list of modules of an ensemble.

Raises:

Type	Description
ValueError	If an artifact is not a joint-views (multiview) model, or the ensemble members disagree on input size, keypoints or the center channel.

Source code in src/deeperfly/pose2d/models.py

def load_models(
    weights: str | Path | Sequence[str | Path],
    *,
    device: str | None = None,
    precision: str = DEFAULT_PRECISION,
) -> tuple[object, torch.device]:
    """Load one multiview artifact, or an ensemble of them, onto a device.

    `precision` defaults to `"auto"` here, not to the float32 reference, which is about
    3x slower (ADR-0069). Returns `(model, device)`: the module, or the list of modules
    of an ensemble.

    Raises:
        ValueError: If an artifact is not a joint-views (multiview) model, or the
            ensemble members disagree on input size, keypoints or the center channel.
    """
    dev = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
    paths = [weights] if isinstance(weights, (str, Path)) else list(weights)
    models = [_load_artifact(w, dev, precision) for w in paths]
    for m in models:
        if not getattr(m, "joint_views", False):
            raise ValueError("predict runs multiview (joint-views) artifacts only")
    head = models[0]
    for m in models[1:]:
        if tuple(m.input_hw) != tuple(head.input_hw) or list(m.point_names) != list(
            head.point_names
        ):
            raise ValueError("ensemble members must share input size and keypoints")
        if bool(m.center_channel) != bool(head.center_channel):
            raise ValueError(
                "ensemble members must agree on the center channel; retrain the odd "
                "one with [data] center_channel to match"
            )
        if marks_others(m) != marks_others(head):
            raise ValueError(
                "ensemble members must agree on the other-animal mark: one batch of "
                "inputs serves them all; retrain the odd one with [data] mark_others "
                "to match"
            )
    return (models if len(models) > 1 else head), dev

marks_others

marks_others(models) -> bool

Whether the detector's inputs mark the other animals' centers (ADR-0148).

An ensemble's members agree (load_models), so the first one answers.

Source code in src/deeperfly/pose2d/models.py

def marks_others(models) -> bool:
    """Whether the detector's inputs mark the other animals' centers (ADR-0148).

    An ensemble's members agree ([load_models][]), so the first one answers.
    """
    return bool(getattr(first_model(models), "mark_others", False))

model_point_names

model_point_names(models) -> list[str]

The keypoint names the artifact was trained on, or [] if it recorded none.

Source code in src/deeperfly/pose2d/models.py

def model_point_names(models) -> list[str]:
    """The keypoint names the artifact was trained on, or `[]` if it recorded none."""
    return list(getattr(first_model(models), "point_names", []) or [])

The knobs of a prediction run, as the library reads them.

A run config writes them as [pose2d] and [pose2d.tracking], where Pose2dStage checks them.

PredictOptions dataclass

How a recording is predicted: a run's [pose2d] knobs, apart from its inputs.

Attributes:

Name	Type	Description
window	str	"static" predicts every frame through the run's fixed crops; "per_frame" plans a window per frame around the animals it finds.
animals	int	How many animals each frame holds; above one the detector needs a center channel and the run tracks them.
tracking	Tracking	The multi-animal tracker's knobs.
crop_aspect	Mapping[str, float] | None	Per named view, the aspect its window is drawn at instead of the input's (a run config's [pose2d] crop_aspect); it must equal the [data] view_aspect the detector trained with.
combine	str	How an ensemble's members are collapsed: "mean" or "median".
candidates	int	Distinct peaks kept per cell besides the decoded point; 0 for none.
tightness	float	How much wider than the animal a discovered window is: the static windows, and the size a per-frame window keeps.
batch_frames	int	Frames per block of the stream, each predicted once per animal.
flip_perm	Any | None	The skeleton's left/right permutation (a torch.Tensor) for flip TTA, or None for no mirrored pass; see flip_permutation.
point_names	Sequence[str] | None	Keypoint names to fall back on when the artifact recorded none.
compile	bool	Put the pose net's forward and readout through torch.compile in the pipelined passes, a tracker's with a locator (ADR-0139) and every static one (ADR-0150): about a third less time on a long run, after seconds to a minute of compiling, and keypoints that move by float rounding.

Source code in src/deeperfly/pose2d/options.py

@dataclass(frozen=True)
class PredictOptions:
    """How a recording is predicted: a run's `[pose2d]` knobs, apart from its inputs.

    Attributes:
        window: `"static"` predicts every frame through the run's fixed crops;
            `"per_frame"` plans a window per frame around the animals it finds.
        animals: How many animals each frame holds; above one the detector needs a
            center channel and the run tracks them.
        tracking: The multi-animal tracker's knobs.
        crop_aspect: Per named view, the aspect its window is drawn at instead of
            the input's (a run config's `[pose2d] crop_aspect`); it must equal the
            `[data] view_aspect` the detector trained with.
        combine: How an ensemble's members are collapsed: `"mean"` or `"median"`.
        candidates: Distinct peaks kept per cell besides the decoded point; 0 for
            none.
        tightness: How much wider than the animal a discovered window is: the
            static windows, and the size a per-frame window keeps.
        batch_frames: Frames per block of the stream, each predicted once per
            animal.
        flip_perm: The skeleton's left/right permutation (a `torch.Tensor`) for
            flip TTA, or `None` for no mirrored pass; see
            [flip_permutation][deeperfly.pose2d.predict.flip_permutation].
        point_names: Keypoint names to fall back on when the artifact recorded none.
        compile: Put the pose net's forward and readout through `torch.compile` in
            the pipelined passes, a tracker's with a locator (ADR-0139) and every
            static one (ADR-0150): about a third less time on a long run, after
            seconds to a minute of compiling, and keypoints that move by float
            rounding.
    """

    window: str = "static"
    animals: int = 1
    tracking: Tracking = field(default_factory=Tracking)
    crop_aspect: Mapping[str, float] | None = None
    combine: str = "mean"
    candidates: int = 0
    tightness: float = 1.5
    batch_frames: int = 8
    flip_perm: Any | None = None
    point_names: Sequence[str] | None = None
    compile: bool = False

    def __post_init__(self) -> None:
        """Refuse a window or a combine rule that is not one of the two."""
        if self.window not in ("static", "per_frame"):
            raise ValueError(
                f"window must be 'static' or 'per_frame', not {self.window!r}"
            )
        if self.combine not in ("mean", "median"):
            raise ValueError(
                f"combine must be 'mean' or 'median', got {self.combine!r}"
            )
        if not self.tightness > 0:
            raise ValueError(f"tightness must be above zero, got {self.tightness}")

__post_init__

__post_init__() -> None

Refuse a window or a combine rule that is not one of the two.

Source code in src/deeperfly/pose2d/options.py

def __post_init__(self) -> None:
    """Refuse a window or a combine rule that is not one of the two."""
    if self.window not in ("static", "per_frame"):
        raise ValueError(
            f"window must be 'static' or 'per_frame', not {self.window!r}"
        )
    if self.combine not in ("mean", "median"):
        raise ValueError(
            f"combine must be 'mean' or 'median', got {self.combine!r}"
        )
    if not self.tightness > 0:
        raise ValueError(f"tightness must be above zero, got {self.tightness}")

Tracking dataclass

How center peaks become animals and stay themselves, in video pixels (ADR-0063).

The defaults are [pose2d.tracking]'s (Tracking).

Attributes:

Name	Type	Description
min_score	float	What a center peak must score to count as a detection.
min_views	int	Views that must see a center before it is an animal.
gate_px	float	A view supports a 3D center when its nearest peak reprojects within this.
min_separation_px	float	Two centers closer than this in every view are one animal found twice (touching flies on hexaview are 85-95 px apart in the views across them, ADR-0131).
max_step_px	float	The frame-to-frame link gate, as the largest reprojected jump over the views.
hold_frames	int	Frames a track keeps its last position without a detection before the slot is freed.

Source code in src/deeperfly/pose2d/options.py

@dataclass(frozen=True)
class Tracking:
    """How center peaks become animals and stay themselves, in video pixels (ADR-0063).

    The defaults are `[pose2d.tracking]`'s ([Tracking][deeperfly.config.run.Tracking]).

    Attributes:
        min_score: What a center peak must score to count as a detection.
        min_views: Views that must see a center before it is an animal.
        gate_px: A view supports a 3D center when its nearest peak reprojects
            within this.
        min_separation_px: Two centers closer than this in every view are one
            animal found twice (touching flies on hexaview are 85-95 px apart in the
            views across them, ADR-0131).
        max_step_px: The frame-to-frame link gate, as the largest reprojected
            jump over the views.
        hold_frames: Frames a track keeps its last position without a detection
            before the slot is freed.
    """

    min_score: float = 0.001
    min_views: int = 2
    gate_px: float = 30.0
    min_separation_px: float = 50.0
    max_step_px: float = 60.0
    hold_frames: int = 12

Where the detector looks: the per-view windows, their warp, and the planners.

WindowGeometry is a moment's windows as one precomputed warp and its inverse (ADR-0069). The planners draw a window per frame around the animals a coarse pass found: Locator for one animal from the pose net's joints, Tracker for several from its center channel (ADR-0063) or a multi-animal locator (ADR-0127), and FullFrameLocator from a locator net over the whole source frame (ADR-0064).

FullFrameLocator

Bases: MeasuredPlanner

Per-frame windows from a locator net of its own, over the whole source frame.

The planner of a rig whose views share a source frame (ADR-0064). Each channel's peak passes through its view's mirror into view pixels and the measured window is centered there; a channel below conf_min keeps the previous box, then the letterbox. Two views of one video whose peaks are within SHARED_PEAK_FRACTION of the window keep only the stronger one's, and the other is missing in that frame. One animal only.

Source code in src/deeperfly/pose2d/windows.py

class FullFrameLocator(MeasuredPlanner):
    """Per-frame windows from a locator net of its own, over the whole source frame.

    The planner of a rig whose views share a source frame (ADR-0064). Each channel's
    peak passes through its view's mirror into view pixels and the measured window is
    centered there; a channel below `conf_min` keeps the previous box, then the
    letterbox. Two views of one video whose peaks are within `SHARED_PEAK_FRACTION`
    of the window keep only the stronger one's, and the other is missing in that
    frame. One animal only.
    """

    def __init__(
        self,
        locator,
        videos,
        input_hw,
        *,
        device,
        crop_aspect=None,
        conf_min: float = 0.1,
    ) -> None:
        """The locator's own source-frame geometry; `input_hw` is the pose net's."""
        self.conf_min = float(conf_min)
        super().__init__(videos, input_hw, crop_aspect)
        self.locator = SourceLocator(locator, videos, device=device)
        self.channel = self.locator.channel
        self.source_of = self.locator.source_of
        self.previous: list[tuple | None] = [None] * len(videos.views)
        #: `(T, V)` of the last planned block: views whose window was another view's.
        self.missing = np.zeros((0, len(videos.views)), dtype=bool)

    def reset(self) -> None:
        """Forget the previous boxes; the next block starts at the letterbox."""
        self.previous = [None] * len(self.previous)

    def coarse(self, blocks):
        """The locator over this block's source frames, at the locator's own input.

        Returns `(T, V, 1, 3)` video-pixel peaks.
        """
        return None, None, self.locator.peaks(blocks)

    def plan(self, points, conf, centers):
        """`(T, V, 1, 3)` video-pixel peaks -> the one animal's window per view."""
        peaks = np.asarray(centers, dtype=float)
        n_t, n_v = peaks.shape[0], peaks.shape[1]
        windows = np.zeros((n_t, 1, n_v, 4))
        self.missing = np.zeros((n_t, n_v), dtype=bool)
        for t in range(n_t):
            seen = [
                peaks[t, v, 0, 2] >= self.conf_min
                and bool(np.isfinite(peaks[t, v, 0, :2]).all())
                for v in range(n_v)
            ]
            self.missing[t] = _shared_windows(
                peaks[t, :, 0], seen, self.sizes, self.source_of
            )
            for v in range(n_v):
                x, y = peaks[t, v, 0, :2]
                if seen[v] and not self.missing[t, v]:
                    w, h = self.sizes[v]
                    box = (x - w / 2, y - h / 2, x + w / 2, y + h / 2)
                elif self.previous[v] is not None:
                    box = self.previous[v]
                else:
                    box = tuple(self.letterbox[v])
                self.previous[v] = box
                windows[t, 0, v] = box
        # One animal, and it is never lost: the fallback box always holds something.
        return windows, None, np.full((n_t, 1), TRACKED, dtype=np.uint8)

__init__

__init__(
    locator,
    videos,
    input_hw,
    *,
    device,
    crop_aspect=None,
    conf_min: float = 0.1,
) -> None

The locator's own source-frame geometry; input_hw is the pose net's.

Source code in src/deeperfly/pose2d/windows.py

def __init__(
    self,
    locator,
    videos,
    input_hw,
    *,
    device,
    crop_aspect=None,
    conf_min: float = 0.1,
) -> None:
    """The locator's own source-frame geometry; `input_hw` is the pose net's."""
    self.conf_min = float(conf_min)
    super().__init__(videos, input_hw, crop_aspect)
    self.locator = SourceLocator(locator, videos, device=device)
    self.channel = self.locator.channel
    self.source_of = self.locator.source_of
    self.previous: list[tuple | None] = [None] * len(videos.views)
    #: `(T, V)` of the last planned block: views whose window was another view's.
    self.missing = np.zeros((0, len(videos.views)), dtype=bool)

coarse

coarse(blocks)

The locator over this block's source frames, at the locator's own input.

Returns (T, V, 1, 3) video-pixel peaks.

Source code in src/deeperfly/pose2d/windows.py

def coarse(self, blocks):
    """The locator over this block's source frames, at the locator's own input.

    Returns `(T, V, 1, 3)` video-pixel peaks.
    """
    return None, None, self.locator.peaks(blocks)

plan

plan(points, conf, centers)

(T, V, 1, 3) video-pixel peaks -> the one animal's window per view.

Source code in src/deeperfly/pose2d/windows.py

def plan(self, points, conf, centers):
    """`(T, V, 1, 3)` video-pixel peaks -> the one animal's window per view."""
    peaks = np.asarray(centers, dtype=float)
    n_t, n_v = peaks.shape[0], peaks.shape[1]
    windows = np.zeros((n_t, 1, n_v, 4))
    self.missing = np.zeros((n_t, n_v), dtype=bool)
    for t in range(n_t):
        seen = [
            peaks[t, v, 0, 2] >= self.conf_min
            and bool(np.isfinite(peaks[t, v, 0, :2]).all())
            for v in range(n_v)
        ]
        self.missing[t] = _shared_windows(
            peaks[t, :, 0], seen, self.sizes, self.source_of
        )
        for v in range(n_v):
            x, y = peaks[t, v, 0, :2]
            if seen[v] and not self.missing[t, v]:
                w, h = self.sizes[v]
                box = (x - w / 2, y - h / 2, x + w / 2, y + h / 2)
            elif self.previous[v] is not None:
                box = self.previous[v]
            else:
                box = tuple(self.letterbox[v])
            self.previous[v] = box
            windows[t, 0, v] = box
    # One animal, and it is never lost: the fallback box always holds something.
    return windows, None, np.full((n_t, 1), TRACKED, dtype=np.uint8)

reset

reset() -> None

Forget the previous boxes; the next block starts at the letterbox.

Source code in src/deeperfly/pose2d/windows.py

def reset(self) -> None:
    """Forget the previous boxes; the next block starts at the letterbox."""
    self.previous = [None] * len(self.previous)

Locator

Bases: WindowPlanner

Per-frame windows: each view's box around the animal the coarse pass found.

Per (frame, view): the box of the confident coarse points, else the previous frame's, else the letterbox; causal, since at 125 fps a walking fly moves about 20 px a frame inside a window of hundreds. The single-animal planner: no rig, no center channel, and its one slot is never lost. On a rig that folds several views into one frame its coarse pass cannot tell the views' images apart (ADR-0064).

Source code in src/deeperfly/pose2d/windows.py

class Locator(WindowPlanner):
    """Per-frame windows: each view's box around the animal the coarse pass found.

    Per `(frame, view)`: the box of the confident coarse points, else the previous
    frame's, else the letterbox; causal, since at 125 fps a walking fly moves about
    20 px a frame inside a window of hundreds. The single-animal planner: no rig, no
    center channel, and its one slot is never lost. On a rig that folds several views
    into one frame its coarse pass cannot tell the views' images apart (ADR-0064).
    """

    def __init__(
        self,
        videos,
        input_hw,
        *,
        crop_aspect=None,
        tightness: float = 1.5,
        conf_min: float = 0.3,
        min_points: int = 4,
    ) -> None:
        """Size the boxes from the pose net's `input_hw` over `videos`' frame sizes."""
        self.frame_hw = list(videos.frame_hw)
        self.aspects = _view_aspects(videos.views, input_hw, crop_aspect)
        self.min_w = input_hw[1] / 2
        self.tightness = tightness
        self.conf_min = conf_min
        self.min_points = min_points
        self.letterbox = [padded_window(hw, input_hw) for hw in self.frame_hw]
        self.previous: list[tuple | None] = [None] * len(self.frame_hw)

    def reset(self) -> None:
        """Forget the previous boxes; the next block starts at the letterbox."""
        self.previous = [None] * len(self.frame_hw)

    def plan(self, points, conf, centers):
        """The one animal's box per `(frame, view)`; no centers, and nothing lost."""
        boxes = self.boxes(points, conf)
        flags = np.full(boxes.shape[:1] + (1,), TRACKED, dtype=np.uint8)
        return boxes[:, None], None, flags

    def boxes(self, pts: np.ndarray, conf: np.ndarray) -> np.ndarray:
        """`(T, V, K, 2)` coarse points + `(T, V, K)` conf -> `(T, V, 4)` boxes."""
        n_t, n_v = pts.shape[:2]
        out = np.zeros((n_t, n_v, 4))
        for t in range(n_t):
            fixes = [
                _frame_box(
                    pts[t, v],
                    conf[t, v],
                    min_points=self.min_points,
                    conf_min=self.conf_min,
                )
                for v in range(n_v)
            ]
            for v in range(n_v):
                if fixes[v] is not None:
                    center, extent = fixes[v]
                    box = _aspect_box(
                        center, extent, self.aspects[v], self.tightness, self.min_w
                    )
                elif self.previous[v] is not None:
                    box = self.previous[v]
                else:
                    box = self.letterbox[v]
                self.previous[v] = box
                out[t, v] = box
        return out

__init__

__init__(
    videos,
    input_hw,
    *,
    crop_aspect=None,
    tightness: float = 1.5,
    conf_min: float = 0.3,
    min_points: int = 4,
) -> None

Size the boxes from the pose net's input_hw over videos' frame sizes.

Source code in src/deeperfly/pose2d/windows.py

def __init__(
    self,
    videos,
    input_hw,
    *,
    crop_aspect=None,
    tightness: float = 1.5,
    conf_min: float = 0.3,
    min_points: int = 4,
) -> None:
    """Size the boxes from the pose net's `input_hw` over `videos`' frame sizes."""
    self.frame_hw = list(videos.frame_hw)
    self.aspects = _view_aspects(videos.views, input_hw, crop_aspect)
    self.min_w = input_hw[1] / 2
    self.tightness = tightness
    self.conf_min = conf_min
    self.min_points = min_points
    self.letterbox = [padded_window(hw, input_hw) for hw in self.frame_hw]
    self.previous: list[tuple | None] = [None] * len(self.frame_hw)

boxes

boxes(pts: ndarray, conf: ndarray) -> np.ndarray

(T, V, K, 2) coarse points + (T, V, K) conf -> (T, V, 4) boxes.

Source code in src/deeperfly/pose2d/windows.py

def boxes(self, pts: np.ndarray, conf: np.ndarray) -> np.ndarray:
    """`(T, V, K, 2)` coarse points + `(T, V, K)` conf -> `(T, V, 4)` boxes."""
    n_t, n_v = pts.shape[:2]
    out = np.zeros((n_t, n_v, 4))
    for t in range(n_t):
        fixes = [
            _frame_box(
                pts[t, v],
                conf[t, v],
                min_points=self.min_points,
                conf_min=self.conf_min,
            )
            for v in range(n_v)
        ]
        for v in range(n_v):
            if fixes[v] is not None:
                center, extent = fixes[v]
                box = _aspect_box(
                    center, extent, self.aspects[v], self.tightness, self.min_w
                )
            elif self.previous[v] is not None:
                box = self.previous[v]
            else:
                box = self.letterbox[v]
            self.previous[v] = box
            out[t, v] = box
    return out

plan

plan(points, conf, centers)

The one animal's box per (frame, view); no centers, and nothing lost.

Source code in src/deeperfly/pose2d/windows.py

def plan(self, points, conf, centers):
    """The one animal's box per `(frame, view)`; no centers, and nothing lost."""
    boxes = self.boxes(points, conf)
    flags = np.full(boxes.shape[:1] + (1,), TRACKED, dtype=np.uint8)
    return boxes[:, None], None, flags

reset

reset() -> None

Forget the previous boxes; the next block starts at the letterbox.

Source code in src/deeperfly/pose2d/windows.py

def reset(self) -> None:
    """Forget the previous boxes; the next block starts at the letterbox."""
    self.previous = [None] * len(self.frame_hw)

MeasuredPlanner

Bases: WindowPlanner

A planner whose window is one measured size per view, around a located animal.

The window travels with the animal and its scale does not, so the detector sees the animal at the scale it trained on. Until discover_sizes runs, sizes is as wide as the input at the view's aspect.

Attributes:

Name	Type	Description
aspects	list[float]	Per view, the width / height its window is drawn at.
min_w	float	Half the input's width, so a view is never upsampled more than 2x.
sizes	ndarray	(V, 2) the (width, height) of each view's window, in video pixels.

Source code in src/deeperfly/pose2d/windows.py

class MeasuredPlanner(WindowPlanner):
    """A planner whose window is one measured size per view, around a located animal.

    The window travels with the animal and its scale does not, so the detector sees the
    animal at the scale it trained on. Until `discover_sizes` runs, `sizes` is as wide
    as the input at the view's aspect.

    Attributes:
        aspects: Per view, the width / height its window is drawn at.
        min_w: Half the input's width, so a view is never upsampled more than 2x.
        sizes: `(V, 2)` the `(width, height)` of each view's window, in video pixels.
    """

    aspects: list[float]
    min_w: float
    sizes: np.ndarray

    def __init__(self, videos, input_hw, crop_aspect=None) -> None:
        """Size the windows of `videos` for a pose net taking `input_hw`."""
        self.aspects = _view_aspects(videos.views, input_hw, crop_aspect)
        self.min_w = input_hw[1] / 2
        self.letterbox = np.array(
            [padded_window(hw, input_hw) for hw in videos.frame_hw], dtype=float
        )
        self.sizes = np.array(
            [[2 * self.min_w, 2 * self.min_w / a] for a in self.aspects], dtype=float
        )

__init__

__init__(videos, input_hw, crop_aspect=None) -> None

Size the windows of videos for a pose net taking input_hw.

Source code in src/deeperfly/pose2d/windows.py

def __init__(self, videos, input_hw, crop_aspect=None) -> None:
    """Size the windows of `videos` for a pose net taking `input_hw`."""
    self.aspects = _view_aspects(videos.views, input_hw, crop_aspect)
    self.min_w = input_hw[1] / 2
    self.letterbox = np.array(
        [padded_window(hw, input_hw) for hw in videos.frame_hw], dtype=float
    )
    self.sizes = np.array(
        [[2 * self.min_w, 2 * self.min_w / a] for a in self.aspects], dtype=float
    )

SourceLocator

A locator net over each whole source frame, its peaks in view pixels (ADR-0064).

Channel k of the net is the view point_names[k], so a peak comes back already attributed to a view. A source is decoded, uploaded and run once however many views read it. Shared by FullFrameLocator, one peak a view, and Tracker, several.

Source code in src/deeperfly/pose2d/windows.py

class SourceLocator:
    """A locator net over each whole source frame, its peaks in view pixels (ADR-0064).

    Channel `k` of the net is the view `point_names[k]`, so a peak comes back already
    attributed to a view. A source is decoded, uploaded and run once however many views
    read it. Shared by [FullFrameLocator][deeperfly.pose2d.windows.FullFrameLocator],
    one peak a view, and [Tracker][deeperfly.pose2d.windows.Tracker], several.
    """

    def __init__(self, locator, videos, *, device) -> None:
        """The locator's own source-frame geometry over `videos`."""
        self.model = locator
        self.device = device
        names = list(getattr(locator, "point_names", []))
        missing = [c for c in videos.views if c not in names]
        if missing:
            raise ValueError(
                f"the locator has no channel for {missing}; its channels are {names}. "
                "Train it on projects that have these views."
            )
        #: Per view, which locator channel carries it.
        self.channel = [names.index(c) for c in videos.views]
        self.source_of = list(videos.source_of)
        loc_hw = tuple(int(v) for v in locator.input_hw)
        self.geom = WindowGeometry.build(
            [padded_window(hw, loc_hw) for hw in videos.source_hw],
            list(videos.source_hw),
            loc_hw,
            device=device,
        )

    def peaks(self, blocks, k: int = 1, animals: int = 1) -> np.ndarray:
        """This block's source frames -> `(T, V, k, 3)` video-pixel `(x, y, score)`.

        `k = 1` is the argmax readout; above it, the `k` strongest distinct maxima of
        each channel, NaN at score 0 where a channel has fewer, a blob that holds
        several of the `animals` split among them
        ([split_shared_peaks][deeperfly.pose2d.candidates.split_shared_peaks]).
        """
        with torch.no_grad():
            images = self.geom.warp(blocks)
            n_t, n_s = images.shape[:2]
            dtype = autocast_dtype(self.model, self.device)
            with torch.autocast(
                self.device.type, dtype=dtype, enabled=dtype is not None
            ):
                out = self.model(images)
            n_c = self.model.num_classes
            field = out.reshape(n_t * n_s, n_c, *out.shape[-2:])
            if k == 1:
                pts, conf = decode_points(self.model, field)
            else:
                pts, conf = decode_candidates(self.model, field, k)
                if animals > 1:
                    pts, conf = self._split(field, pts, conf, animals)
        # Source pixels, per source: the geometry's inverse is the letterbox alone.
        seen = self.geom.to_video_px(pts.float().reshape(n_t, n_s, n_c * k, 2))
        seen = seen.reshape(n_t, n_s, n_c, k, 2).cpu().numpy()
        scored = conf.float().reshape(n_t, n_s, n_c, k).cpu().numpy()
        peaks = np.full((n_t, len(self.channel), k, 3), np.nan)
        for v, (c, s) in enumerate(zip(self.channel, self.source_of, strict=True)):
            peaks[:, v, :, :2] = seen[:, s, c]
            peaks[:, v, :, 2] = scored[:, s, c]
        return peaks

    def _split(self, field, pts, conf, animals: int):
        """Each frame's and channel's peaks through `split_shared_peaks`.

        In the field's cells, and back to input-normalized coordinates.
        """
        geom = self.model.geom
        prob = field.float()
        if not getattr(self.model, "field_is_prob", False):
            prob = torch.softmax(prob.flatten(-2), dim=-1).view(prob.shape)
        prob = prob.cpu().numpy()
        cells = geom.to_field_cells(pts.float().cpu().numpy())
        scores = conf.float().cpu().numpy()
        for n in range(cells.shape[0]):
            for c in range(cells.shape[1]):
                cells[n, c], scores[n, c] = split_shared_peaks(
                    prob[n, c], cells[n, c], scores[n, c], animals
                )
        back = geom.to_input_normalized(cells)
        return torch.as_tensor(back, device=pts.device), torch.as_tensor(
            scores, device=conf.device
        )

__init__

__init__(locator, videos, *, device) -> None

The locator's own source-frame geometry over videos.

Source code in src/deeperfly/pose2d/windows.py

def __init__(self, locator, videos, *, device) -> None:
    """The locator's own source-frame geometry over `videos`."""
    self.model = locator
    self.device = device
    names = list(getattr(locator, "point_names", []))
    missing = [c for c in videos.views if c not in names]
    if missing:
        raise ValueError(
            f"the locator has no channel for {missing}; its channels are {names}. "
            "Train it on projects that have these views."
        )
    #: Per view, which locator channel carries it.
    self.channel = [names.index(c) for c in videos.views]
    self.source_of = list(videos.source_of)
    loc_hw = tuple(int(v) for v in locator.input_hw)
    self.geom = WindowGeometry.build(
        [padded_window(hw, loc_hw) for hw in videos.source_hw],
        list(videos.source_hw),
        loc_hw,
        device=device,
    )

peaks

peaks(blocks, k: int = 1, animals: int = 1) -> np.ndarray

This block's source frames -> (T, V, k, 3) video-pixel (x, y, score).

k = 1 is the argmax readout; above it, the k strongest distinct maxima of each channel, NaN at score 0 where a channel has fewer, a blob that holds several of the animals split among them (split_shared_peaks).

Source code in src/deeperfly/pose2d/windows.py

def peaks(self, blocks, k: int = 1, animals: int = 1) -> np.ndarray:
    """This block's source frames -> `(T, V, k, 3)` video-pixel `(x, y, score)`.

    `k = 1` is the argmax readout; above it, the `k` strongest distinct maxima of
    each channel, NaN at score 0 where a channel has fewer, a blob that holds
    several of the `animals` split among them
    ([split_shared_peaks][deeperfly.pose2d.candidates.split_shared_peaks]).
    """
    with torch.no_grad():
        images = self.geom.warp(blocks)
        n_t, n_s = images.shape[:2]
        dtype = autocast_dtype(self.model, self.device)
        with torch.autocast(
            self.device.type, dtype=dtype, enabled=dtype is not None
        ):
            out = self.model(images)
        n_c = self.model.num_classes
        field = out.reshape(n_t * n_s, n_c, *out.shape[-2:])
        if k == 1:
            pts, conf = decode_points(self.model, field)
        else:
            pts, conf = decode_candidates(self.model, field, k)
            if animals > 1:
                pts, conf = self._split(field, pts, conf, animals)
    # Source pixels, per source: the geometry's inverse is the letterbox alone.
    seen = self.geom.to_video_px(pts.float().reshape(n_t, n_s, n_c * k, 2))
    seen = seen.reshape(n_t, n_s, n_c, k, 2).cpu().numpy()
    scored = conf.float().reshape(n_t, n_s, n_c, k).cpu().numpy()
    peaks = np.full((n_t, len(self.channel), k, 3), np.nan)
    for v, (c, s) in enumerate(zip(self.channel, self.source_of, strict=True)):
        peaks[:, v, :, :2] = seen[:, s, c]
        peaks[:, v, :, 2] = scored[:, s, c]
    return peaks

Tracker

Bases: MeasuredPlanner

The animals of each frame, and the windows the fine pass runs on (ADR-0063).

One per recording, so the tracks and sizes survive the block loop. A step is assign_centers, then link_step, then center_windows, with the run config's [pose2d.tracking] knobs. The center peaks come from the pose net's center channel, or, given a locator, from a locator trained on several animals, which a rig whose views share a frame needs (ADR-0127).

Source code in src/deeperfly/pose2d/windows.py

class Tracker(MeasuredPlanner):
    """The animals of each frame, and the windows the fine pass runs on (ADR-0063).

    One per recording, so the tracks and `sizes` survive the block loop. A step is
    [assign_centers][deeperfly.instances.assign_centers], then
    [link_step][deeperfly.instances.link_step], then
    [center_windows][deeperfly.instances.center_windows], with the run config's
    `[pose2d.tracking]` knobs. The center peaks come from the pose net's center
    channel, or, given a `locator`, from a locator trained on several animals, which a
    rig whose views share a frame needs (ADR-0127).
    """

    def __init__(
        self,
        rig: Rig,
        videos,
        input_hw,
        *,
        animals: int,
        tracking: Tracking,
        crop_aspect=None,
        locator=None,
        device=None,
    ) -> None:
        """Plan `animals` slots through `rig` with the `tracking` knobs.

        Raises:
            ValueError: If `locator` was trained on one animal: its channels hold a
                single peak, so a second animal has nowhere to go.
        """
        self.rig = rig
        self.animals = int(animals)
        self.tracking = tracking
        # Center peaks read per view: the animals, plus room for a spurious one.
        self.n_centers = self.animals + 2
        self.locator = None
        if locator is not None:
            trained = int(getattr(locator, "animals", 1))
            if trained < 2:
                raise ValueError(
                    "locator: trained on one animal, so each channel is a single "
                    f"peak and cannot locate {self.animals}; train one on projects "
                    "that label several animals in a frame"
                )
            self.locator = SourceLocator(locator, videos, device=device)
        super().__init__(videos, input_hw, crop_aspect)
        self.reset()

    def coarse(self, blocks):
        """The locator's `(T, V, n_centers, 3)` peaks, or `None` for the pose net's."""
        if self.locator is None:
            return None
        return None, None, self.locator.peaks(blocks, self.n_centers, self.animals)

    def reset(self) -> None:
        """Forget every track, so the next step starts the slots over."""
        self.last = np.full((self.animals, 3), np.nan)
        self.age = np.zeros(self.animals, dtype=int)

    def plan(self, points, conf, centers):
        """`(T, V, C, 3)` video-pixel center peaks -> this block's animals.

        Returns:
            `(T, N, V, 4)` windows, `(T, N, 3)` centers and `(T, N)` uint8 track flags.
        """
        cands_px = centers
        n_t, n_v = cands_px.shape[0], cands_px.shape[1]
        centers = np.full((n_t, self.animals, 3), np.nan)
        flags = np.full((n_t, self.animals), LOST, dtype=np.uint8)
        windows = np.zeros((n_t, self.animals, n_v, 4))
        knobs = self.tracking
        # The field hands back as many peaks as it is asked for; below the floor they
        # are noise, and noise that three views happen to agree on is an animal.
        peaks = np.array(cands_px, dtype=float)
        peaks[peaks[..., 2] < knobs.min_score] = np.nan
        for t in range(n_t):
            found, _ = assign_centers(
                self.rig,
                peaks[t],
                min_views=knobs.min_views,
                gate_px=knobs.gate_px,
                min_separation_px=knobs.min_separation_px,
                max_animals=self.animals,
            )
            centers[t], flags[t], self.last, self.age = link_step(
                self.rig,
                self.last,
                self.age,
                found,
                max_step_px=knobs.max_step_px,
                hold_frames=knobs.hold_frames,
            )
            windows[t] = center_windows(
                self.rig, centers[t], self.sizes, self.letterbox
            )
        return windows, centers, flags

__init__

__init__(
    rig: Rig,
    videos,
    input_hw,
    *,
    animals: int,
    tracking: Tracking,
    crop_aspect=None,
    locator=None,
    device=None,
) -> None

Plan animals slots through rig with the tracking knobs.

Raises:

Type	Description
ValueError	If locator was trained on one animal: its channels hold a single peak, so a second animal has nowhere to go.

Source code in src/deeperfly/pose2d/windows.py

def __init__(
    self,
    rig: Rig,
    videos,
    input_hw,
    *,
    animals: int,
    tracking: Tracking,
    crop_aspect=None,
    locator=None,
    device=None,
) -> None:
    """Plan `animals` slots through `rig` with the `tracking` knobs.

    Raises:
        ValueError: If `locator` was trained on one animal: its channels hold a
            single peak, so a second animal has nowhere to go.
    """
    self.rig = rig
    self.animals = int(animals)
    self.tracking = tracking
    # Center peaks read per view: the animals, plus room for a spurious one.
    self.n_centers = self.animals + 2
    self.locator = None
    if locator is not None:
        trained = int(getattr(locator, "animals", 1))
        if trained < 2:
            raise ValueError(
                "locator: trained on one animal, so each channel is a single "
                f"peak and cannot locate {self.animals}; train one on projects "
                "that label several animals in a frame"
            )
        self.locator = SourceLocator(locator, videos, device=device)
    super().__init__(videos, input_hw, crop_aspect)
    self.reset()

coarse

coarse(blocks)

The locator's (T, V, n_centers, 3) peaks, or None for the pose net's.

Source code in src/deeperfly/pose2d/windows.py

def coarse(self, blocks):
    """The locator's `(T, V, n_centers, 3)` peaks, or `None` for the pose net's."""
    if self.locator is None:
        return None
    return None, None, self.locator.peaks(blocks, self.n_centers, self.animals)

plan

plan(points, conf, centers)

(T, V, C, 3) video-pixel center peaks -> this block's animals.

Returns:

Type	Description
	(T, N, V, 4) windows, (T, N, 3) centers and (T, N) uint8 track flags.

Source code in src/deeperfly/pose2d/windows.py

def plan(self, points, conf, centers):
    """`(T, V, C, 3)` video-pixel center peaks -> this block's animals.

    Returns:
        `(T, N, V, 4)` windows, `(T, N, 3)` centers and `(T, N)` uint8 track flags.
    """
    cands_px = centers
    n_t, n_v = cands_px.shape[0], cands_px.shape[1]
    centers = np.full((n_t, self.animals, 3), np.nan)
    flags = np.full((n_t, self.animals), LOST, dtype=np.uint8)
    windows = np.zeros((n_t, self.animals, n_v, 4))
    knobs = self.tracking
    # The field hands back as many peaks as it is asked for; below the floor they
    # are noise, and noise that three views happen to agree on is an animal.
    peaks = np.array(cands_px, dtype=float)
    peaks[peaks[..., 2] < knobs.min_score] = np.nan
    for t in range(n_t):
        found, _ = assign_centers(
            self.rig,
            peaks[t],
            min_views=knobs.min_views,
            gate_px=knobs.gate_px,
            min_separation_px=knobs.min_separation_px,
            max_animals=self.animals,
        )
        centers[t], flags[t], self.last, self.age = link_step(
            self.rig,
            self.last,
            self.age,
            found,
            max_step_px=knobs.max_step_px,
            hold_frames=knobs.hold_frames,
        )
        windows[t] = center_windows(
            self.rig, centers[t], self.sizes, self.letterbox
        )
    return windows, centers, flags

reset

reset() -> None

Forget every track, so the next step starts the slots over.

Source code in src/deeperfly/pose2d/windows.py

def reset(self) -> None:
    """Forget every track, so the next step starts the slots over."""
    self.last = np.full((self.animals, 3), np.nan)
    self.age = np.zeros(self.animals, dtype=int)

WindowGeometry dataclass

Everything constant about a moment's warp and its inverse, precomputed on device.

Built by build and shared by every moment of a recording.

Attributes:

Name	Type	Description
input_hw	tuple[int, int]	The network input (height, width) every view is warped onto.
frame_hw	tuple[tuple[int, int], ...]	Per window, the (height, width) of the frame it is read out of.
source_of	tuple[int, ...]	Per window, which of the frames fill is given it reads; views of one video share a row.
shrink_hw	tuple[tuple[int, int], ...]	Per view, the pre-shrink size, frame_hw where none is needed.
canvas_hw	tuple[int, int]	The common (height, width) the shrunk frames are written into.
windows	tuple[Affine2D, ...]	Per view, the window affine (video pixels -> input pixels), mirror included.
grid	Tensor	(V, h, w, 2) grid_sample grid, canvas -> input.
inv_linear	Tensor	(V, 2, 2) linear part of each window's inverse.
inv_offset	Tensor	(V, 2) translation of each window's inverse.
px_scale	Tensor	(2,) [width, height], input-normalized -> input pixels.

Source code in src/deeperfly/pose2d/windows.py

@dataclass(frozen=True)
class WindowGeometry:
    """Everything constant about a moment's warp and its inverse, precomputed on device.

    Built by [build][deeperfly.pose2d.windows.WindowGeometry.build] and shared by every
    moment of a recording.

    Attributes:
        input_hw: The network input `(height, width)` every view is warped onto.
        frame_hw: Per window, the `(height, width)` of the frame it is read out of.
        source_of: Per window, which of the frames `fill` is given it reads; views of
            one video share a row.
        shrink_hw: Per view, the pre-shrink size, `frame_hw` where none is needed.
        canvas_hw: The common `(height, width)` the shrunk frames are written into.
        windows: Per view, the window affine (video pixels -> input pixels), mirror
            included.
        grid: `(V, h, w, 2)` `grid_sample` grid, canvas -> input.
        inv_linear: `(V, 2, 2)` linear part of each window's inverse.
        inv_offset: `(V, 2)` translation of each window's inverse.
        px_scale: `(2,)` `[width, height]`, input-normalized -> input pixels.
    """

    input_hw: tuple[int, int]
    frame_hw: tuple[tuple[int, int], ...]
    source_of: tuple[int, ...]
    shrink_hw: tuple[tuple[int, int], ...]
    canvas_hw: tuple[int, int]
    windows: tuple[Affine2D, ...]
    grid: torch.Tensor
    inv_linear: torch.Tensor
    inv_offset: torch.Tensor
    px_scale: torch.Tensor

    @property
    def n_views(self) -> int:
        """How many views a moment has."""
        return len(self.windows)

    @classmethod
    def build(
        cls,
        boxes: Sequence[Sequence[float]],
        frame_hw: Sequence[Sequence[int]],
        input_hw: Sequence[int],
        *,
        device,
        sources: Sequence[int] | None = None,
        views: Sequence[bool] | None = None,
    ) -> WindowGeometry:
        """Precompute the warp and its inverse for one set of per-view windows.

        `boxes` are `(x0, y0, x1, y1)` slice bounds in each frame's own pixels.
        `sources` says which frame given to `fill` each window reads (`None`: one each,
        in order), and `views` whether each is mirrored (`None`: none is).
        """
        boxes = [tuple(float(v) for v in b) for b in boxes]
        sizes = [(int(hw[0]), int(hw[1])) for hw in frame_hw]
        if len(boxes) != len(sizes):
            raise ValueError(
                f"{len(boxes)} window(s) for {len(sizes)} view(s); the windows are "
                "per view, so one box per view is needed"
            )
        source_of = (
            list(range(len(boxes))) if sources is None else [int(s) for s in sources]
        )
        flips = [False] * len(boxes) if views is None else [bool(v) for v in views]
        for name, given in (("source", source_of), ("mirror", flips)):
            if len(given) != len(boxes):
                raise ValueError(
                    f"{len(given)} {name}(s) for {len(boxes)} window(s); they are per "
                    "window, so a static run repeats them per animal"
                )
        ih, iw = int(input_hw[0]), int(input_hw[1])
        windows = [
            window_affine(b, (ih, iw), m) for b, m in zip(boxes, flips, strict=True)
        ]
        # The mean scale host-side, once: a per-moment `det` would sync the device. A
        # mirror has unit scale, so it changes no view's pre-shrink.
        scales = [
            float(torch.sqrt(torch.abs(torch.det(a.matrix[:, :2].double()))))
            for a in windows
        ]
        shrink = [
            (max(1, round(h * f)), max(1, round(w * f)))
            if f < ANTIALIAS_MAX_SCALE
            else (h, w)
            for (h, w), f in zip(sizes, scales, strict=True)
        ]
        canvas = (max(s[0] for s in shrink), max(s[1] for s in shrink))
        # Each view samples the common canvas, so its theta is normalized against the
        # canvas; in pixel coordinates the map is the same one `warp_image` would build.
        theta = torch.stack(
            [
                (
                    win @ Affine2D.resize(size, shr).inverse() if shr != size else win
                ).grid_theta(canvas, (ih, iw))
                for win, size, shr in zip(windows, sizes, shrink, strict=True)
            ]
        ).to(device=device, dtype=torch.float32)
        grid = F.affine_grid(theta, (len(boxes), 1, ih, iw), align_corners=False)
        inv = [w.inverse().matrix for w in windows]
        return cls(
            input_hw=(ih, iw),
            frame_hw=tuple(sizes),
            source_of=tuple(source_of),
            shrink_hw=tuple(shrink),
            canvas_hw=canvas,
            windows=tuple(windows),
            grid=grid,
            inv_linear=torch.stack([m[:, :2] for m in inv]).to(device),
            inv_offset=torch.stack([m[:, 2] for m in inv]).to(device),
            px_scale=torch.tensor([float(iw), float(ih)], device=device),
        )

    def new_canvas(self, n_frames: int = 1) -> torch.Tensor:
        """A zero `(V, T, H, W)` canvas for `n_frames` moments, on the grid's device.

        It must stay zero outside each view's shrunk frame, which is what makes the
        common canvas identical to per-view buffers; reusing one only overwrites the
        written region.
        """
        return torch.zeros(
            (self.n_views, int(n_frames), *self.canvas_hw),
            device=self.grid.device,
            dtype=torch.float32,
        )

    def fill(self, canvas: torch.Tensor, frames) -> torch.Tensor:
        """Write each window's `(T, H, W)` or `(H, W)` source frame into `canvas`.

        `frames` is indexed by `source_of`: host uint8 arrays, or a stacked on-device
        tensor (the online path's buffer, indexed without a copy). Integer input is
        scaled by 1/255, floats are taken as normalized; a shared source is uploaded,
        and pre-shrunk, once. Returns `canvas`.
        """
        ready: dict[tuple[int, int, int], torch.Tensor] = {}
        for v, s in enumerate(self.source_of):
            sh, sw = self.shrink_hw[v]
            img = ready.get((s, sh, sw))
            if img is None:
                frame = frames[s]
                img = (
                    frame
                    if torch.is_tensor(frame)
                    else torch.from_numpy(np.asarray(frame))
                )
                img = img.to(device=canvas.device, non_blocking=True)
                img = (
                    img.float().div_(255.0)
                    if not img.is_floating_point()
                    else img.float()
                )
                if img.ndim == 2:
                    img = img[None]
                if (sh, sw) != self.frame_hw[v]:
                    img = F.interpolate(
                        img[:, None],
                        size=(sh, sw),
                        mode="bilinear",
                        align_corners=False,
                        antialias=True,
                    )[:, 0]
                ready[(s, sh, sw)] = img
            canvas[v, : img.shape[0], :sh, :sw] = img
        return canvas

    def sample(self, canvas: torch.Tensor) -> torch.Tensor:
        """`(V, T, H, W)` canvas -> the `(T, V, 1, h, w)` network input, in one kernel.

        `T` rides in `grid_sample`'s channel axis, so one grid of `V` rows covers a
        whole block of moments.
        """
        out = F.grid_sample(
            canvas,
            self.grid,
            mode="bilinear",
            padding_mode="zeros",
            align_corners=False,
        )
        return out.permute(1, 0, 2, 3).unsqueeze(2)

    def warp(self, frames) -> torch.Tensor:
        """The sources' `(T, H, W)` uint8 frames -> the `(T, V, 1, h, w)` input."""
        first = frames[0] if hasattr(frames, "__getitem__") else next(iter(frames))
        n = 1 if np.ndim(first) == 2 else int(np.shape(first)[0])
        return self.sample(self.fill(self.new_canvas(n), frames))

    def to_video_px(self, pts_norm: torch.Tensor) -> torch.Tensor:
        """`(..., V, N, 2)` input-normalized points -> video pixels, all views at once.

        `N` is whatever the caller stacks per view: keypoints, or keypoints times
        candidates.
        """
        px = pts_norm * self.px_scale
        return (
            torch.einsum("...vkj,vij->...vki", px, self.inv_linear.to(px.dtype))
            + self.inv_offset.to(px.dtype)[:, None, :]
        )

n_views property

n_views: int

How many views a moment has.

build classmethod

build(
    boxes: Sequence[Sequence[float]],
    frame_hw: Sequence[Sequence[int]],
    input_hw: Sequence[int],
    *,
    device,
    sources: Sequence[int] | None = None,
    views: Sequence[bool] | None = None,
) -> WindowGeometry

Precompute the warp and its inverse for one set of per-view windows.

boxes are (x0, y0, x1, y1) slice bounds in each frame's own pixels. sources says which frame given to fill each window reads (None: one each, in order), and views whether each is mirrored (None: none is).

Source code in src/deeperfly/pose2d/windows.py

@classmethod
def build(
    cls,
    boxes: Sequence[Sequence[float]],
    frame_hw: Sequence[Sequence[int]],
    input_hw: Sequence[int],
    *,
    device,
    sources: Sequence[int] | None = None,
    views: Sequence[bool] | None = None,
) -> WindowGeometry:
    """Precompute the warp and its inverse for one set of per-view windows.

    `boxes` are `(x0, y0, x1, y1)` slice bounds in each frame's own pixels.
    `sources` says which frame given to `fill` each window reads (`None`: one each,
    in order), and `views` whether each is mirrored (`None`: none is).
    """
    boxes = [tuple(float(v) for v in b) for b in boxes]
    sizes = [(int(hw[0]), int(hw[1])) for hw in frame_hw]
    if len(boxes) != len(sizes):
        raise ValueError(
            f"{len(boxes)} window(s) for {len(sizes)} view(s); the windows are "
            "per view, so one box per view is needed"
        )
    source_of = (
        list(range(len(boxes))) if sources is None else [int(s) for s in sources]
    )
    flips = [False] * len(boxes) if views is None else [bool(v) for v in views]
    for name, given in (("source", source_of), ("mirror", flips)):
        if len(given) != len(boxes):
            raise ValueError(
                f"{len(given)} {name}(s) for {len(boxes)} window(s); they are per "
                "window, so a static run repeats them per animal"
            )
    ih, iw = int(input_hw[0]), int(input_hw[1])
    windows = [
        window_affine(b, (ih, iw), m) for b, m in zip(boxes, flips, strict=True)
    ]
    # The mean scale host-side, once: a per-moment `det` would sync the device. A
    # mirror has unit scale, so it changes no view's pre-shrink.
    scales = [
        float(torch.sqrt(torch.abs(torch.det(a.matrix[:, :2].double()))))
        for a in windows
    ]
    shrink = [
        (max(1, round(h * f)), max(1, round(w * f)))
        if f < ANTIALIAS_MAX_SCALE
        else (h, w)
        for (h, w), f in zip(sizes, scales, strict=True)
    ]
    canvas = (max(s[0] for s in shrink), max(s[1] for s in shrink))
    # Each view samples the common canvas, so its theta is normalized against the
    # canvas; in pixel coordinates the map is the same one `warp_image` would build.
    theta = torch.stack(
        [
            (
                win @ Affine2D.resize(size, shr).inverse() if shr != size else win
            ).grid_theta(canvas, (ih, iw))
            for win, size, shr in zip(windows, sizes, shrink, strict=True)
        ]
    ).to(device=device, dtype=torch.float32)
    grid = F.affine_grid(theta, (len(boxes), 1, ih, iw), align_corners=False)
    inv = [w.inverse().matrix for w in windows]
    return cls(
        input_hw=(ih, iw),
        frame_hw=tuple(sizes),
        source_of=tuple(source_of),
        shrink_hw=tuple(shrink),
        canvas_hw=canvas,
        windows=tuple(windows),
        grid=grid,
        inv_linear=torch.stack([m[:, :2] for m in inv]).to(device),
        inv_offset=torch.stack([m[:, 2] for m in inv]).to(device),
        px_scale=torch.tensor([float(iw), float(ih)], device=device),
    )

fill

fill(canvas: Tensor, frames) -> torch.Tensor

Write each window's (T, H, W) or (H, W) source frame into canvas.

frames is indexed by source_of: host uint8 arrays, or a stacked on-device tensor (the online path's buffer, indexed without a copy). Integer input is scaled by 1/255, floats are taken as normalized; a shared source is uploaded, and pre-shrunk, once. Returns canvas.

Source code in src/deeperfly/pose2d/windows.py

def fill(self, canvas: torch.Tensor, frames) -> torch.Tensor:
    """Write each window's `(T, H, W)` or `(H, W)` source frame into `canvas`.

    `frames` is indexed by `source_of`: host uint8 arrays, or a stacked on-device
    tensor (the online path's buffer, indexed without a copy). Integer input is
    scaled by 1/255, floats are taken as normalized; a shared source is uploaded,
    and pre-shrunk, once. Returns `canvas`.
    """
    ready: dict[tuple[int, int, int], torch.Tensor] = {}
    for v, s in enumerate(self.source_of):
        sh, sw = self.shrink_hw[v]
        img = ready.get((s, sh, sw))
        if img is None:
            frame = frames[s]
            img = (
                frame
                if torch.is_tensor(frame)
                else torch.from_numpy(np.asarray(frame))
            )
            img = img.to(device=canvas.device, non_blocking=True)
            img = (
                img.float().div_(255.0)
                if not img.is_floating_point()
                else img.float()
            )
            if img.ndim == 2:
                img = img[None]
            if (sh, sw) != self.frame_hw[v]:
                img = F.interpolate(
                    img[:, None],
                    size=(sh, sw),
                    mode="bilinear",
                    align_corners=False,
                    antialias=True,
                )[:, 0]
            ready[(s, sh, sw)] = img
        canvas[v, : img.shape[0], :sh, :sw] = img
    return canvas

new_canvas

new_canvas(n_frames: int = 1) -> torch.Tensor

A zero (V, T, H, W) canvas for n_frames moments, on the grid's device.

It must stay zero outside each view's shrunk frame, which is what makes the common canvas identical to per-view buffers; reusing one only overwrites the written region.

Source code in src/deeperfly/pose2d/windows.py

def new_canvas(self, n_frames: int = 1) -> torch.Tensor:
    """A zero `(V, T, H, W)` canvas for `n_frames` moments, on the grid's device.

    It must stay zero outside each view's shrunk frame, which is what makes the
    common canvas identical to per-view buffers; reusing one only overwrites the
    written region.
    """
    return torch.zeros(
        (self.n_views, int(n_frames), *self.canvas_hw),
        device=self.grid.device,
        dtype=torch.float32,
    )

sample

sample(canvas: Tensor) -> torch.Tensor

(V, T, H, W) canvas -> the (T, V, 1, h, w) network input, in one kernel.

T rides in grid_sample's channel axis, so one grid of V rows covers a whole block of moments.

Source code in src/deeperfly/pose2d/windows.py

def sample(self, canvas: torch.Tensor) -> torch.Tensor:
    """`(V, T, H, W)` canvas -> the `(T, V, 1, h, w)` network input, in one kernel.

    `T` rides in `grid_sample`'s channel axis, so one grid of `V` rows covers a
    whole block of moments.
    """
    out = F.grid_sample(
        canvas,
        self.grid,
        mode="bilinear",
        padding_mode="zeros",
        align_corners=False,
    )
    return out.permute(1, 0, 2, 3).unsqueeze(2)

to_video_px

to_video_px(pts_norm: Tensor) -> torch.Tensor

(..., V, N, 2) input-normalized points -> video pixels, all views at once.

N is whatever the caller stacks per view: keypoints, or keypoints times candidates.

Source code in src/deeperfly/pose2d/windows.py

def to_video_px(self, pts_norm: torch.Tensor) -> torch.Tensor:
    """`(..., V, N, 2)` input-normalized points -> video pixels, all views at once.

    `N` is whatever the caller stacks per view: keypoints, or keypoints times
    candidates.
    """
    px = pts_norm * self.px_scale
    return (
        torch.einsum("...vkj,vij->...vki", px, self.inv_linear.to(px.dtype))
        + self.inv_offset.to(px.dtype)[:, None, :]
    )

warp

warp(frames) -> torch.Tensor

The sources' (T, H, W) uint8 frames -> the (T, V, 1, h, w) input.

Source code in src/deeperfly/pose2d/windows.py

def warp(self, frames) -> torch.Tensor:
    """The sources' `(T, H, W)` uint8 frames -> the `(T, V, 1, h, w)` input."""
    first = frames[0] if hasattr(frames, "__getitem__") else next(iter(frames))
    n = 1 if np.ndim(first) == 2 else int(np.shape(first)[0])
    return self.sample(self.fill(self.new_canvas(n), frames))

WindowPlanner

One block's coarse pass -> the windows its fine pass runs through.

A planner declares which output of the pose net's coarse pass it reads, the joints or the center peaks, and how many animals it plans; one whose coarse pass is another net overrides coarse.

Attributes:

Name	Type	Description
animals	int	How many slots plan returns.
n_centers	int	Center peaks per view the coarse pass decodes; 0 asks for the joints instead.
letterbox	object	Per view, the whole-frame window the coarse pass runs through.

Source code in src/deeperfly/pose2d/windows.py

class WindowPlanner:
    """One block's coarse pass -> the windows its fine pass runs through.

    A planner declares which output of the pose net's coarse pass it reads, the joints
    or the center peaks, and how many animals it plans; one whose coarse pass is
    another net overrides `coarse`.

    Attributes:
        animals: How many slots `plan` returns.
        n_centers: Center peaks per view the coarse pass decodes; 0 asks for the joints
            instead.
        letterbox: Per view, the whole-frame window the coarse pass runs through.
    """

    animals: int = 1
    n_centers: int = 0
    letterbox: object

    def coarse(self, blocks):
        """A planner's own coarse pass over one block, or `None` for the pose net's.

        The pose net's pass runs over the `letterbox` windows and hands `plan` the
        video-pixel joints and conf, or the center peaks at `n_centers > 0`. A planner
        with a net of its own returns `(points, conf, centers)` itself.
        """
        return

    def reset(self) -> None:
        """Forget what this planner carries from one block to the next."""

    def plan(self, points, conf, centers):
        """One block's coarse output -> its windows, its animals and their flags.

        `points` `(T, V, K, 2)` and `conf` `(T, V, K)` are video-pixel joints, `None` at
        `n_centers > 0`; `centers` `(T, V, C, 3)` are video-pixel peaks as
        `(x, y, score)`, `None` at `n_centers == 0`.

        Returns:
            `(T, N, V, 4)` windows, `(T, N, 3)` world-frame centers or `None`, and
            `(T, N)` uint8 track flags.
        """
        raise NotImplementedError

coarse

coarse(blocks)

A planner's own coarse pass over one block, or None for the pose net's.

The pose net's pass runs over the letterbox windows and hands plan the video-pixel joints and conf, or the center peaks at n_centers > 0. A planner with a net of its own returns (points, conf, centers) itself.

Source code in src/deeperfly/pose2d/windows.py

def coarse(self, blocks):
    """A planner's own coarse pass over one block, or `None` for the pose net's.

    The pose net's pass runs over the `letterbox` windows and hands `plan` the
    video-pixel joints and conf, or the center peaks at `n_centers > 0`. A planner
    with a net of its own returns `(points, conf, centers)` itself.
    """
    return

plan

plan(points, conf, centers)

One block's coarse output -> its windows, its animals and their flags.

points (T, V, K, 2) and conf (T, V, K) are video-pixel joints, None at n_centers > 0; centers (T, V, C, 3) are video-pixel peaks as (x, y, score), None at n_centers == 0.

Returns:

Type	Description
	(T, N, V, 4) windows, (T, N, 3) world-frame centers or None, and
	(T, N) uint8 track flags.

Source code in src/deeperfly/pose2d/windows.py

def plan(self, points, conf, centers):
    """One block's coarse output -> its windows, its animals and their flags.

    `points` `(T, V, K, 2)` and `conf` `(T, V, K)` are video-pixel joints, `None` at
    `n_centers > 0`; `centers` `(T, V, C, 3)` are video-pixel peaks as
    `(x, y, score)`, `None` at `n_centers == 0`.

    Returns:
        `(T, N, V, 4)` windows, `(T, N, 3)` world-frame centers or `None`, and
        `(T, N)` uint8 track flags.
    """
    raise NotImplementedError

reset

reset() -> None

Forget what this planner carries from one block to the next.

Source code in src/deeperfly/pose2d/windows.py

def reset(self) -> None:
    """Forget what this planner carries from one block to the next."""

aspect_window

aspect_window(
    hw: tuple[int, int], input_hw: tuple[int, int]
) -> tuple[float, float, float, float]

The largest centered window of an hw image at the input's aspect.

Source code in src/deeperfly/pose2d/windows.py

def aspect_window(
    hw: tuple[int, int], input_hw: tuple[int, int]
) -> tuple[float, float, float, float]:
    """The largest centered window of an `hw` image at the input's aspect."""
    h, w = hw
    aspect = input_hw[1] / input_hw[0]
    if w / h > aspect:
        cw, ch = h * aspect, float(h)
    else:
        cw, ch = float(w), w / aspect
    x0, y0 = (w - cw) / 2, (h - ch) / 2
    return (x0, y0, x0 + cw, y0 + ch)

follow_affines

follow_affines(
    boxes: Tensor,
    shrink_scale: Tensor,
    canvas_hw: Sequence[int],
    input_hw: Sequence[int],
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

(V, 4) windows -> affine_grid thetas and the inverse map, on their device.

What WindowGeometry.build computes host-side, per moment and without mirrors: shrink_scale (V, 2) is each view's [x, y] pre-shrink (shrunk size over frame size) and canvas_hw the canvas the shrunk frames sit in. Returns theta (V, 2, 3), input-normalized -> canvas- normalized, and scale, offset (V, 2), input pixels -> video pixels as px * scale + offset.

Source code in src/deeperfly/pose2d/windows.py

def follow_affines(
    boxes: torch.Tensor,
    shrink_scale: torch.Tensor,
    canvas_hw: Sequence[int],
    input_hw: Sequence[int],
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """`(V, 4)` windows -> `affine_grid` thetas and the inverse map, on their device.

    What [WindowGeometry.build][deeperfly.pose2d.windows.WindowGeometry.build] computes
    host-side, per moment and without mirrors: `shrink_scale` `(V, 2)` is each view's
    `[x, y]` pre-shrink (shrunk size over frame size) and `canvas_hw` the canvas the
    shrunk frames sit in. Returns `theta` `(V, 2, 3)`, input-normalized -> canvas-
    normalized, and `scale`, `offset` `(V, 2)`, input pixels -> video pixels as
    `px * scale + offset`.
    """
    ih, iw = float(input_hw[0]), float(input_hw[1])
    ch, cw = float(canvas_hw[0]), float(canvas_hw[1])
    x0, y0, x1, y1 = boxes.unbind(dim=-1)
    bw, bh = x1 - x0, y1 - y0
    rx, ry = shrink_scale[:, 0], shrink_scale[:, 1]
    zero = torch.zeros_like(bw)
    theta = torch.stack(
        [
            torch.stack([bw * rx / cw, zero, (bw + 2 * x0) * rx / cw - 1], dim=-1),
            torch.stack([zero, bh * ry / ch, (bh + 2 * y0) * ry / ch - 1], dim=-1),
        ],
        dim=-2,
    )
    scale = torch.stack([bw / iw, bh / ih], dim=-1)
    offset = torch.stack([x0 + 0.5 * bw / iw - 0.5, y0 + 0.5 * bh / ih - 0.5], dim=-1)
    return theta, scale, offset

follow_boxes

follow_boxes(
    points: Tensor,
    conf: Tensor,
    static: Tensor,
    *,
    have_prev: Tensor,
    force_static: Tensor,
    aspects: Tensor,
    min_w: float,
    tightness: float = 1.25,
    conf_min: float = 0.3,
    keep_frac: float = 0.85,
    floor: float = 0.8,
    min_points: int = 4,
) -> torch.Tensor

The last pose -> each view's next window, (V, 4), on the pose's device.

points (V, K, 2) and conf (V, K) are the previous moment's, in video pixels. The box is Locator's (the keep_frac most confident points at or above conf_min), grown to hold floor times the box of every finite point about its own center, then drawn at tightness at the view's aspect and at least min_w wide (ADR-0149). The floor is what stops a tight window from losing a tip for good: a tip it clips comes back unsure, and the confident box alone would leave it out again; floor = 0 is the confident box alone.

A view falls back to its static window, the re-acquisition, when it has fewer than min_points confident points, and every view does before the first pose (have_prev false) or on demand (force_static). Both flags are 0-d bool tensors, so nothing here waits on the device.

Source code in src/deeperfly/pose2d/windows.py

def follow_boxes(
    points: torch.Tensor,
    conf: torch.Tensor,
    static: torch.Tensor,
    *,
    have_prev: torch.Tensor,
    force_static: torch.Tensor,
    aspects: torch.Tensor,
    min_w: float,
    tightness: float = 1.25,
    conf_min: float = 0.3,
    keep_frac: float = 0.85,
    floor: float = 0.8,
    min_points: int = 4,
) -> torch.Tensor:
    """The last pose -> each view's next window, `(V, 4)`, on the pose's device.

    `points` `(V, K, 2)` and `conf` `(V, K)` are the previous moment's, in video
    pixels. The box is [Locator][deeperfly.pose2d.windows.Locator]'s (the `keep_frac`
    most confident points at or above `conf_min`), grown to hold `floor` times the box
    of every finite point about its own center, then drawn at `tightness` at the view's
    aspect and at least `min_w` wide (ADR-0149). The floor is what stops a tight window
    from losing a tip for good: a tip it clips comes back unsure, and the confident box
    alone would leave it out again; `floor = 0` is the confident box alone.

    A view falls back to its `static` window, the re-acquisition, when it has fewer
    than `min_points` confident points, and every view does before the first pose
    (`have_prev` false) or on demand (`force_static`). Both flags are 0-d bool tensors,
    so nothing here waits on the device.
    """
    finite = torch.isfinite(points).all(dim=-1)
    ok = finite & (conf >= conf_min)
    nan = torch.full_like(conf, float("nan"))
    thr = torch.nanquantile(torch.where(ok, conf, nan), 1.0 - keep_frac, dim=-1)
    sel = (ok & (conf >= thr[:, None]))[..., None]
    inf = torch.full_like(points, float("inf"))
    lo = torch.where(sel, points, inf).amin(dim=1)
    hi = torch.where(sel, points, -inf).amax(dim=1)
    if floor > 0:
        every = finite[..., None]
        lo_a = torch.where(every, points, inf).amin(dim=1)
        hi_a = torch.where(every, points, -inf).amax(dim=1)
        mid, half = (lo_a + hi_a) / 2, (hi_a - lo_a) / 2 * floor
        lo = torch.minimum(lo, mid - half)
        hi = torch.maximum(hi, mid + half)
    center, extent = (lo + hi) / 2, hi - lo
    w, h = extent[:, 0], extent[:, 1]
    bw = torch.clamp_min(
        torch.clamp_min(torch.maximum(w, h * aspects), 1.0) * tightness, min_w
    )
    bh = bw / aspects
    box = torch.stack(
        [
            center[:, 0] - bw / 2,
            center[:, 1] - bh / 2,
            center[:, 0] + bw / 2,
            center[:, 1] + bh / 2,
        ],
        dim=-1,
    )
    fall_back = ~have_prev | force_static | (ok.sum(dim=-1) < min_points)
    return torch.where(fall_back[:, None], static.to(box.dtype), box)

mark_others

mark_others(
    img: Tensor,
    centers,
    own: int | None,
    *,
    jitter: float = 0.0,
)

Draw a white disk in a black ring at every center but own's, in input pixels.

Two animals on top of each other give two nearly identical windows, and nothing in either says which animal is wanted: both skeletons land on one of them. The mark is that missing word (ADR-0148). img is (..., H, W) in [0, 1], centers (A, 2), NaN where an animal is not seen; jitter (pixels, training) stands for the error of the center a window is placed on. Inference draws the same mark at the other windows' centers, so the two stay one function.

Source code in src/deeperfly/pose2d/windows.py

def mark_others(img: torch.Tensor, centers, own: int | None, *, jitter: float = 0.0):
    """Draw a white disk in a black ring at every center but `own`'s, in input pixels.

    Two animals on top of each other give two nearly identical windows, and nothing
    in either says which animal is wanted: both skeletons land on one of them. The
    mark is that missing word (ADR-0148). `img` is `(..., H, W)` in `[0, 1]`,
    `centers` `(A, 2)`, NaN where an animal is not seen; `jitter` (pixels, training)
    stands for the error of the center a window is placed on. Inference draws the
    same mark at the other windows' centers, so the two stay one function.
    """
    h, w = img.shape[-2:]
    yy, xx = torch.meshgrid(
        torch.arange(h, dtype=torch.float32, device=img.device),
        torch.arange(w, dtype=torch.float32, device=img.device),
        indexing="ij",
    )
    for a, (cx, cy) in enumerate(torch.as_tensor(centers).tolist()):
        if a == own or not (math.isfinite(cx) and math.isfinite(cy)):
            continue
        if jitter:
            cx += float(torch.randn(())) * jitter
            cy += float(torch.randn(())) * jitter
        d = torch.sqrt((xx - cx) ** 2 + (yy - cy) ** 2)
        img = torch.where(d <= MARK_RING_PX, torch.zeros_like(img), img)
        img = torch.where(d <= MARK_DISK_PX, torch.ones_like(img), img)
    return img

padded_window

padded_window(
    hw: tuple[int, int], input_hw: tuple[int, int]
) -> tuple[float, float, float, float]

The smallest centered box at the input's aspect holding an hw frame.

The letterbox: the whole frame, the rest black, which an inference given no crop feeds the network instead of a squashed frame.

Source code in src/deeperfly/pose2d/windows.py

def padded_window(
    hw: tuple[int, int], input_hw: tuple[int, int]
) -> tuple[float, float, float, float]:
    """The smallest centered box at the input's aspect holding an `hw` frame.

    The letterbox: the whole frame, the rest black, which an inference given no crop
    feeds the network instead of a squashed frame.
    """
    h, w = hw
    aspect = input_hw[1] / input_hw[0]
    if w / h > aspect:
        cw, ch = float(w), w / aspect
    else:
        cw, ch = h * aspect, float(h)
    x0, y0 = (w - cw) / 2, (h - ch) / 2
    return (x0, y0, x0 + cw, y0 + ch)

window_affine

window_affine(
    box: Sequence[float],
    input_hw: Sequence[int],
    mirrored: bool,
) -> Affine2D

Frame pixels -> input pixels for one window: the crop, mirrored where it is.

The one definition of what the detector sees, in training and inference. box is (x0, y0, x1, y1) in the frame's pixels; a mirrored view's crop comes out reflected (ADR-0053), which is a flip of the input, so the frame's width never enters.

Source code in src/deeperfly/pose2d/windows.py

def window_affine(
    box: Sequence[float], input_hw: Sequence[int], mirrored: bool
) -> Affine2D:
    """Frame pixels -> input pixels for one window: the crop, mirrored where it is.

    The one definition of what the detector sees, in training and inference. `box` is
    `(x0, y0, x1, y1)` in the frame's pixels; a mirrored view's crop comes out
    reflected (ADR-0053), which is a flip of the input, so the frame's width never
    enters.
    """
    ih, iw = int(input_hw[0]), int(input_hw[1])
    crop = Affine2D.crop(tuple(float(v) for v in box), (ih, iw))
    return Affine2D.flip_x(float(iw)) @ crop if mirrored else crop

One block of frames of every video through the detector, static or planned.

Videos yields, per block of frames, every source's luma in lockstep, so a moment's views arrive together and a frame carrying three views is decoded once. block_pass is one detector pass over a block, through fixed windows or a planner's; run_models collapses an ensemble and averages flip TTA at the coordinate level (the field grid is not mirror-symmetric). Everything comes back in video pixels.

BlockOut dataclass

One batch of frames through the detector, in video pixels, on the host.

Every per-animal array carries the animal axis behind the frame; N is 1 at one animal, squeezed off when the recording's arrays are assembled.

Attributes:

Name	Type	Description
points	Tensor	(T, N, V, K, 2).
conf	Tensor	(T, N, V, K).
candidates	Tensor | None	(T, N, V, K, C, 2), or None without candidates.
candidate_scores	Tensor | None	(T, N, V, K, C), or None.
windows	ndarray | None	(T, N, V, 4) the planner's boxes, or None in static mode.
flags	ndarray	(T, N) uint8 track flags, all TRACKED where nothing can be lost.
centers	ndarray | None	(T, N, 3) each animal's tracked center, or None from a planner that reads the joints.
center_candidates	ndarray | None	(T, V, C, 3) the center peaks the planner read, or None.

Source code in src/deeperfly/pose2d/blocks.py

@dataclass(frozen=True)
class BlockOut:
    """One batch of frames through the detector, in video pixels, on the host.

    Every per-animal array carries the animal axis behind the frame; `N` is 1 at one
    animal, squeezed off when the recording's arrays are assembled.

    Attributes:
        points: `(T, N, V, K, 2)`.
        conf: `(T, N, V, K)`.
        candidates: `(T, N, V, K, C, 2)`, or `None` without candidates.
        candidate_scores: `(T, N, V, K, C)`, or `None`.
        windows: `(T, N, V, 4)` the planner's boxes, or `None` in static mode.
        flags: `(T, N)` uint8 track flags, all `TRACKED` where nothing can be lost.
        centers: `(T, N, 3)` each animal's tracked center, or `None` from a planner
            that reads the joints.
        center_candidates: `(T, V, C, 3)` the center peaks the planner read, or
            `None`.
    """

    points: torch.Tensor
    conf: torch.Tensor
    candidates: torch.Tensor | None
    candidate_scores: torch.Tensor | None
    windows: np.ndarray | None
    flags: np.ndarray
    centers: np.ndarray | None = None
    center_candidates: np.ndarray | None = None

Videos

Open readers over every video of the recording; frames come out per video.

Views that share a video read the same block, each through its own mirror, so a frame carrying three views is decoded and uploaded once. frames(indices) is random access and stream_blocks the forward pass.

Source code in src/deeperfly/pose2d/blocks.py

class Videos:
    """Open readers over every video of the recording; frames come out per video.

    Views that share a video read the same block, each through its own mirror, so a
    frame carrying three views is decoded and uploaded once. `frames(indices)` is random
    access and `stream_blocks` the forward pass.
    """

    def __init__(
        self, recording: Recording, views: Sequence[str], layout: Layout
    ) -> None:
        """Open one reader per video `views` read; `close` releases them.

        `layout` says which video carries each view and which views are mirrored:
        everything a detector needs of the rig, and nothing of its calibration.
        """
        self.views = list(views)
        # The videos these views need, deduplicated: six hexaview views open two.
        names = layout.videos_of(self.views)
        self.videos = [recording.videos[name] for name in names]
        self.source_of = [names.index(layout.video_of(v)) for v in self.views]
        self.readers = [video.open() for video in self.videos]
        counts = [r.count() for r in self.readers]
        if any(c is None for c in counts):
            raise ValueError(f"cannot determine the frame count of {recording.name!r}")
        self.n_frames = min(int(c) for c in counts)
        self.source_hw = [
            tuple(int(v) for v in np.asarray(r[0]).shape[:2]) for r in self.readers
        ]
        # A view's frame is its video's frame: nothing is cut out and nothing is
        # flipped, so the two sizes are the same one.
        self.frame_hw = [self.source_hw[s] for s in self.source_of]
        self.mirrors = [layout.is_mirrored(name) for name in self.views]

    @staticmethod
    def _luma(block: np.ndarray) -> np.ndarray:
        """One source's `(T, H, W)` luma out of its `(T, H, W, C)` block."""
        arr = np.asarray(block)
        return np.ascontiguousarray(arr[..., 0] if arr.ndim == 4 else arr)

    def window(self, view: int, box, input_hw) -> Affine2D:
        """Video pixels -> input pixels for one window of one view, mirror included."""
        return window_affine(box, input_hw, self.mirrors[view])

    def geometry(self, boxes, input_hw, *, device) -> WindowGeometry:
        """The warp for `boxes`: one per view, or `(N, V)` with the views inside."""
        n = len(boxes) // len(self.views)
        return WindowGeometry.build(
            boxes,
            [self.source_hw[s] for s in self.source_of] * n,
            input_hw,
            device=device,
            sources=self.source_of * n,
            views=list(self.mirrors) * n,
        )

    def frames(self, indices: Sequence[int]) -> list[np.ndarray]:
        """Per source, the `(T, H, W)` uint8 luma of the given frame indices."""
        return [self._luma(r[list(indices)]) for r in self.readers]

    def stream_blocks(
        self,
        *,
        block_size: int,
        thread_count: int = 2,
        start: int = 0,
        stop: int | None = None,
    ):
        """Every source's `(T, H, W)` luma per block, in lockstep: one decode each."""
        streams = [
            r.stream_blocks(
                start=start,
                stop=self.n_frames if stop is None else min(int(stop), self.n_frames),
                block_size=block_size,
                gray_ok=True,
                thread_count=thread_count,
            )
            for r in self.readers
        ]
        for blocks in zip(*streams, strict=True):
            yield [self._luma(b) for b in blocks]

    def close(self) -> None:
        """Release every reader."""
        for r in self.readers:
            r.close()

__init__

__init__(
    recording: Recording,
    views: Sequence[str],
    layout: Layout,
) -> None

Open one reader per video views read; close releases them.

layout says which video carries each view and which views are mirrored: everything a detector needs of the rig, and nothing of its calibration.

Source code in src/deeperfly/pose2d/blocks.py

def __init__(
    self, recording: Recording, views: Sequence[str], layout: Layout
) -> None:
    """Open one reader per video `views` read; `close` releases them.

    `layout` says which video carries each view and which views are mirrored:
    everything a detector needs of the rig, and nothing of its calibration.
    """
    self.views = list(views)
    # The videos these views need, deduplicated: six hexaview views open two.
    names = layout.videos_of(self.views)
    self.videos = [recording.videos[name] for name in names]
    self.source_of = [names.index(layout.video_of(v)) for v in self.views]
    self.readers = [video.open() for video in self.videos]
    counts = [r.count() for r in self.readers]
    if any(c is None for c in counts):
        raise ValueError(f"cannot determine the frame count of {recording.name!r}")
    self.n_frames = min(int(c) for c in counts)
    self.source_hw = [
        tuple(int(v) for v in np.asarray(r[0]).shape[:2]) for r in self.readers
    ]
    # A view's frame is its video's frame: nothing is cut out and nothing is
    # flipped, so the two sizes are the same one.
    self.frame_hw = [self.source_hw[s] for s in self.source_of]
    self.mirrors = [layout.is_mirrored(name) for name in self.views]

close

close() -> None

Release every reader.

Source code in src/deeperfly/pose2d/blocks.py

def close(self) -> None:
    """Release every reader."""
    for r in self.readers:
        r.close()

frames

frames(indices: Sequence[int]) -> list[np.ndarray]

Per source, the (T, H, W) uint8 luma of the given frame indices.

Source code in src/deeperfly/pose2d/blocks.py

def frames(self, indices: Sequence[int]) -> list[np.ndarray]:
    """Per source, the `(T, H, W)` uint8 luma of the given frame indices."""
    return [self._luma(r[list(indices)]) for r in self.readers]

geometry

geometry(boxes, input_hw, *, device) -> WindowGeometry

The warp for boxes: one per view, or (N, V) with the views inside.

Source code in src/deeperfly/pose2d/blocks.py

def geometry(self, boxes, input_hw, *, device) -> WindowGeometry:
    """The warp for `boxes`: one per view, or `(N, V)` with the views inside."""
    n = len(boxes) // len(self.views)
    return WindowGeometry.build(
        boxes,
        [self.source_hw[s] for s in self.source_of] * n,
        input_hw,
        device=device,
        sources=self.source_of * n,
        views=list(self.mirrors) * n,
    )

stream_blocks

stream_blocks(
    *,
    block_size: int,
    thread_count: int = 2,
    start: int = 0,
    stop: int | None = None,
)

Every source's (T, H, W) luma per block, in lockstep: one decode each.

Source code in src/deeperfly/pose2d/blocks.py

def stream_blocks(
    self,
    *,
    block_size: int,
    thread_count: int = 2,
    start: int = 0,
    stop: int | None = None,
):
    """Every source's `(T, H, W)` luma per block, in lockstep: one decode each."""
    streams = [
        r.stream_blocks(
            start=start,
            stop=self.n_frames if stop is None else min(int(stop), self.n_frames),
            block_size=block_size,
            gray_ok=True,
            thread_count=thread_count,
        )
        for r in self.readers
    ]
    for blocks in zip(*streams, strict=True):
        yield [self._luma(b) for b in blocks]

window

window(view: int, box, input_hw) -> Affine2D

Video pixels -> input pixels for one window of one view, mirror included.

Source code in src/deeperfly/pose2d/blocks.py

def window(self, view: int, box, input_hw) -> Affine2D:
    """Video pixels -> input pixels for one window of one view, mirror included."""
    return window_affine(box, input_hw, self.mirrors[view])

block_pass

block_pass(
    models,
    blocks: Sequence[ndarray],
    videos: Videos,
    geom: WindowGeometry,
    *,
    planner: WindowPlanner | None,
    device,
    flip_perm,
    combine: str,
    candidates: int,
    apart: bool = False,
) -> BlockOut

Per-source (T, H, W) luma blocks -> one detector pass, static or planned.

Static (no planner), geom holds the run's fixed N * V windows and the block goes through them in one batched warp. With a planner, its coarse pass gives one box per (frame, animal, view) and the fine pass warps per cell; a slot the planner lost is skipped and comes back NaN at confidence 0, which on single-fly footage run at two animals saves half the work. Every point comes back in video pixels on the pixel-center grid, whatever grid the detector was trained on (ADR-0078). apart says the block's frames are stretches of their own (ADR-0082): the planner is reset before each, so nothing carries from one to the next.

Source code in src/deeperfly/pose2d/blocks.py

def block_pass(
    models,
    blocks: Sequence[np.ndarray],
    videos: Videos,
    geom: WindowGeometry,
    *,
    planner: WindowPlanner | None,
    device,
    flip_perm,
    combine: str,
    candidates: int,
    apart: bool = False,
) -> BlockOut:
    """Per-source `(T, H, W)` luma blocks -> one detector pass, static or planned.

    Static (no `planner`), `geom` holds the run's fixed `N * V` windows and the block
    goes through them in one batched warp. With a planner, its `coarse` pass gives one
    box per `(frame, animal, view)` and the fine pass warps per cell; a slot the
    planner lost is skipped and comes back NaN at confidence 0, which on single-fly
    footage run at two animals saves half the work. Every point comes back in video
    pixels on the pixel-center grid, whatever grid the detector was trained on
    (ADR-0078). `apart` says the block's frames are stretches of their own (ADR-0082):
    the planner is reset before each, so nothing carries from one to the next.
    """
    n_v = len(videos.views)
    n_t = int(np.shape(blocks[0])[0])
    input_hw = geom.input_hw
    windows = centers = center_cands = missing = None
    if planner is None:
        n_a = geom.n_views // n_v
        flags = np.full((n_t, n_a), TRACKED, dtype=np.uint8)
        frame_affines = None
        # Every animal's windows are in one geometry, so the same blocks serve them.
        images = _split_animals(geom.warp(blocks), n_a)
    else:
        n_a = planner.animals
        found = planner.coarse(blocks)
        if found is None:
            found = _coarse_pass(
                models, blocks, geom, combine=combine, centers=planner.n_centers
            )
        if apart:
            windows, centers, flags, missing = _plan_apart(planner, found, n_t)
        else:
            windows, centers, flags = planner.plan(*found)
            missing = getattr(planner, "missing", None)
        center_cands = found[2]
        frame_affines = [
            [
                videos.window(v, tuple(float(x) for x in windows[t, n, v]), input_hw)
                for v in range(n_v)
            ]
            for t in range(n_t)
            for n in range(n_a)
            if flags[t, n] != LOST
        ]
        images = None
    # Only the moments with an animal in them: the cells of the (frame, animal) grid
    # whose slot is not lost, in frame-major order.
    live = np.flatnonzero((flags != LOST).reshape(-1))
    if frame_affines is not None and live.size:
        images = _warp_frames(
            blocks, frame_affines, input_hw, device, live // n_a, videos
        )
        if n_a > 1 and marks_others(models):
            images = _mark_other_slots(images, windows, flags, live, n_a, frame_affines)
    back = scores = None
    if images is not None and frame_affines is None:
        with torch.no_grad():
            pts, conf, cand, _ = run_models(
                models, images, flip_perm, combine, candidates
            )
        foot, conf, back, scores = static_video_px(models, geom, n_a, pts, conf, cand)
        foot, conf = foot.cpu(), conf.cpu()
        if back is not None:
            back, scores = back.cpu(), scores.cpu()
    elif images is None:
        # Every slot of the block is lost: nothing to run, and the grid below is
        # filled in whole.
        n_k = first_model(models).n_points
        foot, conf = torch.empty(0, n_v, n_k, 2), torch.empty(0, n_v, n_k)
        if candidates:
            back = torch.empty(0, n_v, n_k, candidates, 2)
            scores = torch.empty(0, n_v, n_k, candidates)
    else:
        with torch.no_grad():
            pts, conf, cand, _ = run_models(
                models, images, flip_perm, combine, candidates
            )
        shift = grid_shift(models)
        foot = _to_video_px_frames(pts, frame_affines, input_hw) + shift
        if cand is not None:
            cxy, csc = cand
            back = _to_video_px_frames(cxy, frame_affines, input_hw)
            back, scores = (back + shift).cpu(), csc.cpu()
        foot, conf = foot.cpu(), conf.cpu()
    # Back onto the full (frame, animal) grid: a lost slot is NaN at confidence 0.
    shape = (n_t, n_a)
    foot = _scatter(foot, live, n_t * n_a, torch.nan).unflatten(0, shape)
    conf = _scatter(conf, live, n_t * n_a, 0.0).unflatten(0, shape)
    if missing is not None and missing.any():
        # A view windowed on another view's image (see FullFrameLocator) has no
        # observation of its own in that frame.
        mask = torch.as_tensor(missing)[:, None, :]  # (T, 1, V)
        foot = torch.where(mask[..., None, None], torch.nan, foot)
        conf = torch.where(mask[..., None], 0.0, conf)
    if back is not None:
        back = _scatter(back, live, n_t * n_a, torch.nan).unflatten(0, shape)
        scores = _scatter(scores, live, n_t * n_a, 0.0).unflatten(0, shape)
    return BlockOut(foot, conf, back, scores, windows, flags, centers, center_cands)

run_model

run_model(
    model,
    images: Tensor,
    flip_perm=None,
    candidates: int = 0,
    centers: int = 0,
)

(T, V, 1, h, w) -> points, conf, pools and centers of one member.

(T, V, K, 2) input-normalized points and (T, V, K) conf; the pools are this member's candidate peaks, one set per pass, empty when candidates is 0. centers asks for that many center-channel peaks, ((T, V, C, 2), (T, V, C)), read off the unmirrored pass only, or None.

Source code in src/deeperfly/pose2d/blocks.py

def run_model(
    model,
    images: torch.Tensor,
    flip_perm=None,
    candidates: int = 0,
    centers: int = 0,
):
    """`(T, V, 1, h, w)` -> points, conf, pools and centers of one member.

    `(T, V, K, 2)` input-normalized points and `(T, V, K)` conf; the pools are this
    member's candidate peaks, one set per pass, empty when `candidates` is 0.
    `centers` asks for that many center-channel peaks, `((T, V, C, 2), (T, V, C))`,
    read off the unmirrored pass only, or `None`.
    """
    if centers and not model.center_channel:
        raise ValueError(
            "this detector has no center channel to locate animals with; train one "
            "with `[data] center_channel = true`"
        )
    device = images.device
    pools = []
    center = None
    dtype = autocast_dtype(model, device)
    with torch.autocast(device.type, dtype=dtype, enabled=dtype is not None):
        out = model(images)
        t, v = images.shape[:2]
        # A center channel is the last of each view's block; every joint readout takes
        # the joint channels alone, so the single-animal path is the same numbers.
        c, k = model.num_classes, model.n_points
        by_view = out.reshape(t * v, c, *out.shape[-2:])
        field = by_view[:, :k]
        pts, conf = decode_points(model, field)
        pts, conf = pts.float(), conf.float()
        if candidates:
            pools.append(decode_candidates(model, field, candidates))
        if centers:
            cxy, csc = decode_candidates(
                model, by_view[:, k:], centers, radius_px=CANDIDATE_RADIUS_PX
            )
            center = (
                cxy.float().reshape(t, v, centers, 2),
                csc.float().reshape(t, v, centers),
            )
        if flip_perm is not None:
            out_m = model(images.flip(-1))
            field_m = out_m.reshape(t * v, c, *out_m.shape[-2:])[:, :k]
            pm, cm = decode_points(model, field_m)
            perm = flip_perm.to(device)
            w = model.input_hw[1]
            pm = _unmirror(pm.float().index_select(1, perm), w)
            pts = 0.5 * (pts + pm)
            conf = torch.minimum(conf, cm.float().index_select(1, perm))
            if candidates:
                cx, cs = decode_candidates(model, field_m, candidates)
                pools.append(
                    (_unmirror(cx.index_select(1, perm), w), cs.index_select(1, perm))
                )
    pools = [(x.reshape(t, v, k, -1, 2), s.reshape(t, v, k, -1)) for x, s in pools]
    return pts.reshape(t, v, k, 2), conf.reshape(t, v, k), pools, center

run_models

run_models(
    models,
    images: Tensor,
    flip_perm=None,
    combine: str = "mean",
    candidates: int = 0,
    centers: int = 0,
)

run_model over one module or an ensemble, collapsed by combine.

"mean" averages the members' coordinates and confidences; "median" takes the per-coordinate median, which one badly placed member cannot move. The candidates are every member's and pass's peaks pooled around the collapsed point (slot 0), and the centers are member 0's: an ensemble locates the animals with one member.

Source code in src/deeperfly/pose2d/blocks.py

def run_models(
    models,
    images: torch.Tensor,
    flip_perm=None,
    combine: str = "mean",
    candidates: int = 0,
    centers: int = 0,
):
    """`run_model` over one module or an ensemble, collapsed by `combine`.

    `"mean"` averages the members' coordinates and confidences; `"median"` takes the
    per-coordinate median, which one badly placed member cannot move. The candidates
    are every member's and pass's peaks pooled around the collapsed point (slot 0), and
    the centers are member 0's: an ensemble locates the animals with one member.
    """
    members = models if isinstance(models, (list, tuple)) else [models]
    outs = [
        run_model(m, images, flip_perm, candidates, centers if i == 0 else 0)
        for i, m in enumerate(members)
    ]
    center = outs[0][3]
    if len(outs) == 1:
        pts, conf = outs[0][0], outs[0][1]
    else:
        pts = torch.stack([o[0] for o in outs])
        conf = torch.stack([o[1] for o in outs])
        if combine == "median":
            pts, conf = pts.median(0).values, conf.median(0).values
        elif combine == "mean":
            pts, conf = pts.mean(0), conf.mean(0)
        else:
            raise ValueError(f"combine must be 'mean' or 'median', got {combine!r}")
    if not candidates:
        return pts, conf, None, center
    h, w = first_model(members).input_hw
    scale = pts.new_tensor([w, h], dtype=torch.float32)
    t, v, k = conf.shape
    pools = [
        ((x * scale).reshape(t * v * k, -1, 2), s.reshape(t * v * k, -1))
        for o in outs
        for x, s in o[2]
    ]
    cxy, csc = merge_candidates(
        (pts * scale).reshape(t * v * k, 2),
        conf.reshape(t * v * k),
        pools,
        candidates,
        radius=CANDIDATE_RADIUS_PX,
    )
    return (
        pts,
        conf,
        ((cxy / scale).reshape(t, v, k, -1, 2), csc.reshape(t, v, k, -1)),
        center,
    )

static_video_px

static_video_px(
    models, geom: WindowGeometry, n_a: int, pts, conf, cand
)

A static pass's readout -> video pixels, on the readout's device.

pts (T * N, V, K, 2) input-normalized, conf and cand as run_models returns them, through geom's fixed N * V windows. Returns (points, conf, candidates, candidate_scores), the last two None without candidates: what block_pass lays out, before the host copy.

Source code in src/deeperfly/pose2d/blocks.py

def static_video_px(models, geom: WindowGeometry, n_a: int, pts, conf, cand):
    """A static pass's readout -> video pixels, on the readout's device.

    `pts (T * N, V, K, 2)` input-normalized, `conf` and `cand` as `run_models` returns
    them, through `geom`'s fixed `N * V` windows. Returns `(points, conf, candidates,
    candidate_scores)`, the last two `None` without candidates: what `block_pass` lays
    out, before the host copy.
    """
    shift = grid_shift(models)
    foot = _split_animals(geom.to_video_px(_join_animals(pts, n_a)), n_a) + shift
    if cand is None:
        return foot, conf, None, None
    cxy, csc = cand
    # One batched inverse over the keypoints and their candidates folded together,
    # then unfolded again.
    n_m, n_v, n_k, n_c = csc.shape
    flat = _join_animals(cxy.reshape(n_m, n_v, n_k * n_c, 2), n_a)
    back = _split_animals(geom.to_video_px(flat), n_a).reshape(n_m, n_v, n_k, n_c, 2)
    return foot, conf, back + shift, csc

A run's static window per view, discovered with the detector.

The network sees a window around the animal, never a stretched frame (ADR-0046). A run has no labels to draw it from, so discover_crops finds it with the detector on sampled frames in three passes: the whole frame letterboxed, a 2x window around what that found, then the final window at 1.5x the animal's box. discover_static_crops is the one entry, which a run calls on the frames it was asked for and keeps what it found in its results (ADR-0082). Above one animal there is a window per animal, and a per-frame run measures one window size per view instead (discover_sizes).

crop_boxes

crop_boxes(
    crops, views: Sequence[str], animals: int
) -> np.ndarray

A crops mapping as the (N, V, 4) windows a static run predicts through.

A view's entry is a window per animal, or one every animal shares (a per-frame run's letterbox). A run config refuses a shared window for a project of more than one animal: one window over two tethered flies would let the detector pose either.

Source code in src/deeperfly/pose2d/crops.py

def crop_boxes(crops, views: Sequence[str], animals: int) -> np.ndarray:
    """A `crops` mapping as the `(N, V, 4)` windows a static run predicts through.

    A view's entry is a window per animal, or one every animal shares (a per-frame
    run's letterbox). A run config refuses a shared window for a project of more than
    one animal: one window over two tethered flies would let the detector pose either.
    """
    out = np.zeros((animals, len(views), 4), dtype=float)
    for v, view in enumerate(views):
        if view not in crops:
            raise ValueError(f"no window for view {view!r}")
        boxes = np.atleast_2d(np.asarray(crops[view], dtype=float))
        if boxes.shape[-1] != 4 or len(boxes) not in (1, animals):
            raise ValueError(
                f"view {view!r}: {len(boxes)} window(s) of {boxes.shape[-1]} "
                f"numbers, expected 1 or {animals} of 4"
            )
        out[:, v] = boxes
    return out

discover_crops

discover_crops(
    models,
    videos: Videos,
    *,
    device,
    n_frames: int = 24,
    tightness: float = 1.5,
    crop_aspect=None,
    flip_perm=None,
    batch_frames: int = 8,
    frames=None,
    log: Callable[[str], None] = print,
) -> dict[str, tuple[int, int, int, int]]

Find each view's window around the animal with the detector itself.

Three passes over n_frames evenly spaced frames of frames (every frame by default); a view whose animal cannot be located keeps the letterbox. tightness is 1.5, not the evaluation's 1.25: the window is one per run and the animal moves, and a joint leaving the window is not localized at all (15 of 7,323 held-out cells, 323 px mean error, 0.66 px on the overall mean). The model trained on tightness up to 2.2.

Source code in src/deeperfly/pose2d/crops.py

def discover_crops(
    models,
    videos: Videos,
    *,
    device,
    n_frames: int = 24,
    tightness: float = 1.5,
    crop_aspect=None,
    flip_perm=None,
    batch_frames: int = 8,
    frames=None,
    log: Callable[[str], None] = print,
) -> dict[str, tuple[int, int, int, int]]:
    """Find each view's window around the animal with the detector itself.

    Three passes over `n_frames` evenly spaced frames of `frames` (every frame by
    default); a view whose animal cannot be located keeps the letterbox. `tightness` is
    1.5, not the evaluation's 1.25: the window is one per run and the animal
    moves, and a joint leaving the window is not localized at all (15 of 7,323
    held-out cells, 323 px mean error, 0.66 px on the overall mean). The model trained
    on tightness up to 2.2.
    """
    input_hw = tuple(first_model(models).input_hw)
    aspects = _view_aspects(videos.views, input_hw, crop_aspect)
    min_w = input_hw[1] / 2
    indices = _sample_frames(videos, frames, n_frames)
    boxes = [padded_window(hw, input_hw) for hw in videos.frame_hw]
    for stage, t in (("letterbox", 2.0), ("2x", tightness)):
        pts, conf = _predict_windows(
            models,
            videos,
            indices,
            boxes,
            device=device,
            flip_perm=flip_perm,
            batch_frames=batch_frames,
        )
        new = []
        for v, view in enumerate(videos.views):
            box = _reduce_box(pts[:, v], conf[:, v], aspects[v], t, min_w)
            if box is None:
                log(f"  {view}: no confident detections after the {stage} pass")
                box = boxes[v]
            new.append(box)
        boxes = new
    return {
        view: tuple(round(v) for v in box)
        for view, box in zip(videos.views, boxes, strict=True)
    }

discover_sizes

discover_sizes(
    models,
    videos: Videos,
    tracker: MeasuredPlanner,
    *,
    device,
    combine: str = "mean",
    batch_frames: int = 8,
    n_frames: int = 24,
    tightness: float = 1.5,
    frames=None,
    log: Callable[[str], None] = print,
) -> np.ndarray

Measure the animals on sampled frames; one window size per view.

discover_crops for animals that walk: every animal of every sample goes into one _reduce_box per view, and only the size comes out, stored on the tracker too. A view whose animals were never found keeps its provisional window. The samples are drawn from frames, every frame by default.

Source code in src/deeperfly/pose2d/crops.py

def discover_sizes(
    models,
    videos: Videos,
    tracker: MeasuredPlanner,
    *,
    device,
    combine: str = "mean",
    batch_frames: int = 8,
    n_frames: int = 24,
    tightness: float = 1.5,
    frames=None,
    log: Callable[[str], None] = print,
) -> np.ndarray:
    """Measure the animals on sampled frames; one window size per view.

    `discover_crops` for animals that walk: every animal of every sample goes into one
    `_reduce_box` per view, and only the size comes out, stored on the tracker too. A
    view whose animals were never found keeps its provisional window. The samples are
    drawn from `frames`, every frame by default.
    """
    pts, conf, _ = _sample_planned(
        models,
        videos,
        tracker,
        device=device,
        combine=combine,
        batch_frames=batch_frames,
        indices=_sample_frames(videos, frames, n_frames),
    )
    # Every animal of every sample is one measurement of the same size.
    seen = pts.reshape(-1, *pts.shape[2:])  # (S * N, V, K, 2)
    scored = conf.reshape(-1, *conf.shape[2:])
    sizes = tracker.sizes.copy()
    for v, view in enumerate(videos.views):
        box = _reduce_box(
            seen[:, v],
            scored[:, v],
            tracker.aspects[v],
            tightness,
            tracker.min_w,
            center=(0.0, 0.0),
        )
        if box is None:
            log(f"  {view}: no animal located; keeping the provisional window")
            continue
        sizes[v] = (box[2] - box[0], box[3] - box[1])
    tracker.sizes = sizes
    tracker.reset()
    return sizes

discover_static_crops

discover_static_crops(
    models,
    videos: Videos,
    *,
    device,
    options: PredictOptions | None = None,
    frames=None,
    rig: Rig | None = None,
    log: Callable[[str], None] = print,
) -> np.ndarray

A static run's crop: one rectangle per animal and view, (N, V, 4).

Found with the detector on frames sampled from frames, the ones the run was asked for (every frame by default), and rounded to whole pixels; a run discovers its own and keeps it in its results (ADR-0082). Above one animal each view gets a rectangle per animal, measured through a Tracker on rig, the one reader of a calibration here.

Raises:

Type	Description
ValueError	If an animal turns out to walk (see _discover_static_windows), or several animals are to be measured without rig.

Source code in src/deeperfly/pose2d/crops.py

def discover_static_crops(
    models,
    videos: Videos,
    *,
    device,
    options: PredictOptions | None = None,
    frames=None,
    rig: Rig | None = None,
    log: Callable[[str], None] = print,
) -> np.ndarray:
    """A static run's crop: one rectangle per animal and view, `(N, V, 4)`.

    Found with the detector on frames sampled from `frames`, the ones the run was asked
    for (every frame by default), and rounded to whole pixels; a run discovers its own
    and keeps it in its results (ADR-0082). Above one animal each view gets a rectangle
    per animal, measured through a `Tracker` on `rig`, the one reader of a calibration
    here.

    Raises:
        ValueError: If an animal turns out to walk (see `_discover_static_windows`),
            or several animals are to be measured without `rig`.
    """
    options = options or PredictOptions()
    if options.animals > 1:
        if rig is None:
            raise ValueError(
                "measuring several animals triangulates their centers, so it needs "
                "the calibrated rig"
            )
        found = _discover_static_windows(
            models,
            videos,
            Tracker(
                rig,
                videos,
                tuple(first_model(models).input_hw),
                animals=options.animals,
                tracking=options.tracking,
                crop_aspect=options.crop_aspect,
            ),
            device=device,
            combine=options.combine,
            batch_frames=options.batch_frames,
            tightness=options.tightness,
            frames=frames,
            log=log,
        )
    else:
        window = discover_crops(
            models,
            videos,
            device=device,
            crop_aspect=options.crop_aspect,
            flip_perm=options.flip_perm,
            batch_frames=options.batch_frames,
            tightness=options.tightness,
            frames=frames,
            log=log,
        )
        found = np.array([[window[view] for view in videos.views]], dtype=float)
    return np.round(found)

Run a multiview detector over a project's recording, whole or by chosen frames.

predict_recording streams a frame set through the model, one synchronized group per moment, through fixed windows or windows planned per frame, one stretch of contiguous frames at a time, in video pixels; predict_frames reads chosen frames by random access.

FramePredictions dataclass

Detections at chosen frames of the recording, in video pixels.

What predict_frames returns: the detector's arrays over the frames asked for instead of every frame.

Attributes:

Name	Type	Description
view_names	tuple[str, ...]	The views, in order (the V axis).
frames	tuple[int, ...]	The frame indices predicted, in the order given.
points	ndarray	(V, N, K, 2).
conf	ndarray	(V, N, K).
windows	ndarray | None	(V, N, 4) per-frame boxes in per-frame mode, else None.

Source code in src/deeperfly/pose2d/predict.py

@dataclass(frozen=True)
class FramePredictions:
    """Detections at chosen frames of the recording, in video pixels.

    What [predict_frames][deeperfly.pose2d.predict.predict_frames] returns: the
    detector's arrays over the `frames` asked for instead of every frame.

    Attributes:
        view_names: The views, in order (the `V` axis).
        frames: The frame indices predicted, in the order given.
        points: `(V, N, K, 2)`.
        conf: `(V, N, K)`.
        windows: `(V, N, 4)` per-frame boxes in per-frame mode, else `None`.
    """

    view_names: tuple[str, ...]
    frames: tuple[int, ...]
    points: np.ndarray
    conf: np.ndarray
    windows: np.ndarray | None = None

StretchPredictions dataclass

One stretch of a pass: the detections over contiguous frames, in video pixels.

What predict_recording yields, one per stretch of the frames asked for (ADR-0082). Every per-animal array leads with N, one animal included, and has the stretch's S frames where a results file has T. crops is the rectangle each animal was read through, per view and frame, as slice bounds on the raw frame (ADR-0076): the planner's, or the static crop at every frame.

Attributes:

Name	Type	Description
view_names	tuple[str, ...]	The views, in order (the V axis).
keypoint_names	tuple[str, ...]	The detector's points, in order (the K axis).
start	int	The stretch's first frame.
stop	int	One past its last frame.
points	ndarray	(N, V, S, K, 2).
conf	ndarray	(N, V, S, K).
crops	ndarray	(N, V, S, 4).
candidates	ndarray | None	(N, V, S, K, C, 2), or None without candidates.
candidate_scores	ndarray | None	(N, V, S, K, C), or None.
centers	ndarray | None	(N, S, 3) the tracked centers, or None without a tracker.
track_flags	ndarray | None	(N, S) uint8 track flags, or None without a tracker.
center_candidates	ndarray | None	(V, S, C, 3) the center peaks the tracker read, or None.

Source code in src/deeperfly/pose2d/predict.py

@dataclass(frozen=True)
class StretchPredictions:
    """One stretch of a pass: the detections over contiguous frames, in video pixels.

    What [predict_recording][deeperfly.pose2d.predict.predict_recording] yields, one
    per stretch of the frames asked for (ADR-0082). Every per-animal array leads with
    `N`, one animal included, and has the stretch's `S` frames where a results file
    has `T`. `crops` is the rectangle each animal was read through, per view and
    frame, as slice bounds on the raw frame (ADR-0076): the planner's, or the static
    crop at every frame.

    Attributes:
        view_names: The views, in order (the `V` axis).
        keypoint_names: The detector's points, in order (the `K` axis).
        start: The stretch's first frame.
        stop: One past its last frame.
        points: `(N, V, S, K, 2)`.
        conf: `(N, V, S, K)`.
        crops: `(N, V, S, 4)`.
        candidates: `(N, V, S, K, C, 2)`, or `None` without candidates.
        candidate_scores: `(N, V, S, K, C)`, or `None`.
        centers: `(N, S, 3)` the tracked centers, or `None` without a tracker.
        track_flags: `(N, S)` uint8 track flags, or `None` without a tracker.
        center_candidates: `(V, S, C, 3)` the center peaks the tracker read, or
            `None`.
    """

    view_names: tuple[str, ...]
    keypoint_names: tuple[str, ...]
    start: int
    stop: int
    points: np.ndarray
    conf: np.ndarray
    crops: np.ndarray
    candidates: np.ndarray | None = None
    candidate_scores: np.ndarray | None = None
    centers: np.ndarray | None = None
    track_flags: np.ndarray | None = None
    center_candidates: np.ndarray | None = None

    @property
    def frames(self) -> np.ndarray:
        """The stretch's absolute frame indices."""
        return np.arange(self.start, self.stop)

frames property

frames: ndarray

The stretch's absolute frame indices.

flip_permutation

flip_permutation(skeleton, flip_tta: bool)

The skeleton's left/right permutation as a tensor, or None without TTA.

Source code in src/deeperfly/pose2d/predict.py

def flip_permutation(skeleton, flip_tta: bool):
    """The skeleton's left/right permutation as a tensor, or `None` without TTA."""
    if not flip_tta:
        return None
    if skeleton is None:
        raise ValueError("flip TTA needs the skeleton's left/right pairs")
    return torch.as_tensor(skeleton.flip_perm(), dtype=torch.long)

predict_frames

predict_frames(
    models,
    recording: Recording,
    frames: Sequence[int],
    *,
    views: Sequence[str],
    layout: Layout | None = None,
    device,
    options: PredictOptions | None = None,
    crops: dict[str, tuple[int, int, int, int]]
    | None = None,
    locator=None,
) -> FramePredictions

Predict chosen frames of recording by random access, static or per frame.

The evaluation counterpart of predict_recording over a list of frames. Static windows are crops, or each view's letterbox when None; per frame, the planner runs over the frames in the order given, so consecutive frames keep its fallback. A locator places windows but cannot size them, so it costs the recording's measurement pass first. views are the V axis, read through layout; nothing of a calibration is needed. options.animals, tracking and candidates are not read.

Source code in src/deeperfly/pose2d/predict.py

def predict_frames(
    models,
    recording: Recording,
    frames: Sequence[int],
    *,
    views: Sequence[str],
    layout: Layout | None = None,
    device,
    options: PredictOptions | None = None,
    crops: dict[str, tuple[int, int, int, int]] | None = None,
    locator=None,
) -> FramePredictions:
    """Predict chosen `frames` of `recording` by random access, static or per frame.

    The evaluation counterpart of [predict_recording][] over a list of frames. Static
    windows are `crops`, or each view's letterbox when `None`; per frame, the planner
    runs over the frames in the order given, so consecutive frames keep its fallback.
    A `locator` places windows but cannot size them, so it costs the recording's
    measurement pass first. `views` are the `V` axis, read through `layout`; nothing
    of a calibration is needed. `options.animals`, `tracking` and `candidates` are not
    read.
    """
    options = options or PredictOptions()
    per_frame = options.window == "per_frame"
    videos = Videos(recording, views, layout or Layout())
    input_hw = tuple(first_model(models).input_hw)
    try:
        if per_frame or crops is None:
            boxes = [padded_window(hw, input_hw) for hw in videos.frame_hw]
        else:
            boxes = [crops[view] for view in videos.views]
        planner: WindowPlanner | None = None
        if per_frame and locator is not None:
            planner = FullFrameLocator(
                locator,
                videos,
                input_hw,
                device=device,
                crop_aspect=options.crop_aspect,
            )
            # The locator places the window; it has no joints to size it from, so the
            # recording is measured here exactly as a whole-recording run measures it.
            discover_sizes(
                models,
                videos,
                planner,
                device=device,
                combine=options.combine,
                batch_frames=options.batch_frames,
                tightness=options.tightness,
                log=lambda _: None,
            )
        elif per_frame:
            planner = Locator(
                videos,
                input_hw,
                crop_aspect=options.crop_aspect,
                tightness=options.tightness,
            )
        geom = videos.geometry(boxes, input_hw, device=device)
        frames = [int(t) for t in frames]
        consecutive = frames == list(range(frames[0], frames[0] + len(frames)))
        if consecutive:
            # A clip: one linear decode per source instead of a seek per frame.
            blocks_iter = prefetch(
                videos.stream_blocks(
                    block_size=options.batch_frames,
                    start=frames[0],
                    stop=frames[-1] + 1,
                ),
                depth=2,
            )
        else:
            blocks_iter = (
                videos.frames(frames[i : i + options.batch_frames])
                for i in range(0, len(frames), options.batch_frames)
            )
        outs = [
            block_pass(
                models,
                blocks,
                videos,
                geom,
                planner=planner,
                device=device,
                flip_perm=options.flip_perm,
                combine=options.combine,
                candidates=0,
            )
            for blocks in blocks_iter
        ]
    finally:
        videos.close()
    # One animal, so the pass's `N` axis comes straight off again.
    points = torch.cat([o.points for o in outs])[:, 0].numpy().transpose(1, 0, 2, 3)
    conf = torch.cat([o.conf for o in outs])[:, 0].numpy().transpose(1, 0, 2)
    windows = (
        np.concatenate([o.windows for o in outs])[:, 0]
        .transpose(1, 0, 2)
        .astype(np.float32)
        if per_frame
        else None
    )
    return FramePredictions(tuple(videos.views), tuple(frames), points, conf, windows)

predict_recording

predict_recording(
    models,
    recording: Recording,
    *,
    views: Sequence[str],
    layout: Layout | None = None,
    frames: Sequence[int] | ndarray | None = None,
    crops: ndarray | None = None,
    rig: Rig | None = None,
    device,
    options: PredictOptions | None = None,
    locator=None,
    log: Callable[[str], None] = print,
    on_block: Callable[..., None] | None = None,
) -> Iterator[StretchPredictions]

frames of recording through models, one stretch at a time, in video pixels.

frames (every frame by default) is unioned and sorted, and each stretch of it, a maximal contiguous range, is one linear decode (ADR-0069) that the planner starts afresh: a planner's previous box and the tracks never cross a gap (ADR-0082). One StretchPredictions comes back per stretch, as it is computed. One-frame stretches in a row, as a stepped range makes, carry nothing between them and share detector passes of options.batch_frames.

options.window = "static" predicts through the fixed crops, (N, V, 4) one rectangle per animal and view, or, without them, the ones discovered on the frames asked for. "per_frame" ignores crops: a coarse pass over each view's letterbox locates the animals, a planner draws one box per (frame, animal, view), and the frames are predicted again through those boxes, at about twice the cost of static. Above one animal the animals are found by the detector's center channel or by a locator trained on several (ADR-0127); per frame a Tracker follows them, while static windows hold tethered ones. locator replaces the coarse pass of a per-frame run (ADR-0064). The tracker's and the locator's boxes have one size per view, measured on frames spread over the whole recording, so that they do not depend on the frames asked for (ADR-0089). The keypoint names are the detector's, or options.point_names for an artifact that recorded none.

on_block, when given, sees each detector pass as it completes, before its stretch is assembled: on_block(frames, points, conf, boxes, images) with the pass's absolute frame indices (T,), points (T, N, V, K, 2) and conf (T, N, V, K) in video pixels, the windows they ran through (T, N, V, 4), and per video the (T, H, W) uint8 luma decoded. The arrays are the predictor's own: read them, and copy what is kept. What it raises ends the iteration, the videos closed.

views are the V axis, read through layout. Only several animals need a calibration, to triangulate their centers: rig is read by the tracker and by the discovery of several static crops, and nowhere else.

Raises:

Type	Description
ValueError	Up front, for options that contradict each other, a crop of the wrong shape or a frame outside the recording; while iterating, for an animal that walks out of a static crop.

Source code in src/deeperfly/pose2d/predict.py

def predict_recording(
    models,
    recording: Recording,
    *,
    views: Sequence[str],
    layout: Layout | None = None,
    frames: Sequence[int] | np.ndarray | None = None,
    crops: np.ndarray | None = None,
    rig: Rig | None = None,
    device,
    options: PredictOptions | None = None,
    locator=None,
    log: Callable[[str], None] = print,
    on_block: Callable[..., None] | None = None,
) -> Iterator[StretchPredictions]:
    """`frames` of `recording` through `models`, one stretch at a time, in video pixels.

    `frames` (every frame by default) is unioned and sorted, and each stretch of it, a
    maximal contiguous range, is one linear decode (ADR-0069) that the planner starts
    afresh: a planner's previous box and the tracks never cross a gap (ADR-0082). One
    [StretchPredictions][deeperfly.pose2d.predict.StretchPredictions] comes back per
    stretch, as it is computed. One-frame stretches in a row, as a stepped range makes,
    carry nothing between them and share detector passes of `options.batch_frames`.

    `options.window = "static"` predicts through the fixed `crops`, `(N, V, 4)` one
    rectangle per animal and view, or, without them, the ones discovered on the frames
    asked for. `"per_frame"` ignores `crops`: a coarse pass over each view's letterbox
    locates the animals, a planner draws one box per `(frame, animal, view)`, and the
    frames are predicted again through those boxes, at about twice the cost of static.
    Above one animal the animals are found by the detector's center channel or by a
    `locator` trained on several (ADR-0127); per frame a `Tracker` follows them, while
    static windows hold tethered ones. `locator` replaces the coarse pass of a
    per-frame run (ADR-0064). The tracker's and the locator's boxes have one size per
    view, measured on frames spread over the whole recording, so that they do not
    depend on the frames asked for (ADR-0089). The keypoint names are the detector's,
    or `options.point_names` for an artifact that recorded none.

    `on_block`, when given, sees each detector pass as it completes, before its stretch
    is assembled: `on_block(frames, points, conf, boxes, images)` with the pass's
    absolute frame indices `(T,)`, `points (T, N, V, K, 2)` and `conf (T, N, V, K)` in
    video pixels, the windows they ran through `(T, N, V, 4)`, and per video the
    `(T, H, W)` uint8 luma decoded. The arrays are the predictor's own: read them, and
    copy what is kept. What it raises ends the iteration, the videos closed.

    `views` are the `V` axis, read through `layout`. Only several animals need a
    calibration, to triangulate their centers: `rig` is read by the tracker and by the
    discovery of several static crops, and nowhere else.

    Raises:
        ValueError: Up front, for options that contradict each other, a crop of the
            wrong shape or a frame outside the recording; while iterating, for an
            animal that walks out of a static crop.
    """
    options = options or PredictOptions()
    per_frame = options.window == "per_frame"
    if locator is not None and not per_frame:
        raise ValueError(
            "a locator picks each frame's window; it has nothing to do with "
            'window = "static"'
        )
    if options.animals > 1 and rig is None and (per_frame or crops is None):
        raise ValueError(
            "locating several animals triangulates their centers, so it needs the "
            "calibrated rig"
        )
    fixed = None
    if crops is not None and not per_frame:
        fixed = np.asarray(crops, dtype=float)
        want = (options.animals, len(views), 4)
        if fixed.shape != want:
            raise ValueError(
                f"the static crop is {fixed.shape}; {options.animals} animal(s) and "
                f"{len(views)} views want {want}"
            )
    total = recording.n_frames
    requested = (
        np.arange(total) if frames is None else np.unique(np.asarray(frames, int))
    )
    outside = requested[(requested < 0) | (requested >= total)]
    if outside.size:
        raise ValueError(
            f"{recording.name}: frame {int(outside[0])} is outside its {total} frames"
        )
    return _stretches(
        models,
        recording,
        views,
        layout or Layout(),
        requested,
        fixed,
        rig=rig,
        device=device,
        options=options,
        locator=locator,
        log=log,
        on_block=on_block,
    )

Predict one moment at a time, for a rig that is producing frames right now.

MomentPredictor is the online counterpart of predict_recording, over the same geometry and decode (ADR-0069). It freezes every shape at construction, takes each moment from pinned host buffers a camera SDK writes into and reads the points out of a pinned output buffer: no .cpu() per moment and nothing allocated after the first call. "latency" keeps one moment in flight and "throughput" two. With Follow each moment's windows come from the previous moment's pose instead of staying where crops put them (ADR-0149).

Follow dataclass

Windows that follow the animal: each moment's from the previous moment's pose.

On a tethered rig the static windows must hold the legs' full reach for the whole recording, so they are cut loose; a window drawn around the last pose can be cut tight, which gives the network more pixels on the animal (ADR-0149). The box rule is follow_boxes's.

Attributes:

Name	Type	Description
tightness	float	How much wider than the last pose a window is.
conf_min	float	Confidence a keypoint needs to count toward a box.
keep_frac	float	Share of the confident keypoints, the most confident, a box is drawn around.
floor	float	The box holds at least this share of the box of every keypoint.
min_points	int	A view with fewer confident keypoints re-acquires: its next window is its static one.
reacquire_every	int	Every this many moments, every view re-acquires; 0 never.
static_tightness	float	What the static crops were cut at, which sizes the frames' pre-shrink for windows this much tighter.

Source code in src/deeperfly/pose2d/stream.py

@dataclass(frozen=True)
class Follow:
    """Windows that follow the animal: each moment's from the previous moment's pose.

    On a tethered rig the static windows must hold the legs' full reach for the whole
    recording, so they are cut loose; a window drawn around the last pose can be cut
    tight, which gives the network more pixels on the animal (ADR-0149). The box rule
    is [follow_boxes][deeperfly.pose2d.windows.follow_boxes]'s.

    Attributes:
        tightness: How much wider than the last pose a window is.
        conf_min: Confidence a keypoint needs to count toward a box.
        keep_frac: Share of the confident keypoints, the most confident, a box is
            drawn around.
        floor: The box holds at least this share of the box of every keypoint.
        min_points: A view with fewer confident keypoints re-acquires: its next window
            is its static one.
        reacquire_every: Every this many moments, every view re-acquires; 0 never.
        static_tightness: What the static `crops` were cut at, which sizes the
            frames' pre-shrink for windows this much tighter.
    """

    tightness: float = 1.25
    conf_min: float = 0.3
    keep_frac: float = 0.85
    floor: float = 0.8
    min_points: int = 4
    reacquire_every: int = 0
    static_tightness: float = 1.5

MomentPredictor

One moment of V synchronized frames in, (V, K, 2) video-pixel points out.

Attributes:

Name	Type	Description
frames	tuple[Tensor, ...]	One pinned (H, W) uint8 host buffer per view, written in place; in "throughput" mode they belong to the next moment to submit, so write, submit, write again, collect.
windows		The (V, 4) windows the last collected moment ran through.

Source code in src/deeperfly/pose2d/stream.py

class MomentPredictor:
    """One moment of `V` synchronized frames in, `(V, K, 2)` video-pixel points out.

    Attributes:
        frames: One pinned `(H, W)` uint8 host buffer per view, written in place; in
            `"throughput"` mode they belong to the next moment to submit, so write,
            `submit`, write again, `collect`.
        windows: The `(V, 4)` windows the last collected moment ran through.
    """

    def __init__(
        self,
        model,
        *,
        crops,
        frame_hw,
        views=None,
        precision: str = "auto",
        device=None,
        mode: str = "latency",
        compile_mode: str = "off",
        graph: bool = False,
        follow: Follow | None = None,
    ) -> None:
        """Freeze the shapes, allocate the buffers and warm the forward up.

        `model` is a loaded joint-views detector. `crops` is a window per view: a
        mapping from view name (in `views` order, or its own) or a sequence of boxes.
        `frame_hw` is one `(height, width)` or one per view. `compile_mode` is `"off"`
        or a `torch.compile` mode, applied with `dynamic=False`, and `graph` captures
        the step as a CUDA graph; a capture that fails leaves the eager step and
        records why in `graph_error`. `follow` moves the windows with the animal; then
        `crops` are the windows it starts from and re-acquires with.
        """
        if mode not in MODES:
            raise ValueError(f"mode must be one of {sorted(MODES)}, got {mode!r}")
        if not getattr(model, "joint_views", False):
            raise ValueError(
                "MomentPredictor runs multiview (joint-views) detectors; a per-view "
                "detector needs no moment to be assembled before its forward"
            )
        if isinstance(crops, dict):
            names = list(views) if views is not None else list(crops)
            missing = [n for n in names if n not in crops]
            if missing:
                raise ValueError(f"no window for view(s) {missing}")
            boxes = [crops[n] for n in names]
        else:
            boxes = [tuple(b) for b in crops]
            names = (
                list(views)
                if views is not None
                else [str(i) for i in range(len(boxes))]
            )
        self.model = model
        # A Python float, so it is folded into the captured step (ADR-0078).
        self._shift = grid_shift(model)
        self.views = tuple(names)
        self.mode = mode
        self.point_names = tuple(getattr(model, "point_names", ()) or ())
        dev = (
            torch.device(device)
            if device is not None
            else next(model.parameters()).device
        )
        self.device = dev
        self._cuda = dev.type == "cuda"
        set_precision(model, precision)
        self._dtype = autocast_dtype(model, dev)
        self.input_hw = tuple(int(v) for v in model.input_hw)
        n_v, n_k = len(boxes), int(model.num_classes)
        flat = np.ndim(frame_hw) == 1
        self.frame_hw = tuple(
            (int(frame_hw[0]), int(frame_hw[1])) if flat else (int(h), int(w))
            for h, w in ([frame_hw] * n_v if flat else frame_hw)
        )
        if len(self.frame_hw) != n_v:
            raise ValueError(
                f"{len(self.frame_hw)} frame size(s) for {n_v} view(s); give one pair "
                "or one per view"
            )
        self.geom = WindowGeometry.build(
            boxes, self.frame_hw, self.input_hw, device=dev
        )
        self.follow = follow
        self.windows = np.asarray(boxes, dtype=np.float64)
        self._fill_geom = self.geom
        if follow is not None:
            self._setup_follow(boxes, n_v, n_k, dev)
        self._forward = model
        self.compile_mode = compile_mode
        self.graph_error: str | None = None
        self._slots = [self._new_slot(n_v, n_k, dev) for _ in range(MODES[mode])]
        self._moments = 0
        self._copy_stream = torch.cuda.Stream() if self._cuda else None
        self._write = 0
        self._read = 0
        self._warm()
        if compile_mode != "off":
            # Compiled only after an eager warm-up: a trace taken with the backbone's
            # rotary table not yet cached builds it on the host on every call, and a
            # host tensor copied in makes the step uncapturable.
            self._forward = torch.compile(model, mode=compile_mode, dynamic=False)
            self._warm()
        if graph:
            self._capture()
        self.reset()

    def _setup_follow(self, boxes, n_v: int, n_k: int, dev) -> None:
        """The follow state: the start windows, the pre-shrink, the last pose.

        The frames are pre-shrunk once, for windows `tightness / static_tightness` the
        size of `crops`: shrunk for the static windows instead, a tighter window would
        be upsampled out of a frame that had already lost what it gains.
        """
        f = self.follow
        ratio = f.tightness / f.static_tightness
        nominal = []
        for x0, y0, x1, y1 in boxes:
            cx, cy = (x0 + x1) / 2, (y0 + y1) / 2
            hw, hh = (x1 - x0) / 2 * ratio, (y1 - y0) / 2 * ratio
            nominal.append((cx - hw, cy - hh, cx + hw, cy + hh))
        self._fill_geom = WindowGeometry.build(
            nominal, self.frame_hw, self.input_hw, device=dev
        )
        g = self._fill_geom
        self._shrink_scale = torch.tensor(
            [
                [s[1] / fr[1], s[0] / fr[0]]
                for s, fr in zip(g.shrink_hw, g.frame_hw, strict=True)
            ],
            dtype=torch.float32,
            device=dev,
        )
        self._static = torch.tensor(boxes, dtype=torch.float32, device=dev)
        aspect = self.input_hw[1] / self.input_hw[0]
        self._aspects = torch.full((n_v,), aspect, dtype=torch.float32, device=dev)
        self._last_pts = torch.zeros((n_v, n_k, 2), device=dev)
        self._last_conf = torch.zeros((n_v, n_k), device=dev)
        self._have_prev = torch.zeros((), dtype=torch.bool, device=dev)
        self._force = torch.zeros((), dtype=torch.bool, device=dev)

    def reset(self) -> None:
        """Forget the last pose: the next moment runs through the static windows."""
        self._moments = 0
        if self.follow is not None:
            self._have_prev.fill_(False)
            self._force.fill_(False)

    def reacquire(self) -> None:
        """Run the next moment submitted through the static windows, then follow on."""
        if self.follow is not None:
            self._force.fill_(True)

    def _new_slot(self, n_v: int, n_k: int, dev) -> _Slot:
        """One slot's buffers, contiguous per moment where the views agree on a size."""
        uniform = len(set(self.frame_hw)) == 1
        if uniform:
            block = torch.zeros(
                (n_v, *self.frame_hw[0]), dtype=torch.uint8, pin_memory=self._cuda
            )
            staged_block = torch.zeros(
                (n_v, *self.frame_hw[0]), dtype=torch.uint8, device=dev
            )
            frames, staged = tuple(block.unbind(0)), tuple(staged_block.unbind(0))
        else:
            block = staged_block = None
            frames = tuple(
                torch.zeros(hw, dtype=torch.uint8, pin_memory=self._cuda)
                for hw in self.frame_hw
            )
            staged = tuple(
                torch.zeros(hw, dtype=torch.uint8, device=dev) for hw in self.frame_hw
            )
        return _Slot(
            frames=frames,
            block=block,
            staged=staged,
            staged_block=staged_block,
            canvas=self._fill_geom.new_canvas(1),
            points=torch.zeros((n_v, n_k, 2), pin_memory=self._cuda),
            conf=torch.zeros((n_v, n_k), pin_memory=self._cuda),
            windows=torch.zeros((n_v, 4), pin_memory=self._cuda),
            copied=torch.cuda.Event() if self._cuda else None,
            done=torch.cuda.Event() if self._cuda else None,
        )

    # -- the surface ---------------------------------------------------------------

    @property
    def frames(self) -> tuple[torch.Tensor, ...]:
        """The pinned per-view buffers the next `submit` takes its moment from."""
        return self._slots[self._write].frames

    def submit(self) -> None:
        """Stage `frames` and launch the moment, with no sync and nothing read back."""
        slot = self._slots[self._write]
        if slot.in_flight:
            raise RuntimeError(
                f"all {len(self._slots)} slot(s) are in flight; collect() a moment "
                "before submitting another"
            )
        pairs = (
            [(slot.staged_block, slot.block)]
            if slot.block is not None
            else list(zip(slot.staged, slot.frames, strict=True))
        )
        if self._cuda:
            # The copy stream must not overwrite a staging buffer whose step is still
            # reading it, which is the whole reason the slot carries a done event.
            self._copy_stream.wait_event(slot.done)
            with torch.cuda.stream(self._copy_stream):
                for dst, src in pairs:
                    dst.copy_(src, non_blocking=True)
                slot.copied.record(self._copy_stream)
            torch.cuda.current_stream().wait_event(slot.copied)
        else:
            for dst, src in pairs:
                dst.copy_(src)
        every = 0 if self.follow is None else self.follow.reacquire_every
        if every and self._moments and self._moments % every == 0:
            self.reacquire()
        self._moments += 1
        if slot.graph is not None:
            slot.graph.replay()
        else:
            self._step(slot)
        if self._cuda:
            slot.done.record()
        slot.in_flight = True
        self._write = (self._write + 1) % len(self._slots)

    def collect(self) -> tuple[np.ndarray, np.ndarray]:
        """Wait for the oldest moment: `(V, K, 2)` video pixels and `(V, K)` conf."""
        slot = self._slots[self._read]
        if not slot.in_flight:
            raise RuntimeError("nothing is in flight; submit() a moment first")
        if self._cuda:
            slot.done.synchronize()
        slot.in_flight = False
        self._read = (self._read + 1) % len(self._slots)
        if self.follow is not None:
            self.windows = slot.windows.numpy().astype(np.float64)
        # Copied out because the slot's pinned buffers are reused by the next moment.
        return slot.points.numpy().copy(), slot.conf.numpy().copy()

    def predict(self, frames=None) -> tuple[np.ndarray, np.ndarray]:
        """One moment, submitted and collected, for a caller with no pipeline."""
        if frames is not None:
            for buf, frame in zip(self.frames, frames, strict=True):
                buf.copy_(torch.as_tensor(np.asarray(frame)))
        self.submit()
        return self.collect()

    # -- the step ------------------------------------------------------------------

    def _step(self, slot: _Slot) -> None:
        """Staged frames -> the slot's output buffers, all on device.

        One `float(tensor)` in here would make the moment uncapturable.
        """
        with torch.no_grad():
            self._fill_geom.fill(slot.canvas, slot.staged)
            if self.follow is None:
                images = self.geom.sample(slot.canvas)
            else:
                images, scale, offset = self._follow_sample(slot)
            with torch.autocast(
                self.device.type, dtype=self._dtype, enabled=self._dtype is not None
            ):
                out = self._forward(images)
                v = images.shape[1]
                k = out.shape[1] // v
                pts, conf = decode_points(
                    self.model, out.reshape(v, k, *out.shape[-2:])
                )
            if self.follow is None:
                px = self.geom.to_video_px(pts.float()[None])[0] + self._shift
            else:
                px = pts.float() * self.geom.px_scale
                px = px * scale[:, None, :] + offset[:, None, :] + self._shift
                self._last_pts.copy_(px)
                self._last_conf.copy_(conf.float())
                self._have_prev.fill_(True)
                self._force.fill_(False)
            slot.points.copy_(px, non_blocking=self._cuda)
            slot.conf.copy_(conf.float(), non_blocking=self._cuda)

    def _follow_sample(self, slot: _Slot):
        """This moment's windows from the last pose, and the input through them."""
        f = self.follow
        boxes = follow_boxes(
            self._last_pts,
            self._last_conf,
            self._static,
            have_prev=self._have_prev,
            force_static=self._force,
            aspects=self._aspects,
            min_w=self.input_hw[1] / 2,
            tightness=f.tightness,
            conf_min=f.conf_min,
            keep_frac=f.keep_frac,
            floor=f.floor,
            min_points=f.min_points,
        )
        theta, scale, offset = follow_affines(
            boxes, self._shrink_scale, self._fill_geom.canvas_hw, self.input_hw
        )
        grid = F.affine_grid(
            theta, (len(boxes), 1, *self.input_hw), align_corners=False
        )
        out = F.grid_sample(
            slot.canvas,
            grid,
            mode="bilinear",
            padding_mode="zeros",
            align_corners=False,
        )
        slot.windows.copy_(boxes, non_blocking=self._cuda)
        return out.permute(1, 0, 2, 3).unsqueeze(2), scale, offset

    def _warm(self) -> None:
        """Run the step until its shapes settle, so a capture records steady state."""
        for _ in range(WARMUP):
            self._step(self._slots[0])
        if self._cuda:
            torch.cuda.synchronize()

    def _capture(self) -> None:
        """Capture the per-moment step per slot, or record why it could not be."""
        if not self._cuda:
            self.graph_error = "CUDA graphs need a CUDA device"
            return
        try:
            side = torch.cuda.Stream()
            side.wait_stream(torch.cuda.current_stream())
            with torch.cuda.stream(side):
                for slot in self._slots:
                    for _ in range(WARMUP):
                        self._step(slot)
            torch.cuda.current_stream().wait_stream(side)
            torch.cuda.synchronize()
            pool = torch.cuda.graph_pool_handle()
            for slot in self._slots:
                g = torch.cuda.CUDAGraph()
                with torch.cuda.graph(g, pool=pool):
                    self._step(slot)
                slot.graph = g
            torch.cuda.synchronize()
        except (
            RuntimeError
        ) as exc:  # a capture blocker is a fact to report, not a crash
            for slot in self._slots:
                slot.graph = None
            self.graph_error = f"{type(exc).__name__}: {exc}"
            torch.cuda.synchronize()

frames property

frames: tuple[Tensor, ...]

The pinned per-view buffers the next submit takes its moment from.

__init__

__init__(
    model,
    *,
    crops,
    frame_hw,
    views=None,
    precision: str = "auto",
    device=None,
    mode: str = "latency",
    compile_mode: str = "off",
    graph: bool = False,
    follow: Follow | None = None,
) -> None

Freeze the shapes, allocate the buffers and warm the forward up.

model is a loaded joint-views detector. crops is a window per view: a mapping from view name (in views order, or its own) or a sequence of boxes. frame_hw is one (height, width) or one per view. compile_mode is "off" or a torch.compile mode, applied with dynamic=False, and graph captures the step as a CUDA graph; a capture that fails leaves the eager step and records why in graph_error. follow moves the windows with the animal; then crops are the windows it starts from and re-acquires with.

Source code in src/deeperfly/pose2d/stream.py

def __init__(
    self,
    model,
    *,
    crops,
    frame_hw,
    views=None,
    precision: str = "auto",
    device=None,
    mode: str = "latency",
    compile_mode: str = "off",
    graph: bool = False,
    follow: Follow | None = None,
) -> None:
    """Freeze the shapes, allocate the buffers and warm the forward up.

    `model` is a loaded joint-views detector. `crops` is a window per view: a
    mapping from view name (in `views` order, or its own) or a sequence of boxes.
    `frame_hw` is one `(height, width)` or one per view. `compile_mode` is `"off"`
    or a `torch.compile` mode, applied with `dynamic=False`, and `graph` captures
    the step as a CUDA graph; a capture that fails leaves the eager step and
    records why in `graph_error`. `follow` moves the windows with the animal; then
    `crops` are the windows it starts from and re-acquires with.
    """
    if mode not in MODES:
        raise ValueError(f"mode must be one of {sorted(MODES)}, got {mode!r}")
    if not getattr(model, "joint_views", False):
        raise ValueError(
            "MomentPredictor runs multiview (joint-views) detectors; a per-view "
            "detector needs no moment to be assembled before its forward"
        )
    if isinstance(crops, dict):
        names = list(views) if views is not None else list(crops)
        missing = [n for n in names if n not in crops]
        if missing:
            raise ValueError(f"no window for view(s) {missing}")
        boxes = [crops[n] for n in names]
    else:
        boxes = [tuple(b) for b in crops]
        names = (
            list(views)
            if views is not None
            else [str(i) for i in range(len(boxes))]
        )
    self.model = model
    # A Python float, so it is folded into the captured step (ADR-0078).
    self._shift = grid_shift(model)
    self.views = tuple(names)
    self.mode = mode
    self.point_names = tuple(getattr(model, "point_names", ()) or ())
    dev = (
        torch.device(device)
        if device is not None
        else next(model.parameters()).device
    )
    self.device = dev
    self._cuda = dev.type == "cuda"
    set_precision(model, precision)
    self._dtype = autocast_dtype(model, dev)
    self.input_hw = tuple(int(v) for v in model.input_hw)
    n_v, n_k = len(boxes), int(model.num_classes)
    flat = np.ndim(frame_hw) == 1
    self.frame_hw = tuple(
        (int(frame_hw[0]), int(frame_hw[1])) if flat else (int(h), int(w))
        for h, w in ([frame_hw] * n_v if flat else frame_hw)
    )
    if len(self.frame_hw) != n_v:
        raise ValueError(
            f"{len(self.frame_hw)} frame size(s) for {n_v} view(s); give one pair "
            "or one per view"
        )
    self.geom = WindowGeometry.build(
        boxes, self.frame_hw, self.input_hw, device=dev
    )
    self.follow = follow
    self.windows = np.asarray(boxes, dtype=np.float64)
    self._fill_geom = self.geom
    if follow is not None:
        self._setup_follow(boxes, n_v, n_k, dev)
    self._forward = model
    self.compile_mode = compile_mode
    self.graph_error: str | None = None
    self._slots = [self._new_slot(n_v, n_k, dev) for _ in range(MODES[mode])]
    self._moments = 0
    self._copy_stream = torch.cuda.Stream() if self._cuda else None
    self._write = 0
    self._read = 0
    self._warm()
    if compile_mode != "off":
        # Compiled only after an eager warm-up: a trace taken with the backbone's
        # rotary table not yet cached builds it on the host on every call, and a
        # host tensor copied in makes the step uncapturable.
        self._forward = torch.compile(model, mode=compile_mode, dynamic=False)
        self._warm()
    if graph:
        self._capture()
    self.reset()

collect

collect() -> tuple[np.ndarray, np.ndarray]

Wait for the oldest moment: (V, K, 2) video pixels and (V, K) conf.

Source code in src/deeperfly/pose2d/stream.py

def collect(self) -> tuple[np.ndarray, np.ndarray]:
    """Wait for the oldest moment: `(V, K, 2)` video pixels and `(V, K)` conf."""
    slot = self._slots[self._read]
    if not slot.in_flight:
        raise RuntimeError("nothing is in flight; submit() a moment first")
    if self._cuda:
        slot.done.synchronize()
    slot.in_flight = False
    self._read = (self._read + 1) % len(self._slots)
    if self.follow is not None:
        self.windows = slot.windows.numpy().astype(np.float64)
    # Copied out because the slot's pinned buffers are reused by the next moment.
    return slot.points.numpy().copy(), slot.conf.numpy().copy()

predict

predict(frames=None) -> tuple[np.ndarray, np.ndarray]

One moment, submitted and collected, for a caller with no pipeline.

Source code in src/deeperfly/pose2d/stream.py

def predict(self, frames=None) -> tuple[np.ndarray, np.ndarray]:
    """One moment, submitted and collected, for a caller with no pipeline."""
    if frames is not None:
        for buf, frame in zip(self.frames, frames, strict=True):
            buf.copy_(torch.as_tensor(np.asarray(frame)))
    self.submit()
    return self.collect()

reacquire

reacquire() -> None

Run the next moment submitted through the static windows, then follow on.

Source code in src/deeperfly/pose2d/stream.py

def reacquire(self) -> None:
    """Run the next moment submitted through the static windows, then follow on."""
    if self.follow is not None:
        self._force.fill_(True)

reset

reset() -> None

Forget the last pose: the next moment runs through the static windows.

Source code in src/deeperfly/pose2d/stream.py

def reset(self) -> None:
    """Forget the last pose: the next moment runs through the static windows."""
    self._moments = 0
    if self.follow is not None:
        self._have_prev.fill_(False)
        self._force.fill_(False)

submit

submit() -> None

Stage frames and launch the moment, with no sync and nothing read back.

Source code in src/deeperfly/pose2d/stream.py

def submit(self) -> None:
    """Stage `frames` and launch the moment, with no sync and nothing read back."""
    slot = self._slots[self._write]
    if slot.in_flight:
        raise RuntimeError(
            f"all {len(self._slots)} slot(s) are in flight; collect() a moment "
            "before submitting another"
        )
    pairs = (
        [(slot.staged_block, slot.block)]
        if slot.block is not None
        else list(zip(slot.staged, slot.frames, strict=True))
    )
    if self._cuda:
        # The copy stream must not overwrite a staging buffer whose step is still
        # reading it, which is the whole reason the slot carries a done event.
        self._copy_stream.wait_event(slot.done)
        with torch.cuda.stream(self._copy_stream):
            for dst, src in pairs:
                dst.copy_(src, non_blocking=True)
            slot.copied.record(self._copy_stream)
        torch.cuda.current_stream().wait_event(slot.copied)
    else:
        for dst, src in pairs:
            dst.copy_(src)
    every = 0 if self.follow is None else self.follow.reacquire_every
    if every and self._moments and self._moments % every == 0:
        self.reacquire()
    self._moments += 1
    if slot.graph is not None:
        slot.graph.replay()
    else:
        self._step(slot)
    if self._cuda:
        slot.done.record()
    slot.in_flight = True
    self._write = (self._write + 1) % len(self._slots)