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

sleap_nn/training/model_trainer.py

@@ -2,6 +2,7 @@
 
 import os
 import shutil
+import copy
 import attrs
 import torch
 import random
@@ -15,13 +16,14 @@ import yaml
 from pathlib import Path
 from typing import List, Optional
 from datetime import datetime
-from itertools import count
+from itertools import cycle, count
 from omegaconf import DictConfig, OmegaConf
 from lightning.pytorch.loggers import WandbLogger
 from sleap_nn.data.utils import check_cache_memory
 from lightning.pytorch.callbacks import (
     ModelCheckpoint,
     EarlyStopping,
+    LearningRateMonitor,
 )
 from lightning.pytorch.profilers import (
     SimpleProfiler,
@@ -53,11 +55,12 @@ from sleap_nn.config.training_job_config import verify_training_cfg
 from sleap_nn.training.callbacks import (
     ProgressReporterZMQ,
     TrainingControllerZMQ,
+    MatplotlibSaver,
+    WandBPredImageLogger,
+    WandBVizCallback,
+    WandBVizCallbackWithPAFs,
     CSVLoggerCallback,
     SleapProgressBar,
-    EpochEndEvaluationCallback,
-    CentroidEvaluationCallback,
-    UnifiedVizCallback,
 )
 from sleap_nn import RANK
 from sleap_nn.legacy_models import get_keras_first_layer_channels
@@ -487,36 +490,16 @@ class ModelTrainer:
             ckpt_dir = "."
             self.config.trainer_config.ckpt_dir = ckpt_dir
         run_name = self.config.trainer_config.run_name
-        run_name_is_empty = run_name is None or run_name == "" or run_name == "None"
-
-        # Validate: multi-GPU + disk cache requires explicit run_name
-        if run_name_is_empty:
-            is_disk_caching = (
-                self.config.data_config.data_pipeline_fw
-                == "torch_dataset_cache_img_disk"
-            )
-            num_devices = self._get_trainer_devices()
-
-            if is_disk_caching and num_devices > 1:
-                raise ValueError(
-                    f"Multi-GPU training with disk caching requires an explicit `run_name`.\n\n"
-                    f"Detected {num_devices} device(s) with "
-                    f"`data_pipeline_fw='torch_dataset_cache_img_disk'`.\n"
-                    f"Without an explicit run_name, each GPU worker generates a different "
-                    f"timestamp-based directory, causing cache synchronization failures.\n\n"
-                    f"Please provide a run_name using one of these methods:\n"
-                    f"  - CLI: sleap-nn train config.yaml trainer_config.run_name=my_experiment\n"
-                    f"  - Config file: Set `trainer_config.run_name: my_experiment`\n"
-                    f"  - Python API: train(..., run_name='my_experiment')"
-                )
-
-            # Auto-generate timestamp-based run_name (safe for single GPU or non-disk-cache)
+        if run_name is None or run_name == "" or run_name == "None":
             sum_train_lfs = sum([len(train_label) for train_label in self.train_labels])
             sum_val_lfs = sum([len(val_label) for val_label in self.val_labels])
-            run_name = (
-                datetime.now().strftime("%y%m%d_%H%M%S")
-                + f".{self.model_type}.n={sum_train_lfs + sum_val_lfs}"
-            )
+            if self._get_trainer_devices() > 1:
+                run_name = f"{self.model_type}.n={sum_train_lfs + sum_val_lfs}"
+            else:
+                run_name = (
+                    datetime.now().strftime("%y%m%d_%H%M%S")
+                    + f".{self.model_type}.n={sum_train_lfs + sum_val_lfs}"
+                )
 
         # If checkpoint path already exists, add suffix to prevent overwriting
         if (Path(ckpt_dir) / run_name).exists() and (
@@ -659,6 +642,10 @@ class ModelTrainer:
         if self.config.trainer_config.wandb.prv_runid == "":
             self.config.trainer_config.wandb.prv_runid = None
 
+        # Default wandb run name to trainer run_name if not specified
+        if self.config.trainer_config.wandb.name is None:
+            self.config.trainer_config.wandb.name = self.config.trainer_config.run_name
+
         # compute preprocessing parameters from the labels objects and fill in the config
         self._setup_preprocessing_config()
 
@@ -708,14 +695,9 @@ class ModelTrainer:
                 )
             )
 
