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

sleap_nn/system_info.py

@@ -0,0 +1,443 @@
+"""System diagnostics and compatibility checking for sleap-nn."""
+
+import importlib.metadata
+import json
+import platform
+import shutil
+import subprocess
+import sys
+from typing import Optional
+
+# Key packages to check versions for
+PACKAGES = [
+    "sleap-nn",
+    "sleap-io",
+    "torch",
+    "pytorch-lightning",
+    "kornia",
+    "wandb",
+    "numpy",
+    "h5py",
+]
+
+# CUDA version -> (min_driver_linux, min_driver_windows)
+# Source: https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/
+CUDA_DRIVER_REQUIREMENTS = {
+    "12.6": ("560.28.03", "560.76"),
+    "12.8": ("570.26", "570.65"),
+    "13.0": ("580.65.06", "580.00"),
+}
+
+
+def parse_driver_version(version: str) -> tuple[int, ...]:
+    """Parse driver version string into comparable tuple."""
+    try:
+        return tuple(int(x) for x in version.split("."))
+    except ValueError:
+        return (0,)
+
+
+def get_min_driver_for_cuda(cuda_version: str) -> Optional[tuple[str, str]]:
+    """Get minimum driver versions for a CUDA version.
+
+    Args:
+        cuda_version: CUDA version string (e.g., "12.6" or "12.6.1")
+
+    Returns:
+        Tuple of (min_linux, min_windows) or None if unknown version.
+    """
+    if not cuda_version:
+        return None
+    # Match major.minor (e.g., "12.6" from "12.6.1")
+    parts = cuda_version.split(".")
+    if len(parts) >= 2:
+        major_minor = f"{parts[0]}.{parts[1]}"
+        return CUDA_DRIVER_REQUIREMENTS.get(major_minor)
+    return None
+
+
+def check_driver_compatibility(
+    driver_version: str, cuda_version: str
+) -> tuple[bool, Optional[str]]:
+    """Check if driver version is compatible with CUDA version.
+
+    Args:
+        driver_version: Installed driver version string
+        cuda_version: CUDA version from PyTorch
+
+    Returns:
+        Tuple of (is_compatible, min_required_version).
+        If CUDA version is unknown, returns (True, None).
+    """
+    min_versions = get_min_driver_for_cuda(cuda_version)
+    if not min_versions:
+        return True, None  # Unknown CUDA version, skip check
+
+    if sys.platform == "win32":
+        min_version = min_versions[1]
+    else:
+        min_version = min_versions[0]
+
+    current = parse_driver_version(driver_version)
+    required = parse_driver_version(min_version)
+
+    # Pad tuples to same length for comparison
+    max_len = max(len(current), len(required))
+    current = current + (0,) * (max_len - len(current))
+    required = required + (0,) * (max_len - len(required))
+
+    return current >= required, min_version
+
+
+def get_nvidia_driver_version() -> Optional[str]:
+    """Get NVIDIA driver version from nvidia-smi."""
+    if not shutil.which("nvidia-smi"):
+        return None
+    try:
+        result = subprocess.run(
+            ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"],
+            capture_output=True,
+            text=True,
+            timeout=5,
+        )
+        if result.returncode == 0:
+            return result.stdout.strip().split("\n")[0]
+    except Exception:
+        pass
+    return None
+
+
+def _get_package_location(name: str, dist) -> str:
+    """Get the actual installed location of a package.
+
+    For editable installs, returns the source directory.
+    For regular installs, returns the site-packages path.
+    """
+    # Try to import the package and get its __file__
+    try:
+        # Convert package name to module name (e.g., "sleap-nn" -> "sleap_nn")
+        module_name = name.replace("-", "_")
+        module = __import__(module_name)
+        if hasattr(module, "__file__") and module.__file__:
+            # Return parent of the package's __init__.py
+            from pathlib import Path
+
+            return str(Path(module.__file__).parent.parent)
+    except (ImportError, AttributeError):
+        pass
+
+    # Fallback to dist._path
+    if dist._path:
+        path = dist._path.parent
+        # Make absolute if relative
+        if not path.is_absolute():
+            from pathlib import Path
+
+            path = Path.cwd() / path
+        return str(path)
+
+    return ""
+
+
+def get_package_info(name: str) -> dict:
+    """Get package version, location, and install source.
+
+    Args:
+        name: Package name (e.g., "sleap-nn")
+
+    Returns:
+        Dict with version, location, source, and editable fields.
+    """
+    try:
+        dist = importlib.metadata.distribution(name)
+        version = dist.version
+
+        # Check for editable install and source via direct_url.json
+        is_editable = False
+        source = "pip"  # Default assumption
+        try:
+            direct_url_text = dist.read_text("direct_url.json")
+            if direct_url_text:
+                direct_url = json.loads(direct_url_text)
+                is_editable = direct_url.get("dir_info", {}).get("editable", False)
+                if is_editable:
+                    source = "editable"
+                elif "vcs_info" in direct_url:
+                    source = "git"
+                elif direct_url.get("url", "").startswith("file://"):
+                    source = "local"
+        except FileNotFoundError:
+            pass
+
+        # Fallback: detect old-style editable installs (.egg-info not in site-packages)
+        if not is_editable and dist._path:
+            path_str = str(dist._path)
+            # Old-style editable: .egg-info in source dir, not site-packages
+            if ".egg-info" in path_str and "site-packages" not in path_str:
+                is_editable = True
+                source = "editable"
+
+        # Check for conda install via INSTALLER file (only if not already known)
+        if source == "pip":
+            try:
+                installer = dist.read_text("INSTALLER")
+                if installer and installer.strip() == "conda":
+                    source = "conda"
+            except FileNotFoundError:
+                pass
+
+        # Get location (after determining if editable, so we can use the right method)
+        location = _get_package_location(name, dist)
+
+        return {
+            "version": version,
+            "location": location,
+            "source": source,
+            "editable": is_editable,
+        }
+    except importlib.metadata.PackageNotFoundError:
+        return {
+            "version": "not installed",
+            "location": "",
+            "source": "",
+            "editable": False,
+        }
+
+
+def get_system_info_dict() -> dict:
+    """Get system information as a dictionary.
+
+    Returns:
+        Dictionary with system info including Python version, platform,
+        PyTorch version, CUDA availability, GPU details, and package versions.
+    """
+    import torch
+
+    info = {
+        "python_version": sys.version.split()[0],
+        "platform": platform.platform(),
+        "pytorch_version": torch.__version__,
+        "cuda_available": torch.cuda.is_available(),
+        "cuda_version": None,
+        "cudnn_version": None,
+        "driver_version": None,
+        "driver_compatible": None,
+        "driver_min_required": None,
+        "gpu_count": 0,
+        "gpus": [],
+        "mps_available": False,
+        "accelerator": "cpu",  # cpu, cuda, or mps
+        "packages": {},
+    }
+
+    # Driver version (check even if CUDA unavailable - old driver can cause this)
+    driver = get_nvidia_driver_version()
+    if driver:
+        info["driver_version"] = driver
+
+    # CUDA details
+    if torch.cuda.is_available():
+        info["cuda_version"] = torch.version.cuda
+        info["cudnn_version"] = str(torch.backends.cudnn.version())
+        info["gpu_count"] = torch.cuda.device_count()
+        info["accelerator"] = "cuda"
+
+        # Check driver compatibility
+        if driver and info["cuda_version"]:
+            is_compatible, min_required = check_driver_compatibility(
+                driver, info["cuda_version"]
+            )
+            info["driver_compatible"] = is_compatible
+            info["driver_min_required"] = min_required
+
+        # GPU details
+        for i in range(torch.cuda.device_count()):
+            props = torch.cuda.get_device_properties(i)
+            info["gpus"].append(
+                {
+                    "id": i,
+                    "name": props.name,
+                    "compute_capability": f"{props.major}.{props.minor}",
+                    "memory_gb": round(props.total_memory / (1024**3), 1),
+                }
+            )
+
+    # MPS (Apple Silicon)
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        info["mps_available"] = True
+        info["accelerator"] = "mps"
+        info["gpu_count"] = 1
+
+    # Package versions
+    for pkg in PACKAGES:
+        info["packages"][pkg] = get_package_info(pkg)
+
+    return info
+
+
+def test_gpu_operations() -> tuple[bool, Optional[str]]:
+    """Test that GPU tensor operations work.
+
+    Returns:
+        Tuple of (success, error_message).
+    """
+    import torch
+
+    if torch.cuda.is_available():
+        try:
+            x = torch.randn(100, 100, device="cuda")
+            _ = torch.mm(x, x)
+            return True, None
+        except Exception as e:
+            return False, str(e)
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        try:
+            x = torch.randn(100, 100, device="mps")
+            _ = torch.mm(x, x)
+            return True, None
+        except Exception as e:
+            return False, str(e)
+    return False, "No GPU available"
+
+
+def _shorten_path(path: str, max_len: int = 40) -> str:
+    """Shorten a path for display, keeping the end."""
+    if len(path) <= max_len:
+        return path
+    return "..." + path[-(max_len - 3) :]
+
+
+def print_system_info() -> None:
+    """Print comprehensive system diagnostics to console."""
+    from rich.console import Console
+    from rich.table import Table
+
+    console = Console()
+    info = get_system_info_dict()
+
+    # System info table (with GPU details integrated)
+    table = Table(title="System Information", show_header=False)
+    table.add_column("Property", style="cyan")
+    table.add_column("Value", style="white")
+
+    table.add_row("Python", info["python_version"])
+    table.add_row("Platform", info["platform"])
+    table.add_row("PyTorch", info["pytorch_version"])
+
+    # GPU/Accelerator info
+    if info["accelerator"] == "cuda":
+        table.add_row("Accelerator", "CUDA")
+        table.add_row("CUDA version", info["cuda_version"] or "N/A")
+        table.add_row("cuDNN version", info["cudnn_version"] or "N/A")
+        table.add_row("Driver version", info["driver_version"] or "N/A")
+        table.add_row("GPU count", str(info["gpu_count"]))
+        # GPU details inline
+        for gpu in info["gpus"]:
+            gpu_str = f"{gpu['name']} ({gpu['memory_gb']} GB, compute {gpu['compute_capability']})"
+            table.add_row(f"GPU {gpu['id']}", gpu_str)
+    elif info["accelerator"] == "mps":
+        table.add_row("Accelerator", "MPS (Apple Silicon)")
+    else:
+        table.add_row("Accelerator", "CPU only")
+        # Show driver if present but CUDA unavailable (helps diagnose issues)
+        if info["driver_version"]:
+            table.add_row("Driver version", info["driver_version"])
+
+    console.print(table)
+
+    # Package versions table
+    console.print()
+    pkg_table = Table(title="Package Versions")
+    pkg_table.add_column("Package", style="cyan")
+    pkg_table.add_column("Version", style="white")
+    pkg_table.add_column("Source", style="yellow")
+    pkg_table.add_column("Location", style="dim")
+
+    for pkg, pkg_info in info["packages"].items():
+        if pkg_info["version"] == "not installed":
+            version_display = f"[dim]{pkg_info['version']}[/dim]"
+            pkg_table.add_row(pkg, version_display, "", "")
+        else:
+            location_display = _shorten_path(pkg_info["location"])
+            pkg_table.add_row(
+                pkg, pkg_info["version"], pkg_info["source"], location_display
+            )
+
+    console.print(pkg_table)
+
+    # Actionable diagnostics
+    console.print()
+
+    # Driver compatibility check (CUDA only)
+    if info["accelerator"] == "cuda" and info["driver_version"]:
+        if info["driver_min_required"]:
+            if info["driver_compatible"]:
+                console.print(
+                    f"[green]OK[/green] Driver is compatible: "
+                    f"{info['driver_version']} >= {info['driver_min_required']} "
+                    f"(required for CUDA {info['cuda_version']})"
+                )
+            else:
+                console.print(
+                    f"[red]FAIL[/red] Driver is too old: "
+                    f"{info['driver_version']} < {info['driver_min_required']} "
+                    f"(required for CUDA {info['cuda_version']})"
+                )
+                console.print(
+                    "  [yellow]Update your driver: https://www.nvidia.com/drivers[/yellow]"
+                )
+        else:
+            # Unknown CUDA version, can't check compatibility
+            console.print(
+                f"[yellow]![/yellow] Driver version: {info['driver_version']} "
+                f"(CUDA {info['cuda_version']} compatibility unknown)"
+            )
+    elif info["accelerator"] == "cpu" and info["driver_version"]:
+        # Has driver but no CUDA - might be a problem
+        console.print(
+            f"[yellow]![/yellow] NVIDIA driver found ({info['driver_version']}) "
+            "but CUDA is not available"
+        )
+        console.print(
+            "  [dim]This may indicate a driver/PyTorch version mismatch[/dim]"
+        )
+
+    # GPU connection test
+    success, error = test_gpu_operations()
+    if info["accelerator"] == "cuda":
+        if success:
+            console.print("[green]OK[/green] PyTorch can use GPU")
+        else:
+            console.print(f"[red]FAIL[/red] PyTorch cannot use GPU: {error}")
+    elif info["accelerator"] == "mps":
+        if success:
+            console.print("[green]OK[/green] PyTorch can use GPU")
+        else:
+            console.print(f"[red]FAIL[/red] PyTorch cannot use GPU: {error}")
+    else:
+        console.print("[dim]--[/dim] No GPU available (using CPU)")
+
+
+def get_startup_info_string() -> str:
+    """Get a concise system info string for startup logging.
+
+    Returns:
+        Single-line string with key system info.
+    """
+    import torch
+
+    from sleap_nn import __version__
+
+    parts = [f"sleap-nn {__version__}"]
+    parts.append(f"Python {sys.version.split()[0]}")
+    parts.append(f"PyTorch {torch.__version__}")
+
+    if torch.cuda.is_available():
+        parts.append(f"CUDA {torch.version.cuda}")
+        parts.append(f"{torch.cuda.device_count()} GPU(s)")
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        parts.append("MPS")
+    else:
+        parts.append("CPU only")
+
+    return " | ".join(parts)

