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

sleap_nn/training/schedulers.py

@@ -0,0 +1,191 @@
+"""Custom learning rate schedulers for sleap-nn training.
+
+This module provides learning rate schedulers with warmup phases that are commonly
+used in deep learning for pose estimation and computer vision tasks.
+"""
+
+import math
+from torch.optim.lr_scheduler import LRScheduler
+
+
+class LinearWarmupCosineAnnealingLR(LRScheduler):
+    """Cosine annealing learning rate scheduler with linear warmup.
+
+    The learning rate increases linearly from `warmup_start_lr` to the optimizer's
+    base learning rate over `warmup_epochs`, then decreases following a cosine
+    curve to `eta_min` over the remaining epochs.
+
+    This schedule is widely used in vision transformers and modern CNN architectures
+    as it provides stable early training (warmup) and smooth convergence (cosine decay).
+
+    Args:
+        optimizer: Wrapped optimizer.
+        warmup_epochs: Number of epochs for the linear warmup phase.
+        max_epochs: Total number of training epochs.
+        warmup_start_lr: Learning rate at the start of warmup. Default: 0.0.
+        eta_min: Minimum learning rate at the end of the schedule. Default: 0.0.
+        last_epoch: The index of the last epoch. Default: -1.
+
+    Example:
+        >>> optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+        >>> scheduler = LinearWarmupCosineAnnealingLR(
+        ...     optimizer, warmup_epochs=5, max_epochs=100, eta_min=1e-6
+        ... )
+        >>> for epoch in range(100):
+        ...     train(...)
+        ...     scheduler.step()
+    """
+
+    def __init__(
+        self,
+        optimizer,
+        warmup_epochs: int,
+        max_epochs: int,
+        warmup_start_lr: float = 0.0,
+        eta_min: float = 0.0,
+        last_epoch: int = -1,
+    ):
+        """Initialize the scheduler.
+
+        Args:
+            optimizer: Wrapped optimizer.
+            warmup_epochs: Number of epochs for the linear warmup phase.
+            max_epochs: Total number of training epochs.
+            warmup_start_lr: Learning rate at the start of warmup. Default: 0.0.
+            eta_min: Minimum learning rate at the end of the schedule. Default: 0.0.
+            last_epoch: The index of the last epoch. Default: -1.
+        """
+        if warmup_epochs < 0:
+            raise ValueError(f"warmup_epochs must be >= 0, got {warmup_epochs}")
+        if max_epochs <= 0:
+            raise ValueError(f"max_epochs must be > 0, got {max_epochs}")
+        if warmup_epochs >= max_epochs:
+            raise ValueError(
+                f"warmup_epochs ({warmup_epochs}) must be < max_epochs ({max_epochs})"
+            )
+        if warmup_start_lr < 0:
+            raise ValueError(f"warmup_start_lr must be >= 0, got {warmup_start_lr}")
+        if eta_min < 0:
+            raise ValueError(f"eta_min must be >= 0, got {eta_min}")
+
+        self.warmup_epochs = warmup_epochs
+        self.max_epochs = max_epochs
+        self.warmup_start_lr = warmup_start_lr
+        self.eta_min = eta_min
+        super().__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        """Compute the learning rate at the current epoch."""
+        if self.last_epoch < self.warmup_epochs:
+            # Linear warmup phase
+            if self.warmup_epochs == 0:
+                return list(self.base_lrs)
+            alpha = self.last_epoch / self.warmup_epochs
+            return [
+                self.warmup_start_lr + alpha * (base_lr - self.warmup_start_lr)
+                for base_lr in self.base_lrs
+            ]
+        else:
+            # Cosine annealing phase
+            decay_epochs = self.max_epochs - self.warmup_epochs
+            if decay_epochs == 0:
+                return [self.eta_min for _ in self.base_lrs]
+            progress = (self.last_epoch - self.warmup_epochs) / decay_epochs
+            # Clamp progress to [0, 1] to handle epochs beyond max_epochs
+            progress = min(1.0, progress)
+            return [
+                self.eta_min
+                + (base_lr - self.eta_min) * (1 + math.cos(math.pi * progress)) / 2
+                for base_lr in self.base_lrs
+            ]
+
+
+class LinearWarmupLinearDecayLR(LRScheduler):
+    """Linear warmup followed by linear decay learning rate scheduler.
+
+    The learning rate increases linearly from `warmup_start_lr` to the optimizer's
+    base learning rate over `warmup_epochs`, then decreases linearly to `end_lr`
+    over the remaining epochs.
+
+    This schedule provides a simple, interpretable learning rate trajectory and is
+    commonly used in transformer-based models and NLP tasks.
+
+    Args:
+        optimizer: Wrapped optimizer.
+        warmup_epochs: Number of epochs for the linear warmup phase.
+        max_epochs: Total number of training epochs.
+        warmup_start_lr: Learning rate at the start of warmup. Default: 0.0.
+        end_lr: Learning rate at the end of training. Default: 0.0.
+        last_epoch: The index of the last epoch. Default: -1.
+
+    Example:
+        >>> optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+        >>> scheduler = LinearWarmupLinearDecayLR(
+        ...     optimizer, warmup_epochs=5, max_epochs=100, end_lr=1e-6
+        ... )
+        >>> for epoch in range(100):
+        ...     train(...)
+        ...     scheduler.step()
+    """
+
+    def __init__(
+        self,
+        optimizer,
+        warmup_epochs: int,
+        max_epochs: int,
+        warmup_start_lr: float = 0.0,
+        end_lr: float = 0.0,
+        last_epoch: int = -1,
+    ):
+        """Initialize the scheduler.
+
+        Args:
+            optimizer: Wrapped optimizer.
+            warmup_epochs: Number of epochs for the linear warmup phase.
+            max_epochs: Total number of training epochs.
+            warmup_start_lr: Learning rate at the start of warmup. Default: 0.0.
+            end_lr: Learning rate at the end of training. Default: 0.0.
+            last_epoch: The index of the last epoch. Default: -1.
+        """
+        if warmup_epochs < 0:
+            raise ValueError(f"warmup_epochs must be >= 0, got {warmup_epochs}")
+        if max_epochs <= 0:
+            raise ValueError(f"max_epochs must be > 0, got {max_epochs}")
+        if warmup_epochs >= max_epochs:
+            raise ValueError(
+                f"warmup_epochs ({warmup_epochs}) must be < max_epochs ({max_epochs})"
+            )
+        if warmup_start_lr < 0:
+            raise ValueError(f"warmup_start_lr must be >= 0, got {warmup_start_lr}")
+        if end_lr < 0:
+            raise ValueError(f"end_lr must be >= 0, got {end_lr}")
+
+        self.warmup_epochs = warmup_epochs
+        self.max_epochs = max_epochs
+        self.warmup_start_lr = warmup_start_lr
+        self.end_lr = end_lr
+        super().__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        """Compute the learning rate at the current epoch."""
+        if self.last_epoch < self.warmup_epochs:
+            # Linear warmup phase
+            if self.warmup_epochs == 0:
+                return list(self.base_lrs)
+            alpha = self.last_epoch / self.warmup_epochs
+            return [
+                self.warmup_start_lr + alpha * (base_lr - self.warmup_start_lr)
+                for base_lr in self.base_lrs
+            ]
+        else:
+            # Linear decay phase
+            decay_epochs = self.max_epochs - self.warmup_epochs
+            if decay_epochs == 0:
+                return [self.end_lr for _ in self.base_lrs]
+            progress = (self.last_epoch - self.warmup_epochs) / decay_epochs
+            # Clamp progress to [0, 1] to handle epochs beyond max_epochs
+            progress = min(1.0, progress)
+            return [
+                base_lr + progress * (self.end_lr - base_lr)
+                for base_lr in self.base_lrs
+            ]

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.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sleap-nn
-Version: 0.0.5
+Version: 0.1.0
 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,10 +13,10 @@ 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.2
 Requires-Dist: numpy
 Requires-Dist: lightning