-        # setup checkpoint path (generates run_name if not specified)
+        # setup checkpoint path
         self._setup_ckpt_path()
 
-        # Default wandb run name to trainer run_name if not specified
-        # Note: This must come after _setup_ckpt_path() which generates run_name
-        if self.config.trainer_config.wandb.name is None:
-            self.config.trainer_config.wandb.name = self.config.trainer_config.run_name
-
         # verify input_channels in model_config based on input image and pretrained model weights
         self._verify_model_input_channels()
 
@@ -727,15 +709,15 @@ class ModelTrainer:
         ).as_posix()
         logger.info(f"Setting up model ckpt dir: `{ckpt_path}`...")
 
-        # Only rank 0 (or non-distributed) should create directories and save files
+        if not Path(ckpt_path).exists():
+            try:
+                Path(ckpt_path).mkdir(parents=True, exist_ok=True)
+            except OSError as e:
+                message = f"Cannot create a new folder in {ckpt_path}.\n {e}"
+                logger.error(message)
+                raise OSError(message)
+
         if RANK in [0, -1]:
-            if not Path(ckpt_path).exists():
-                try:
-                    Path(ckpt_path).mkdir(parents=True, exist_ok=True)
-                except OSError as e:
-                    message = f"Cannot create a new folder in {ckpt_path}.\n {e}"
-                    logger.error(message)
-                    raise OSError(message)
             # Check if we should filter to user-labeled frames only
             user_instances_only = OmegaConf.select(
                 self.config, "data_config.user_instances_only", default=True
@@ -808,40 +790,10 @@ class ModelTrainer:
         base_cache_img_path = None
         if self.config.data_config.data_pipeline_fw == "torch_dataset_cache_img_memory":
             # check available memory. If insufficient memory, default to disk caching.
-            # Account for DataLoader worker memory overhead
-            train_num_workers = self.config.trainer_config.train_data_loader.num_workers
-            val_num_workers = self.config.trainer_config.val_data_loader.num_workers
-            max_num_workers = max(train_num_workers, val_num_workers)
-
             mem_available = check_cache_memory(
-                self.train_labels,
-                self.val_labels,
-                memory_buffer=MEMORY_BUFFER,
-                num_workers=max_num_workers,
+                self.train_labels, self.val_labels, memory_buffer=MEMORY_BUFFER
             )
             if not mem_available:
-                # Validate: multi-GPU + auto-generated run_name + fallback to disk cache
-                original_run_name = self._initial_config.trainer_config.run_name
-                run_name_was_auto = (
-                    original_run_name is None
-                    or original_run_name == ""
-                    or original_run_name == "None"
-                )
-                if run_name_was_auto and self.trainer.num_devices > 1:
-                    raise ValueError(
-                        f"Memory caching failed and disk caching fallback requires an "
-                        f"explicit `run_name` for multi-GPU training.\n\n"
-                        f"Detected {self.trainer.num_devices} device(s) with insufficient "
-                        f"memory for in-memory caching.\n"
-                        f"Without an explicit run_name, each GPU worker generates a different "
-                        f"timestamp-based directory, causing cache synchronization failures.\n\n"
-                        f"Please provide a run_name using one of these methods:\n"
-                        f"  - CLI: sleap-nn train config.yaml trainer_config.run_name=my_experiment\n"
-                        f"  - Config file: Set `trainer_config.run_name: my_experiment`\n"
-                        f"  - Python API: train(..., run_name='my_experiment')\n\n"
-                        f"Alternatively, use `data_pipeline_fw='torch_dataset'` to disable caching."
-                    )
-
                 self.config.data_config.data_pipeline_fw = (
                     "torch_dataset_cache_img_disk"
                 )
@@ -885,7 +837,7 @@ class ModelTrainer:
                     / self.config.trainer_config.run_name
                 ).as_posix(),
                 filename="best",
-                monitor="val/loss",
+                monitor="val_loss",
                 mode="min",
             )
             callbacks.append(checkpoint_callback)
@@ -893,52 +845,18 @@ class ModelTrainer:
             # csv log callback
             csv_log_keys = [
                 "epoch",
-                "train/loss",
-                "val/loss",
+                "train_loss",
+                "val_loss",
                 "learning_rate",
-                "train/time",
-                "val/time",
+                "train_time",
+                "val_time",
             ]