sleap_nn/tracking/tracker.py

@@ -898,7 +898,14 @@ def run_tracker(
             tracking_target_instance_count is None
             or tracking_target_instance_count == 0
         ):
-            message = "tracking_target_instance_count is None or 0. To connect single breaks, tracking_target_instance_count should be set to an integer."
+            if max_tracks is not None:
+                suggestion = f"Add --tracking_target_instance_count {max_tracks} to your command (using your --max_tracks value)."
+            else:
+                suggestion = "Add --tracking_target_instance_count N where N is the expected number of instances per frame."
+            message = (
+                f"--post_connect_single_breaks requires --tracking_target_instance_count to be set. "
+                f"{suggestion}"
+            )
             logger.error(message)
             raise ValueError(message)
         start_final_pass_time = time()

sleap_nn/train.py

@@ -6,6 +6,7 @@ from datetime import datetime
 from time import time
 from omegaconf import DictConfig, OmegaConf
 from typing import Any, Dict, Optional, List, Tuple, Union
+import sleap_io as sio
 from sleap_nn.config.training_job_config import TrainingJobConfig
 from sleap_nn.training.model_trainer import ModelTrainer
 from sleap_nn.predict import run_inference as predict
@@ -15,15 +16,31 @@ from sleap_nn.config.get_config import (
     get_model_config,
     get_data_config,
 )