-Requires-Dist: kornia
+Requires-Dist: skia-python>=87.0
 Requires-Dist: jsonpickle
 Requires-Dist: scipy
 Requires-Dist: attrs
@@ -32,37 +32,33 @@ Requires-Dist: hydra-core
 Requires-Dist: jupyter
 Requires-Dist: jupyterlab
 Requires-Dist: pyzmq
+Requires-Dist: rich-click>=1.9.5
 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"
+Provides-Extra: export
+Requires-Dist: onnx>=1.15.0; extra == "export"
+Requires-Dist: onnxruntime>=1.16.0; extra == "export"
+Requires-Dist: onnxscript>=0.1.0; extra == "export"
+Provides-Extra: export-gpu
+Requires-Dist: onnx>=1.15.0; extra == "export-gpu"
+Requires-Dist: onnxruntime-gpu>=1.16.0; extra == "export-gpu"
+Requires-Dist: onnxscript>=0.1.0; extra == "export-gpu"
+Provides-Extra: tensorrt
+Requires-Dist: tensorrt>=10.13.0; (sys_platform == "linux" or sys_platform == "win32") and extra == "tensorrt"
+Requires-Dist: torch-tensorrt>=2.5.0; (sys_platform == "linux" or sys_platform == "win32") and extra == "tensorrt"
 Dynamic: license-file
 
 # sleap-nn
@@ -120,22 +116,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 +154,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`.