-            # Add model-specific keys for wandb parity
             if self.model_type in [
                 "single_instance",
                 "centered_instance",
                 "multi_class_topdown",
             ]:
-                csv_log_keys.extend(
-                    [f"train/confmaps/{name}" for name in self.skeletons[0].node_names]
-                )
-            if self.model_type == "bottomup":
-                csv_log_keys.extend(
-                    [
-                        "train/confmaps_loss",
-                        "train/paf_loss",
-                        "val/confmaps_loss",
-                        "val/paf_loss",
-                    ]
-                )
-            if self.model_type == "multi_class_bottomup":
-                csv_log_keys.extend(
-                    [
-                        "train/confmaps_loss",
-                        "train/classmap_loss",
-                        "train/class_accuracy",
-                        "val/confmaps_loss",
-                        "val/classmap_loss",
-                        "val/class_accuracy",
-                    ]
-                )
-            if self.model_type == "multi_class_topdown":
-                csv_log_keys.extend(
-                    [
-                        "train/confmaps_loss",
-                        "train/classvector_loss",
-                        "train/class_accuracy",
-                        "val/confmaps_loss",
-                        "val/classvector_loss",
-                        "val/class_accuracy",
-                    ]
-                )
+                csv_log_keys.extend(self.skeletons[0].node_names)
             csv_logger = CSVLoggerCallback(
                 filepath=Path(self.config.trainer_config.ckpt_dir)
                 / self.config.trainer_config.run_name
@@ -951,7 +869,7 @@ class ModelTrainer:
             # early stopping callback
             callbacks.append(
                 EarlyStopping(
-                    monitor="val/loss",
+                    monitor="val_loss",
                     mode="min",
                     verbose=False,
                     min_delta=self.config.trainer_config.early_stopping.min_delta,
@@ -980,16 +898,9 @@ class ModelTrainer:
             )
             loggers.append(wandb_logger)
 
-            # Log message about wandb local logs cleanup
-            should_delete_wandb_logs = wandb_config.delete_local_logs is True or (
-                wandb_config.delete_local_logs is None
-                and wandb_config.wandb_mode != "offline"
-            )
-            if should_delete_wandb_logs:
-                logger.info(
-                    "WandB local logs will be deleted after training completes. "
-                    "To keep logs, set trainer_config.wandb.delete_local_logs=false"
-                )
+            # Learning rate monitor callback - logs LR at each step for dynamic schedulers
+            # Only added when wandb is enabled since it requires a logger
+            callbacks.append(LearningRateMonitor(logging_interval="step"))
 
             # save the configs as yaml in the checkpoint dir
             # Mask API key in both configs to prevent saving to disk
@@ -1009,8 +920,11 @@ class ModelTrainer:
             )
             callbacks.append(ProgressReporterZMQ(address=publish_address))
 
-        # viz callbacks - use unified callback for all visualization outputs
+        # viz callbacks
         if self.config.trainer_config.visualize_preds_during_training:
+            train_viz_pipeline = cycle(viz_train_dataset)
+            val_viz_pipeline = cycle(viz_val_dataset)
+
             viz_dir = (
                 Path(self.config.trainer_config.ckpt_dir)
                 / self.config.trainer_config.run_name
@@ -1020,77 +934,147 @@ class ModelTrainer:
                 if RANK in [0, -1]:
                     Path(viz_dir).mkdir(parents=True, exist_ok=True)
 
-            # Get wandb viz config options
-            log_wandb = self.config.trainer_config.use_wandb and OmegaConf.select(
-                self.config, "trainer_config.wandb.save_viz_imgs_wandb", default=False
-            )
-            wandb_modes = []
-            if log_wandb:
-                if OmegaConf.select(
-                    self.config, "trainer_config.wandb.viz_enabled", default=True
-                ):
-                    wandb_modes.append("direct")
-                if OmegaConf.select(
-                    self.config, "trainer_config.wandb.viz_boxes", default=False
-                ):
-                    wandb_modes.append("boxes")
-                if OmegaConf.select(
-                    self.config, "trainer_config.wandb.viz_masks", default=False
-                ):
-                    wandb_modes.append("masks")
-
-            # Single unified callback handles all visualization outputs
             callbacks.append(
-                UnifiedVizCallback(
-                    model_trainer=self,
-                    train_dataset=viz_train_dataset,
-                    val_dataset=viz_val_dataset,
-                    model_type=self.model_type,
-                    save_local=self.config.trainer_config.save_ckpt,
-                    local_save_dir=viz_dir,
-                    log_wandb=log_wandb,
-                    wandb_modes=wandb_modes if wandb_modes else ["direct"],
-                    wandb_box_size=OmegaConf.select(
-                        self.config, "trainer_config.wandb.viz_box_size", default=5.0
-                    ),
-                    wandb_confmap_threshold=OmegaConf.select(
-                        self.config,
-                        "trainer_config.wandb.viz_confmap_threshold",
-                        default=0.1,
+                MatplotlibSaver(
+                    save_folder=viz_dir,
+                    plot_fn=lambda: self.lightning_model.visualize_example(
+                        next(train_viz_pipeline)
                     ),
-                    log_wandb_table=OmegaConf.select(
-                        self.config, "trainer_config.wandb.log_viz_table", default=False
+                    prefix="train",
+                )
+            )
+            callbacks.append(
+                MatplotlibSaver(
+                    save_folder=viz_dir,
+                    plot_fn=lambda: self.lightning_model.visualize_example(
+                        next(val_viz_pipeline)
                     ),
+                    prefix="validation",
                 )
             )
 
-        # Add custom progress bar with better metric formatting
-        if self.config.trainer_config.enable_progress_bar:
-            callbacks.append(SleapProgressBar())
+            if self.model_type == "bottomup":
+                train_viz_pipeline1 = cycle(copy.deepcopy(viz_train_dataset))
+                val_viz_pipeline1 = cycle(copy.deepcopy(viz_val_dataset))
+                callbacks.append(
+                    MatplotlibSaver(
+                        save_folder=viz_dir,
+                        plot_fn=lambda: self.lightning_model.visualize_pafs_example(
+                            next(train_viz_pipeline1)
+                        ),
+                        prefix="train.pafs_magnitude",
+                    )
+                )
+                callbacks.append(
+                    MatplotlibSaver(
+                        save_folder=viz_dir,
+                        plot_fn=lambda: self.lightning_model.visualize_pafs_example(
+                            next(val_viz_pipeline1)
+                        ),
+                        prefix="validation.pafs_magnitude",
+                    )
+                )
 
-        # Add epoch-end evaluation callback if enabled
-        if self.config.trainer_config.eval.enabled:
-            if self.model_type == "centroid":
-                # Use centroid-specific evaluation with distance-based metrics
+            if self.model_type == "multi_class_bottomup":
+                train_viz_pipeline1 = cycle(copy.deepcopy(viz_train_dataset))
+                val_viz_pipeline1 = cycle(copy.deepcopy(viz_val_dataset))
                 callbacks.append(
-                    CentroidEvaluationCallback(
-                        videos=self.val_labels[0].videos,
-                        eval_frequency=self.config.trainer_config.eval.frequency,
-                        match_threshold=self.config.trainer_config.eval.match_threshold,
+                    MatplotlibSaver(
+                        save_folder=viz_dir,
+                        plot_fn=lambda: self.lightning_model.visualize_class_maps_example(
+                            next(train_viz_pipeline1)
+                        ),
+                        prefix="train.class_maps",
                     )
                 )
-            else:
-                # Use standard OKS/PCK evaluation for pose models
                 callbacks.append(
-                    EpochEndEvaluationCallback(
-                        skeleton=self.skeletons[0],
-                        videos=self.val_labels[0].videos,
-                        eval_frequency=self.config.trainer_config.eval.frequency,
-                        oks_stddev=self.config.trainer_config.eval.oks_stddev,
-                        oks_scale=self.config.trainer_config.eval.oks_scale,
+                    MatplotlibSaver(
+                        save_folder=viz_dir,
+                        plot_fn=lambda: self.lightning_model.visualize_class_maps_example(
+                            next(val_viz_pipeline1)
+                        ),
+                        prefix="validation.class_maps",
                     )
                 )
 
+            if self.config.trainer_config.use_wandb and OmegaConf.select(
+                self.config, "trainer_config.wandb.save_viz_imgs_wandb", default=False
+            ):
+                # Get wandb viz config options
+                viz_enabled = OmegaConf.select(
+                    self.config, "trainer_config.wandb.viz_enabled", default=True
+                )
+                viz_boxes = OmegaConf.select(
+                    self.config, "trainer_config.wandb.viz_boxes", default=False
+                )
+                viz_masks = OmegaConf.select(
+                    self.config, "trainer_config.wandb.viz_masks", default=False
+                )
+                viz_box_size = OmegaConf.select(
+                    self.config, "trainer_config.wandb.viz_box_size", default=5.0
+                )
+                viz_confmap_threshold = OmegaConf.select(
+                    self.config,
+                    "trainer_config.wandb.viz_confmap_threshold",
+                    default=0.1,
+                )
+                log_viz_table = OmegaConf.select(
+                    self.config, "trainer_config.wandb.log_viz_table", default=False
+                )
+
+                # Create viz data pipelines for wandb callback
+                wandb_train_viz_pipeline = cycle(copy.deepcopy(viz_train_dataset))
+                wandb_val_viz_pipeline = cycle(copy.deepcopy(viz_val_dataset))
+
+                if self.model_type == "bottomup":
+                    # Bottom-up model needs PAF visualizations
+                    wandb_train_pafs_pipeline = cycle(copy.deepcopy(viz_train_dataset))
+                    wandb_val_pafs_pipeline = cycle(copy.deepcopy(viz_val_dataset))
+                    callbacks.append(
+                        WandBVizCallbackWithPAFs(
+                            train_viz_fn=lambda: self.lightning_model.get_visualization_data(
+                                next(wandb_train_viz_pipeline)
+                            ),
+                            val_viz_fn=lambda: self.lightning_model.get_visualization_data(
+                                next(wandb_val_viz_pipeline)
+                            ),
+                            train_pafs_viz_fn=lambda: self.lightning_model.get_visualization_data(
+                                next(wandb_train_pafs_pipeline), include_pafs=True
+                            ),
+                            val_pafs_viz_fn=lambda: self.lightning_model.get_visualization_data(
+                                next(wandb_val_pafs_pipeline), include_pafs=True
+                            ),
+                            viz_enabled=viz_enabled,
+                            viz_boxes=viz_boxes,
+                            viz_masks=viz_masks,
+                            box_size=viz_box_size,
+                            confmap_threshold=viz_confmap_threshold,
+                            log_table=log_viz_table,
+                        )
+                    )
+                else:
+                    # Standard models
+                    callbacks.append(
+                        WandBVizCallback(
+                            train_viz_fn=lambda: self.lightning_model.get_visualization_data(
+                                next(wandb_train_viz_pipeline)
+                            ),
+                            val_viz_fn=lambda: self.lightning_model.get_visualization_data(
+                                next(wandb_val_viz_pipeline)
+                            ),
+                            viz_enabled=viz_enabled,
+                            viz_boxes=viz_boxes,
+                            viz_masks=viz_masks,
+                            box_size=viz_box_size,
+                            confmap_threshold=viz_confmap_threshold,
+                            log_table=log_viz_table,
+                        )
+                    )
+
+        # Add custom progress bar with better metric formatting
+        if self.config.trainer_config.enable_progress_bar:
+            callbacks.append(SleapProgressBar())
+
         return loggers, callbacks
 
     def _delete_cache_imgs(self):
@@ -1179,11 +1163,6 @@ class ModelTrainer:
                     : self.config.trainer_config.trainer_devices
                 ]
             ]
-            # Sort device indices in ascending order for NCCL compatibility.
-            # NCCL expects devices in consistent ascending order across ranks
-            # to properly set up communication rings. Without sorting, DDP may
-            # assign multiple ranks to the same GPU, causing "Duplicate GPU detected" errors.
-            devices.sort()
             logger.info(f"Using GPUs with most available memory: {devices}")
 
         # create lightning.Trainer instance.
@@ -1205,10 +1184,6 @@ class ModelTrainer:
         # setup datasets
         train_dataset, val_dataset = self._setup_datasets()
 
-        # Barrier after dataset creation to ensure all workers wait for disk caching
-        # (rank 0 caches to disk, others must wait before reading cached files)
-        self.trainer.strategy.barrier()
-
         # set-up steps per epoch
         train_steps_per_epoch = self.config.trainer_config.train_steps_per_epoch
         if train_steps_per_epoch is None:
@@ -1282,21 +1257,18 @@ class ModelTrainer:
                 # Define custom x-axes for wandb metrics
                 # Epoch-level metrics use epoch as x-axis, step-level use default global_step
                 wandb.define_metric("epoch")
-
-                # Training metrics (train/ prefix for grouping) - all use epoch x-axis
-                wandb.define_metric("train/*", step_metric="epoch")
-                wandb.define_metric("train/confmaps/*", step_metric="epoch")
-
-                # Validation metrics (val/ prefix for grouping)
-                wandb.define_metric("val/*", step_metric="epoch")
-
-                # Evaluation metrics (eval/ prefix for grouping)
-                wandb.define_metric("eval/*", step_metric="epoch")
-
-                # Visualization images (need explicit nested paths)
-                wandb.define_metric("viz/*", step_metric="epoch")
-                wandb.define_metric("viz/train/*", step_metric="epoch")
-                wandb.define_metric("viz/val/*", step_metric="epoch")
+                wandb.define_metric("val_loss", step_metric="epoch")
+                wandb.define_metric("val_time", step_metric="epoch")
+                wandb.define_metric("train_time", step_metric="epoch")
+                # Per-node losses use epoch as x-axis
+                for node_name in self.skeletons[0].node_names:
+                    wandb.define_metric(node_name, step_metric="epoch")
+
+                # Visualization images use epoch as x-axis
+                wandb.define_metric("train_predictions*", step_metric="epoch")
+                wandb.define_metric("val_predictions*", step_metric="epoch")
+                wandb.define_metric("train_pafs*", step_metric="epoch")
+                wandb.define_metric("val_pafs*", step_metric="epoch")
 
                 self.config.trainer_config.wandb.current_run_id = wandb.run.id
                 wandb.config["run_name"] = self.config.trainer_config.wandb.name
@@ -1339,7 +1311,8 @@ class ModelTrainer:
             logger.info(
                 f"Finished training loop. [{(time.time() - start_train_time) / 60:.1f} min]"
             )
-            # Note: wandb.finish() is called in train.py after post-training evaluation
+            if self.trainer.global_rank == 0 and self.config.trainer_config.use_wandb:
+                wandb.finish()
 
             # delete image disk caching
             if (

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sleap-nn
-Version: 0.1.0
+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,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.7.0,>=0.6.2
+Requires-Dist: sleap-io<0.7.0,>=0.6.0
 Requires-Dist: numpy
 Requires-Dist: lightning
-Requires-Dist: skia-python>=87.0
+Requires-Dist: kornia
 Requires-Dist: jsonpickle
 Requires-Dist: scipy
 Requires-Dist: attrs
@@ -32,7 +32,6 @@ 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.20.0; extra == "torch"
@@ -48,17 +47,6 @@ 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

sleap_nn-0.1.0a0.dist-info/RECORD

@@ -0,0 +1,65 @@
+sleap_nn/.DS_Store,sha256=HY8amA79eHkt7o5VUiNsMxkc9YwW6WIPyZbYRj_JdSU,6148
+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=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
+sleap_nn/architectures/encoder_decoder.py,sha256=f3DUFJo6RrIUposdC3Ytyblr5J0tAeZ_si9dm_m_PhM,28339
+sleap_nn/architectures/heads.py,sha256=5E-7kQ-b2gsL0EviQ8z3KS1DAAMT4F2ZnEzx7eSG5gg,21001
+sleap_nn/architectures/model.py,sha256=1_dsP_4T9fsEVJjDt3er0haMKtbeM6w6JC6tc2jD0Gw,7139
+sleap_nn/architectures/swint.py,sha256=S66Wd0j8Hp-rGlv1C60WSw3AwGyAyGetgfwpL0nIK_M,14687
+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=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=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=Kqw_DayPth_DBsmaO1G8Voou_-cYZuSPOjSQWSajgRI,13618
+sleap_nn/data/confidence_maps.py,sha256=PTRqZWSAz1S7viJhxu7QgIC1aHiek97c_dCUsKUwG1o,6217
+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=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=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=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=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=_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=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=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,,