+from sleap_nn.system_info import get_startup_info_string
 
 
-def run_training(config: DictConfig):
-    """Create ModelTrainer instance and start training."""
+def run_training(
+    config: DictConfig,
+    train_labels: Optional[List[sio.Labels]] = None,
+    val_labels: Optional[List[sio.Labels]] = None,
+):
+    """Create ModelTrainer instance and start training.
+
+    Args:
+        config: Training configuration as a DictConfig.
+        train_labels: List of Labels objects for training.
+        val_labels: List of Labels objects for validation.
+            If not provided, the labels will be loaded from paths in the config.
+    """
     start_train_time = time()
     start_timestamp = str(datetime.now())
     logger.info(f"Started training at: {start_timestamp}")
+    logger.info(get_startup_info_string())
 
-    trainer = ModelTrainer.get_model_trainer_from_config(config)
+    # provide the labels as the train labels, val labels will be split from the train labels
+    trainer = ModelTrainer.get_model_trainer_from_config(
+        config, train_labels=train_labels, val_labels=val_labels
+    )
     trainer.train()
 
     finish_timestamp = str(datetime.now())
@@ -39,48 +56,44 @@ def run_training(config: DictConfig):
         # run inference on val dataset
         if trainer.config.trainer_config.save_ckpt:
             data_paths = {}
