sleap-nn 0.0.5__py3-none-any.whl → 0.1.0a0__py3-none-any.whl

sleap_nn/training/utils.py

@@ -1,13 +1,19 @@
 """Miscellaneous utility functions for training."""
 
+from dataclasses import dataclass, field
+from io import BytesIO
 import numpy as np
+import matplotlib
+
+matplotlib.use(
+    "Agg"
+)  # Use non-interactive backend to avoid tkinter issues on Windows CI
 import matplotlib.pyplot as plt
 from loguru import logger
 from torch import nn
 import torch.distributed as dist
-import matplotlib
 import seaborn as sns
-from typing import List
+from typing import List, Optional
 import shutil
 import os
 import subprocess
@@ -236,3 +242,362 @@ def plot_peaks(
                 )
             )
     return handles
+
+
+@dataclass
+class VisualizationData:
+    """Container for visualization data from a single sample.
+
+    This dataclass decouples data extraction from rendering, allowing the same
+    data to be rendered to different output targets (matplotlib, wandb, etc.).
+
+    Attributes:
+        image: Input image as (H, W, C) numpy array, normalized to [0, 1].
+        pred_confmaps: Predicted confidence maps as (H, W, nodes) array, values in [0, 1].
+        pred_peaks: Predicted keypoints as (instances, nodes, 2) or (nodes, 2) array.
+        pred_peak_values: Confidence values as (instances, nodes) or (nodes,) array.
+        gt_instances: Ground truth keypoints, same shape as pred_peaks.
+        node_names: List of node/keypoint names, e.g., ["head", "thorax", ...].
+        output_scale: Ratio of confmap size to image size (confmap_h / image_h).
+        is_paired: Whether GT and predictions can be paired for error visualization.
+        pred_pafs: Part affinity fields for bottom-up models, optional.
+        pred_class_maps: Class maps for multi-class models, optional.
+    """
+
+    image: np.ndarray
+    pred_confmaps: np.ndarray
+    pred_peaks: np.ndarray
+    pred_peak_values: np.ndarray
+    gt_instances: np.ndarray
+    node_names: List[str] = field(default_factory=list)
+    output_scale: float = 1.0
+    is_paired: bool = True
+    pred_pafs: Optional[np.ndarray] = None
+    pred_class_maps: Optional[np.ndarray] = None
+
+
+class MatplotlibRenderer:
+    """Renders VisualizationData to matplotlib figures."""
+
+    def render(self, data: VisualizationData) -> matplotlib.figure.Figure:
+        """Render visualization data to a matplotlib figure.
+
+        Args:
+            data: VisualizationData containing image, confmaps, peaks, etc.
+
+        Returns:
+            A matplotlib Figure object.
+        """
+        img = data.image
+        scale = 1.0
+        if img.shape[0] < 512:
+            scale = 2.0
+        if img.shape[0] < 256:
+            scale = 4.0
+
+        fig = plot_img(img, dpi=72 * scale, scale=scale)
+        plot_confmaps(data.pred_confmaps, output_scale=data.output_scale)
+        plot_peaks(data.gt_instances, data.pred_peaks, paired=data.is_paired)
+        return fig
+
+    def render_pafs(self, data: VisualizationData) -> matplotlib.figure.Figure:
+        """Render PAF magnitude visualization.
+
+        Args:
+            data: VisualizationData with pred_pafs populated.
+
+        Returns:
+            A matplotlib Figure object showing PAF magnitudes.
+        """
+        if data.pred_pafs is None:
+            raise ValueError("pred_pafs is None, cannot render PAFs")
+
+        img = data.image
+        scale = 1.0
+        if img.shape[0] < 512:
+            scale = 2.0
+        if img.shape[0] < 256:
+            scale = 4.0
+
+        # Compute PAF magnitude
+        pafs = data.pred_pafs  # (H, W, 2*edges) or (H, W, edges, 2)
+        if pafs.ndim == 3:
+            n_edges = pafs.shape[-1] // 2
+            pafs = pafs.reshape(pafs.shape[0], pafs.shape[1], n_edges, 2)
+        magnitude = np.sqrt(pafs[..., 0] ** 2 + pafs[..., 1] ** 2)
+        magnitude = magnitude.max(axis=-1)  # Max over edges
+
+        fig = plot_img(img, dpi=72 * scale, scale=scale)
+        ax = plt.gca()
+
+        # Calculate PAF output scale from actual PAF dimensions, not confmap output_scale
+        # PAFs may have a different output_stride than confmaps
+        paf_output_scale = magnitude.shape[0] / img.shape[0]
+
+        ax.imshow(
+            magnitude,
+            alpha=0.5,
+            origin="upper",
+            cmap="viridis",
+            extent=[
+                -0.5,
+                magnitude.shape[1] / paf_output_scale - 0.5,
+                magnitude.shape[0] / paf_output_scale - 0.5,
+                -0.5,
+            ],
+        )
+        return fig
+
+
+class WandBRenderer:
+    """Renders VisualizationData to wandb.Image objects.
+
+    Supports multiple rendering modes:
+    - "direct": Pre-render with matplotlib, convert to wandb.Image
+    - "boxes": Use wandb boxes for interactive keypoint visualization
+    - "masks": Use wandb masks for confidence map overlay
+    """
+
+    def __init__(
+        self,
+        mode: str = "direct",
+        box_size: float = 5.0,
+        confmap_threshold: float = 0.1,
+        min_size: int = 512,
+    ):
+        """Initialize the renderer.
+
+        Args:
+            mode: Rendering mode - "direct", "boxes", or "masks".
+            box_size: Size of keypoint boxes in pixels (for "boxes" mode).
+            confmap_threshold: Threshold for confmap mask (for "masks" mode).
+            min_size: Minimum image dimension. Smaller images will be upscaled.
+        """
+        self.mode = mode
+        self.box_size = box_size
+        self.confmap_threshold = confmap_threshold
+        self.min_size = min_size
+        self._mpl_renderer = MatplotlibRenderer()
+
+    def render(
+        self, data: VisualizationData, caption: Optional[str] = None
+    ) -> "wandb.Image":
+        """Render visualization data to a wandb.Image.
+
+        Args:
+            data: VisualizationData containing image, confmaps, peaks, etc.
+            caption: Optional caption for the image.
+
+        Returns:
+            A wandb.Image object.
+        """
+        import wandb
+
+        if self.mode == "boxes":
+            return self._render_with_boxes(data, caption)
+        elif self.mode == "masks":
+            return self._render_with_masks(data, caption)
+        else:  # "direct"
+            return self._render_direct(data, caption)
+
+    def _get_scale_factor(self, img_h: int, img_w: int) -> int:
+        """Calculate scale factor to ensure minimum image size."""
+        min_dim = min(img_h, img_w)
+        if min_dim >= self.min_size:
+            return 1
+        return int(np.ceil(self.min_size / min_dim))
+
+    def _render_direct(
+        self, data: VisualizationData, caption: Optional[str] = None
+    ) -> "wandb.Image":
+        """Pre-render with matplotlib, return as wandb.Image."""
+        import wandb
+        from PIL import Image
+
+        fig = self._mpl_renderer.render(data)
+
+        # Convert figure to PIL Image
+        buf = BytesIO()
+        fig.savefig(buf, format="png", bbox_inches="tight", pad_inches=0)
+        buf.seek(0)
+        plt.close(fig)
+
+        pil_image = Image.open(buf)
+        return wandb.Image(pil_image, caption=caption)
+
+    def _render_with_boxes(
+        self, data: VisualizationData, caption: Optional[str] = None
+    ) -> "wandb.Image":
+        """Use wandb boxes for interactive keypoint visualization."""
+        import wandb
+        from PIL import Image
+
+        # Prepare class labels from node names
+        class_labels = {i: name for i, name in enumerate(data.node_names)}
+        if not class_labels:
+            class_labels = {i: f"node_{i}" for i in range(data.pred_peaks.shape[-2])}
+
+        # Convert image to uint8
+        img_uint8 = (np.clip(data.image, 0, 1) * 255).astype(np.uint8)
+        # Handle single-channel images: squeeze (H, W, 1) -> (H, W)
+        if img_uint8.ndim == 3 and img_uint8.shape[2] == 1:
+            img_uint8 = img_uint8.squeeze(axis=2)
+        img_h, img_w = img_uint8.shape[:2]
+
+        # Scale up small images for better visibility in wandb
+        scale = self._get_scale_factor(img_h, img_w)
+        if scale > 1:
+            pil_img = Image.fromarray(img_uint8)
+            pil_img = pil_img.resize(
+                (img_w * scale, img_h * scale), resample=Image.BILINEAR
+            )
+            img_uint8 = np.array(pil_img)
+
+        # Build ground truth boxes (use percent domain for proper thumbnail scaling)
+        gt_box_data = self._peaks_to_boxes(
+            data.gt_instances, data.node_names, img_w, img_h, is_gt=True
+        )
+
+        # Build prediction boxes
+        pred_box_data = self._peaks_to_boxes(
+            data.pred_peaks,
+            data.node_names,
+            img_w,
+            img_h,
+            peak_values=data.pred_peak_values,
+            is_gt=False,
+        )
+
+        return wandb.Image(
+            img_uint8,
+            boxes={
+                "ground_truth": {"box_data": gt_box_data, "class_labels": class_labels},
+                "predictions": {
+                    "box_data": pred_box_data,
+                    "class_labels": class_labels,
+                },
+            },
+            caption=caption,
+        )
+
+    def _peaks_to_boxes(
+        self,
+        peaks: np.ndarray,
+        node_names: List[str],
+        img_w: int,
+        img_h: int,
+        peak_values: Optional[np.ndarray] = None,
+        is_gt: bool = False,
+    ) -> List[dict]:
+        """Convert peaks array to wandb box_data format.
+
+        Args:
+            peaks: Keypoints as (instances, nodes, 2) or (nodes, 2).
+            node_names: List of node names.
+            img_w: Image width in pixels.
+            img_h: Image height in pixels.
+            peak_values: Optional confidence values.
+            is_gt: Whether these are ground truth points.
+
+        Returns:
+            List of box dictionaries for wandb.
+        """
+        box_data = []
+
+        # Normalize shape to (instances, nodes, 2)
+        if peaks.ndim == 2:
+            peaks = peaks[np.newaxis, ...]
+            if peak_values is not None and peak_values.ndim == 1:
+                peak_values = peak_values[np.newaxis, ...]
+
+        # Convert box_size from pixels to percent
+        box_w_pct = self.box_size / img_w
+        box_h_pct = self.box_size / img_h
+
+        for inst_idx, instance in enumerate(peaks):
+            for node_idx, (x, y) in enumerate(instance):
+                if np.isnan(x) or np.isnan(y):
+                    continue
+
+                node_name = (
+                    node_names[node_idx]
+                    if node_idx < len(node_names)
+                    else f"node_{node_idx}"
+                )
+
+                # Convert pixel coordinates to percent (0-1 range)
+                x_pct = float(x) / img_w
+                y_pct = float(y) / img_h
+
+                box = {
+                    "position": {
+                        "middle": [x_pct, y_pct],
+                        "width": box_w_pct,
+                        "height": box_h_pct,
+                    },
+                    "domain": "percent",
+                    "class_id": node_idx,
+                }
+
+                if is_gt:
+                    box["box_caption"] = f"GT: {node_name}"
+                else:
+                    if peak_values is not None:
+                        conf = float(peak_values[inst_idx, node_idx])
+                        box["box_caption"] = f"{node_name} ({conf:.2f})"
+                        box["scores"] = {"confidence": conf}
+                    else:
+                        box["box_caption"] = node_name
+
+                box_data.append(box)
+
+        return box_data
+
+    def _render_with_masks(
+        self, data: VisualizationData, caption: Optional[str] = None
+    ) -> "wandb.Image":
+        """Use wandb masks for confidence map overlay.
+
+        Uses argmax approach: each pixel shows the dominant node.
+        """
+        import wandb
+
+        # Prepare class labels (0 = background, 1+ = nodes)
+        class_labels = {0: "background"}
+        for i, name in enumerate(data.node_names):
+            class_labels[i + 1] = name
+        if not data.node_names:
+            n_nodes = data.pred_confmaps.shape[-1]
+            for i in range(n_nodes):
+                class_labels[i + 1] = f"node_{i}"
+
+        # Create argmax mask from confmaps
+        confmaps = data.pred_confmaps  # (H/stride, W/stride, nodes)
+        max_vals = confmaps.max(axis=-1)
+        argmax_map = confmaps.argmax(axis=-1) + 1  # +1 for background offset
+        argmax_map[max_vals < self.confmap_threshold] = 0  # Background
+
+        # Convert image to uint8
+        img_uint8 = (np.clip(data.image, 0, 1) * 255).astype(np.uint8)
+        # Handle single-channel images: (H, W, 1) -> (H, W)
+        if img_uint8.ndim == 3 and img_uint8.shape[2] == 1:
+            img_uint8 = img_uint8.squeeze(axis=2)
+        img_h, img_w = img_uint8.shape[:2]
+
+        # Resize mask to match image dimensions (confmaps are H/stride, W/stride)
+        from PIL import Image
+
+        mask_pil = Image.fromarray(argmax_map.astype(np.uint8))
+        mask_pil = mask_pil.resize((img_w, img_h), resample=Image.NEAREST)
+        argmax_map = np.array(mask_pil)
+
+        return wandb.Image(
+            img_uint8,
+            masks={
+                "confidence_maps": {
+                    "mask_data": argmax_map.astype(np.uint8),
+                    "class_labels": class_labels,
+                }
+            },
+            caption=caption,
+        )