-            for index, path in enumerate(trainer.config.data_config.train_labels_path):
-                logger.info(
-                    f"Training labels path for index {index}: {(Path(trainer.config.trainer_config.ckpt_dir) / trainer.config.trainer_config.run_name).as_posix()}"
-                )
-                data_paths[f"train_{index}"] = (
-                    Path(trainer.config.trainer_config.ckpt_dir)
-                    / trainer.config.trainer_config.run_name
-                    / f"labels_train_gt_{index}.slp"
+            run_path = (
+                Path(trainer.config.trainer_config.ckpt_dir)
+                / trainer.config.trainer_config.run_name
+            )
+            for index, _ in enumerate(trainer.train_labels):
+                logger.info(f"Run path for index {index}: {run_path.as_posix()}")
+                data_paths[f"train.{index}"] = (
+                    run_path / f"labels_gt.train.{index}.slp"
                 ).as_posix()
-                data_paths[f"val_{index}"] = (
-                    Path(trainer.config.trainer_config.ckpt_dir)
-                    / trainer.config.trainer_config.run_name
-                    / f"labels_val_gt_{index}.slp"
+                data_paths[f"val.{index}"] = (
+                    run_path / f"labels_gt.val.{index}.slp"
                 ).as_posix()
 
-            if (
-                OmegaConf.select(config, "data_config.test_file_path", default=None)
-                is not None
-            ):
-                data_paths["test"] = config.data_config.test_file_path
+            # Handle test_file_path as either a string or list of strings
+            test_file_path = OmegaConf.select(
+                config, "data_config.test_file_path", default=None
+            )
+            if test_file_path is not None:
+                # Normalize to list of strings
+                if isinstance(test_file_path, str):
+                    test_paths = [test_file_path]
+                else:
+                    test_paths = list(test_file_path)
+                # Add each test path to data_paths (always use index for consistency)
+                for idx, test_path in enumerate(test_paths):
+                    data_paths[f"test.{idx}"] = test_path
 
             for d_name, path in data_paths.items():
-                pred_path = (
-                    Path(trainer.config.trainer_config.ckpt_dir)
-                    / trainer.config.trainer_config.run_name
-                    / f"pred_{d_name}.slp"
-                )
-                metrics_path = (
-                    Path(trainer.config.trainer_config.ckpt_dir)
-                    / trainer.config.trainer_config.run_name
-                    / f"{d_name}_pred_metrics.npz"
-                )
+                # d_name is now in format: "train.0", "val.0", "test.0", etc.
+                pred_path = run_path / f"labels_pr.{d_name}.slp"
+                metrics_path = run_path / f"metrics.{d_name}.npz"
 
                 pred_labels = predict(
                     data_path=path,
-                    model_paths=[
-                        Path(trainer.config.trainer_config.ckpt_dir)
-                        / trainer.config.trainer_config.run_name
-                    ],
+                    model_paths=[run_path],
                     peak_threshold=0.2,
                     make_labels=True,
-                    device=trainer.trainer.strategy.root_device,
+                    device=str(trainer.trainer.strategy.root_device),
                     output_path=pred_path,
                     ensure_rgb=config.data_config.preprocessing.ensure_rgb,
                     ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
@@ -110,7 +123,8 @@ def train(
     train_labels_path: Optional[List[str]] = None,
     val_labels_path: Optional[List[str]] = None,
     validation_fraction: float = 0.1,
-    test_file_path: Optional[str] = None,
+    use_same_data_for_val: bool = False,
+    test_file_path: Optional[Union[str, List[str]]] = None,
     provider: str = "LabelsReader",
     user_instances_only: bool = True,
     data_pipeline_fw: str = "torch_dataset",
@@ -124,6 +138,7 @@ def train(
     max_width: Optional[int] = None,
     crop_size: Optional[int] = None,
     min_crop_size: Optional[int] = 100,
+    crop_padding: Optional[int] = None,
     use_augmentations_train: bool = False,
     intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
     geometry_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
@@ -188,7 +203,11 @@ def train(
             training set to sample for generating the validation set. The remaining
             labeled frames will be left in the training set. If the `validation_labels`
             are already specified, this has no effect. Default: 0.1.
-        test_file_path: Path to test dataset (`.slp` file or `.mp4` file).
+        use_same_data_for_val: If `True`, use the same data for both training and
+            validation (train = val). Useful for intentional overfitting on small
+            datasets. When enabled, `val_labels_path` and `validation_fraction` are
+            ignored. Default: False.
+        test_file_path: Path or list of paths to test dataset(s) (`.slp` file(s) or `.mp4` file(s)).
             Note: This is used to get evaluation on test set after training is completed.
         provider: Provider class to read the input sleap files. Only "LabelsReader"
             supported for the training pipeline. Default: "LabelsReader".
@@ -210,14 +229,17 @@ def train(
             is set to True, then we convert the image to grayscale (single-channel)
             image. If the source image has only one channel and this is set to False, then we retain the single channel input. Default: `False`.
         scale: Factor to resize the image dimensions by, specified as a float. Default: 1.0.
-        max_height: Maximum height the image should be padded to. If not provided, the
+        max_height: Maximum height the original image should be resized and padded to. If not provided, the
             original image size will be retained. Default: None.
-        max_width: Maximum width the image should be padded to. If not provided, the
+        max_width: Maximum width the original image should be resized and padded to. If not provided, the
             original image size will be retained. Default: None.
         crop_size: Crop size of each instance for centered-instance model.
             If `None`, this would be automatically computed based on the largest instance
-            in the `sio.Labels` file. Default: None.
+            in the `sio.Labels` file. If `scale` is provided, then the cropped image will be resized according to `scale`. Default: None.
         min_crop_size: Minimum crop size to be used if `crop_size` is `None`. Default: 100.
+        crop_padding: Padding in pixels to add around instance bounding box when computing
+            crop size. If `None`, padding is auto-computed based on augmentation settings.
+            Only used when `crop_size` is `None`. Default: None.
         use_augmentations_train: True if the data augmentation should be applied to the
             training data, else False. Default: False.
         intensity_aug: One of ["uniform_noise", "gaussian_noise", "contrast", "brightness"]
@@ -376,6 +398,7 @@ def train(
         train_labels_path=train_labels_path,
         val_labels_path=val_labels_path,
         validation_fraction=validation_fraction,
+        use_same_data_for_val=use_same_data_for_val,
         test_file_path=test_file_path,
         provider=provider,
         user_instances_only=user_instances_only,
@@ -390,6 +413,7 @@ def train(
         max_width=max_width,
         crop_size=crop_size,
         min_crop_size=min_crop_size,
+        crop_padding=crop_padding,
         use_augmentations_train=use_augmentations_train,
         intensity_aug=intensity_aug,
         geometry_aug=geometry_aug,