{sleap_nn-0.0.5.dist-info → sleap_nn-0.1.0a0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sleap-nn
-Version: 0.0.5
+Version: 0.1.0a0
 Summary: Neural network backend for training and inference for animal pose estimation.
 Author-email: Divya Seshadri Murali <dimurali@salk.edu>, Elizabeth Berrigan <eberrigan@salk.edu>, Vincent Tu <vitu@ucsd.edu>, Liezl Maree <lmaree@salk.edu>, David Samy <davidasamy@gmail.com>, Talmo Pereira <talmo@salk.edu>
 License: BSD-3-Clause
@@ -13,7 +13,7 @@ Classifier: Programming Language :: Python :: 3.13
 Requires-Python: <3.14,>=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: sleap-io>=0.5.7
+Requires-Dist: sleap-io<0.7.0,>=0.6.0
 Requires-Dist: numpy
 Requires-Dist: lightning
 Requires-Dist: kornia
@@ -34,35 +34,19 @@ Requires-Dist: jupyterlab
 Requires-Dist: pyzmq
 Provides-Extra: torch
 Requires-Dist: torch; extra == "torch"
-Requires-Dist: torchvision<0.24.0,>=0.20.0; extra == "torch"
+Requires-Dist: torchvision>=0.20.0; extra == "torch"
 Provides-Extra: torch-cpu
 Requires-Dist: torch; extra == "torch-cpu"
-Requires-Dist: torchvision<0.24.0,>=0.20.0; extra == "torch-cpu"
+Requires-Dist: torchvision>=0.20.0; extra == "torch-cpu"
 Provides-Extra: torch-cuda118
 Requires-Dist: torch; extra == "torch-cuda118"
-Requires-Dist: torchvision<0.24.0,>=0.20.0; extra == "torch-cuda118"
+Requires-Dist: torchvision>=0.20.0; extra == "torch-cuda118"
 Provides-Extra: torch-cuda128
 Requires-Dist: torch; extra == "torch-cuda128"
-Requires-Dist: torchvision<0.24.0,>=0.20.0; extra == "torch-cuda128"
-Provides-Extra: dev
-Requires-Dist: pytest; extra == "dev"
-Requires-Dist: pytest-cov; extra == "dev"
-Requires-Dist: black; extra == "dev"
-Requires-Dist: pydocstyle; extra == "dev"
-Requires-Dist: toml; extra == "dev"
-Requires-Dist: twine; extra == "dev"
-Requires-Dist: build; extra == "dev"
-Requires-Dist: ipython; extra == "dev"
-Requires-Dist: ruff; extra == "dev"
-Provides-Extra: docs
-Requires-Dist: mkdocs; extra == "docs"
-Requires-Dist: mkdocs-material; extra == "docs"
-Requires-Dist: mkdocs-jupyter; extra == "docs"
-Requires-Dist: mike; extra == "docs"
-Requires-Dist: mkdocstrings[python]; extra == "docs"
-Requires-Dist: mkdocs-gen-files; extra == "docs"
-Requires-Dist: mkdocs-literate-nav; extra == "docs"
-Requires-Dist: mkdocs-section-index; extra == "docs"
+Requires-Dist: torchvision>=0.20.0; extra == "torch-cuda128"
+Provides-Extra: torch-cuda130
+Requires-Dist: torch; extra == "torch-cuda130"
+Requires-Dist: torchvision>=0.20.0; extra == "torch-cuda130"
 Dynamic: license-file
 
 # sleap-nn
@@ -120,22 +104,28 @@ powershell -c "irm https://astral.sh/uv/install.ps1 | iex"
 > Replace `...` with the rest of your install command as needed.
 
 - Sync all dependencies based on your correct wheel using `uv sync`. `uv sync` creates a `.venv` (virtual environment) inside your current working directory. This environment is only active within that directory and can't be directly accessed from outside. To use all installed packages, you must run commands with `uv run` (e.g., `uv run sleap-nn train ...` or `uv run pytest ...`).
-   - **Windows/Linux with NVIDIA GPU (CUDA 11.8):**
+   - **Windows/Linux with NVIDIA GPU (CUDA 13.0):**
 
    ```bash
-   uv sync --extra dev --extra torch-cuda118
+   uv sync --extra torch-cuda130
    ```
 
    - **Windows/Linux with NVIDIA GPU (CUDA 12.8):**
 
    ```bash
-   uv sync --extra dev --extra torch-cuda128
+   uv sync --extra torch-cuda128
    ```
-   
-   - **macOS with Apple Silicon (M1, M2, M3, M4) or CPU-only (no GPU or unsupported GPU):** 
+
+   - **Windows/Linux with NVIDIA GPU (CUDA 11.8):**
+
+   ```bash
+   uv sync --extra torch-cuda118
+   ```
+
+   - **macOS with Apple Silicon (M1, M2, M3, M4) or CPU-only (no GPU or unsupported GPU):**
    Note: Even if torch-cpu is used on macOS, the MPS backend will be available.
    ```bash
-   uv sync --extra dev --extra torch-cpu
+   uv sync --extra torch-cpu
    ```
 
 4. **Run tests**  
@@ -152,6 +142,6 @@ powershell -c "irm https://astral.sh/uv/install.ps1 | iex"
 > **Upgrading All Dependencies**
 > To ensure you have the latest versions of all dependencies, use the `--upgrade` flag with `uv sync`:
 > ```bash
-> uv sync --extra dev --upgrade
+> uv sync --upgrade
 > ```
 > This will upgrade all installed packages in your environment to the latest available versions compatible with your `pyproject.toml`.

{sleap_nn-0.0.5.dist-info → sleap_nn-0.1.0a0.dist-info}/RECORD

@@ -1,10 +1,11 @@
 sleap_nn/.DS_Store,sha256=HY8amA79eHkt7o5VUiNsMxkc9YwW6WIPyZbYRj_JdSU,6148
-sleap_nn/__init__.py,sha256=9Duogtzl1UgBKtc92YBcgpY264Hn1KEnat_HRc7vKUw,1213
-sleap_nn/cli.py,sha256=O4yKepcqarJfS-oz2zYG6DAaDRFTplcfn_Cj10xgUuc,16196
-sleap_nn/evaluation.py,sha256=yfkgiBXi6CnRYdJtPe6_0EJDXmFpB_eBNs5np_hfWwk,28494
+sleap_nn/__init__.py,sha256=DzQeiZIFUmfhpf6mk4j1AKAY2bofVMyIa31xbiSu-ls,1317
+sleap_nn/cli.py,sha256=U4hpEcOxK7a92GeItY95E2DRm5P1ME1GqU__mxaDcW0,21167
+sleap_nn/evaluation.py,sha256=3u7y85wFoBgCwOB2xOGTJIDrd2dUPWOo4m0s0oW3da4,31095
 sleap_nn/legacy_models.py,sha256=8aGK30DZv3pW2IKDBEWH1G2mrytjaxPQD4miPUehj0M,20258
-sleap_nn/predict.py,sha256=A7vMZOMlznolrZojJCBypwYMkSxKw8RDjM7TVWdWUvQ,27813
-sleap_nn/train.py,sha256=RjTXGf_KhXW3I-th1bkkP7ir5jpF__OQey-6Ty3dLOg,25799
+sleap_nn/predict.py,sha256=8QKjRbS-L-6HF1NFJWioBPv3HSzUpFr2oGEB5hRJzQA,35523
+sleap_nn/system_info.py,sha256=7tWe3y6s872nDbrZoHIdSs-w4w46Z4dEV2qCV-Fe7No,14711
+sleap_nn/train.py,sha256=fWx_b1HqkadQ-GM_VEM1frCd8WkzJLqRARBNn8UoUbo,27181
 sleap_nn/architectures/__init__.py,sha256=w0XxQcx-CYyooszzvxRkKWiJkUg-26IlwQoGna8gn40,46
 sleap_nn/architectures/common.py,sha256=MLv-zdHsWL5Q2ct_Wv6SQbRS-5hrFtjK_pvBEfwx-vU,3660
 sleap_nn/architectures/convnext.py,sha256=l9lMJDxIMb-9MI3ShOtVwbOUMuwOLtSQlxiVyYHqjvE,13953
@@ -15,49 +16,50 @@ sleap_nn/architectures/swint.py,sha256=S66Wd0j8Hp-rGlv1C60WSw3AwGyAyGetgfwpL0nIK
 sleap_nn/architectures/unet.py,sha256=rAy2Omi6tv1MNW2nBn0Tw-94Nw_-1wFfCT3-IUyPcgo,11723
 sleap_nn/architectures/utils.py,sha256=L0KVs0gbtG8U75Sl40oH_r_w2ySawh3oQPqIGi54HGo,2171
 sleap_nn/config/__init__.py,sha256=l0xV1uJsGJfMPfWAqlUR7Ivu4cSCWsP-3Y9ueyPESuk,42
-sleap_nn/config/data_config.py,sha256=bf-JTdQR0JQNJERnllvEVtHn2T1YgkkaVgxtGg25FUg,21489
-sleap_nn/config/get_config.py,sha256=42SHD4GWHiW5O8uaQJ5oe9tMMCf7_P77tTWvHtPh72g,41628
+sleap_nn/config/data_config.py,sha256=5a5YlXm4V9qGvkqgFNy6o0XJ_Q06UFjpYJXmNHfvXEI,24021
+sleap_nn/config/get_config.py,sha256=vN_aOPTj9F-QBqGGfVSv8_aFSAYl-RfXY0pdbdcqjcM,42021
 sleap_nn/config/model_config.py,sha256=XFIbqFno7IkX0Se5WF_2_7aUalAlC2SvpDe-uP2TttM,57582
-sleap_nn/config/trainer_config.py,sha256=zoqGwzsmjuwYUPZyqe76I64QAzEL7phwpMg3FItDp2w,27282
+sleap_nn/config/trainer_config.py,sha256=PaoNtRSNc2xgzwN955aR9kTZL8IxCWdevGljLxS6jOk,28073
 sleap_nn/config/training_job_config.py,sha256=v12_ME_tBUg8JFwOxJNW4sDQn-SedDhiJOGz-TlRwT0,5861
 sleap_nn/config/utils.py,sha256=GgWgVs7_N7ifsJ5OQG3_EyOagNyN3Dx7wS2BAlkaRkg,5553
 sleap_nn/data/__init__.py,sha256=eMNvFJFa3gv5Rq8oK5wzo6zt1pOlwUGYf8EQii6bq7c,54
-sleap_nn/data/augmentation.py,sha256=9A9eDZrLfaQNX3-kwcJ5bTlagVNKHaG9pc5-9n0nE1A,11438
+sleap_nn/data/augmentation.py,sha256=Kqw_DayPth_DBsmaO1G8Voou_-cYZuSPOjSQWSajgRI,13618
 sleap_nn/data/confidence_maps.py,sha256=PTRqZWSAz1S7viJhxu7QgIC1aHiek97c_dCUsKUwG1o,6217
-sleap_nn/data/custom_datasets.py,sha256=ZgqBUJUtA1WaG3sIKKjr3s-F4_ovqUqJI_xEjJpNEf4,98218
+sleap_nn/data/custom_datasets.py,sha256=2qAyLeiCPI9uudFFP7zlj6d_tbxc5OVzpnzT23mRkVw,98472
 sleap_nn/data/edge_maps.py,sha256=75qG_7zHRw7fC8JUCVI2tzYakIoxxneWWmcrTwjcHPo,12519
 sleap_nn/data/identity.py,sha256=7vNup6PudST4yDLyDT9wDO-cunRirTEvx4sP77xrlfk,5193
 sleap_nn/data/instance_centroids.py,sha256=SF-3zJt_VMTbZI5ssbrvmZQZDd3684bn55EAtvcbQ6o,2172
-sleap_nn/data/instance_cropping.py,sha256=I0Maq2EJvZkbfAAkz0GwgFfJI0X-1cfBFYhowIFvEyY,5816
-sleap_nn/data/normalization.py,sha256=M3DWPVY9ZU84NocXE4enbdtS3BxQLKd7T0gwkZZVh6M,1231
-sleap_nn/data/providers.py,sha256=skQL_ftKIJeFXgaOka4rPVC09Yb5Ob79BzWw9edBtQU,13287
+sleap_nn/data/instance_cropping.py,sha256=2dYq5OTwkFN1PdMjoxyuMuHq1OEe03m3Vzqvcs_dkPE,8304
+sleap_nn/data/normalization.py,sha256=5xEvcguG-fvAGObl4nWPZ9TEM5gvv0uYPGDuni34XII,2930
+sleap_nn/data/providers.py,sha256=0x6GFP1s1c08ji4p0M5V6p-dhT4Z9c-SI_Aw1DWX-uM,14272
 sleap_nn/data/resizing.py,sha256=YFpSQduIBkRK39FYmrqDL-v8zMySlEs6TJxh6zb_0ZU,5076
 sleap_nn/data/utils.py,sha256=rT0w7KMOTlzaeKWq1TqjbgC4Lvjz_G96McllvEOqXx8,5641
-sleap_nn/inference/__init__.py,sha256=8RXWgwPcNL4CMJSJ4rRaRF40CZkc93fK6zCoSpCM5c0,33
+sleap_nn/inference/__init__.py,sha256=eVkCmKrxHlDFJIlZTf8B5XEOcSyw-gPQymXMY5uShOM,170
 sleap_nn/inference/bottomup.py,sha256=NqN-G8TzAOsvCoL3bttEjA1iGsuveLOnOCXIUeFCdSA,13684
 sleap_nn/inference/identity.py,sha256=GjNDL9MfGqNyQaK4AE8JQCAE8gpMuE_Y-3r3Gpa53CE,6540
 sleap_nn/inference/paf_grouping.py,sha256=7Fo9lCAj-zcHgv5rI5LIMYGcixCGNt_ZbSNs8Dik7l8,69973
-sleap_nn/inference/peak_finding.py,sha256=uS4mb_T5ozjiCL5KhWRiiphRA50J3WXvTuJGRjPa29I,12722
-sleap_nn/inference/predictors.py,sha256=15x9n4zFgjXXtm_fAmJtkjWgPoWnjXOUmqryN1qTmhM,153223
+sleap_nn/inference/peak_finding.py,sha256=L9LdYKt_Bfw7cxo6xEpgF8wXcZAwq5plCfmKJ839N40,13014
+sleap_nn/inference/predictors.py,sha256=U114RlgOXKGm5iz1lnTfE3aN9S0WCh6gWhVP3KVewfc,158046
+sleap_nn/inference/provenance.py,sha256=0BekXyvpLMb0Vv6DjpctlLduG9RN-Q8jt5zDm783eZE,11204
 sleap_nn/inference/single_instance.py,sha256=rOns_5TsJ1rb-lwmHG3ZY-pOhXGN2D-SfW9RmBxxzcI,4089
-sleap_nn/inference/topdown.py,sha256=LX1o7HQmIinvTde75B0n3dcX4Gr8wX8nxMY9UF38l0Y,34872
+sleap_nn/inference/topdown.py,sha256=Ha0Nwx-XCH_rebIuIGhP0qW68QpjLB3XRr9rxt05JLs,35108
 sleap_nn/inference/utils.py,sha256=JnaJK4S_qLtHkWOSkHf4oRZjOmgnU9BGADQnntgGxxs,4689
 sleap_nn/tracking/__init__.py,sha256=rGR35wpSW-n5d3cMiQUzQQ_Dy5II5DPjlXAoPw2QhmM,31
 sleap_nn/tracking/track_instance.py,sha256=9k0uVy9VmpleaLcJh7sVWSeFUPXiw7yj95EYNdXJcks,1373
-sleap_nn/tracking/tracker.py,sha256=eAaFFzu5EdvC_lOgM18Wo9IdpbDPU_dMa4zq0Sa0U7s,41625
+sleap_nn/tracking/tracker.py,sha256=_WT-HFruzyOsvcq3AtLm3vnI9MYSwyBmq-HlQvj1vmU,41955
 sleap_nn/tracking/utils.py,sha256=uHVd_mzzZjviVDdLSKXJJ1T96n5ObKvkqIuGsl9Yy8U,11276
 sleap_nn/tracking/candidates/__init__.py,sha256=1O7NObIwshM7j1rLHmImbFphvkM9wY1j4j1TvO5scSE,49
 sleap_nn/tracking/candidates/fixed_window.py,sha256=D80KMlTnenuQveQVVhk9j0G8yx6K324C7nMLHgG76e0,6296
 sleap_nn/tracking/candidates/local_queues.py,sha256=Nx3R5wwEwq0gbfH-fi3oOumfkQo8_sYe5GN47pD9Be8,7305
 sleap_nn/training/__init__.py,sha256=vNTKsIJPZHJwFSKn5PmjiiRJunR_9e7y4_v0S6rdF8U,32
-sleap_nn/training/callbacks.py,sha256=eIRSFMI65tkrrPBXmD7lzdvr_lBntaoPx6i2WJhi0iU,12576
-sleap_nn/training/lightning_modules.py,sha256=XIWxDdeZdM1G0oVmdsMKiuyrh0TYsgqTiJhqEs9hFYU,80852
+sleap_nn/training/callbacks.py,sha256=TVnQ6plNC2MnlTiY2rSCRuw2WRk5cQSziek_VPUcOEg,25994
+sleap_nn/training/lightning_modules.py,sha256=G3c4xJkYWW-iSRawzkgTqkGd4lTsbPiMTcB5Nvq7jes,85512
 sleap_nn/training/losses.py,sha256=gbdinUURh4QUzjmNd2UJpt4FXwecqKy9gHr65JZ1bZk,1632
-sleap_nn/training/model_trainer.py,sha256=G7n2a4D4ztpCmsYrpORD084hhYL0BoX-0_U_N58Ag8k,44237
-sleap_nn/training/utils.py,sha256=WSUpB793hTj0rd9Pc1xcUrQ8rEKnXt_7UOpt3OFRNvo,7252
-sleap_nn-0.0.5.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-sleap_nn-0.0.5.dist-info/METADATA,sha256=TRaCfutydVkWdlamLhmrfD7k0RQ5fPiJv0GkhgzT3uA,6232
-sleap_nn-0.0.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-sleap_nn-0.0.5.dist-info/entry_points.txt,sha256=zfl5Y3hidZxWBvo8qXvu5piJAXJ_l6v7xVFm0gNiUoI,46
-sleap_nn-0.0.5.dist-info/top_level.txt,sha256=Kz68iQ55K75LWgSeqz4V4SCMGeFFYH-KGBOyhQh3xZE,9
-sleap_nn-0.0.5.dist-info/RECORD,,
+sleap_nn/training/model_trainer.py,sha256=InDKHrQxBwbltZKutW4yrBR9NThLdRpWNUGhmB0xAi4,57863
+sleap_nn/training/utils.py,sha256=ivdkZEI0DkTCm6NPszsaDOh9jSfozkONZdl6TvvQUWI,20398
+sleap_nn-0.1.0a0.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+sleap_nn-0.1.0a0.dist-info/METADATA,sha256=lxSmGNTUg9eetqHCvhw8Tv5zJtia-dIM5RzOeoDccj8,5637
+sleap_nn-0.1.0a0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+sleap_nn-0.1.0a0.dist-info/entry_points.txt,sha256=zfl5Y3hidZxWBvo8qXvu5piJAXJ_l6v7xVFm0gNiUoI,46
+sleap_nn-0.1.0a0.dist-info/top_level.txt,sha256=Kz68iQ55K75LWgSeqz4V4SCMGeFFYH-KGBOyhQh3xZE,9
+sleap_nn-0.1.0a0.dist-info/RECORD,,