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

sleap_nn/inference/predictors.py

@@ -56,8 +56,6 @@ from rich.progress import (
     MofNCompleteColumn,
 )
 from time import time
-import json
-import sys
 
 
 def _filter_user_labeled_frames(
@@ -135,8 +133,6 @@ class Predictor(ABC):
             `backbone_config`. This determines the downsampling factor applied by the backbone,
             and is used to ensure that input images are padded or resized to be compatible
             with the model's architecture. Default: 16.
-        gui: If True, outputs JSON progress lines for GUI integration instead of
-            Rich progress bars. Default: False.
     """
 
     preprocess: bool = True
@@ -156,7 +152,6 @@ class Predictor(ABC):
     ] = None
     instances_key: bool = False
     max_stride: int = 16
-    gui: bool = False
 
     @classmethod
     def from_model_paths(
@@ -360,7 +355,7 @@ class Predictor(ABC):
     def make_pipeline(
         self,
         data_path: str,
-        queue_maxsize: int = 32,
+        queue_maxsize: int = 8,
         frames: Optional[list] = None,
         only_labeled_frames: bool = False,
         only_suggested_frames: bool = False,
@@ -386,102 +381,6 @@ class Predictor(ABC):
                     v[n] = v[n].cpu().numpy()
         return output
 
-    def _process_batch(self) -> tuple:
-        """Process a single batch of frames from the pipeline.
-
-        Returns:
-            Tuple of (imgs, fidxs, vidxs, org_szs, instances, eff_scales, done)
-            where done is True if the pipeline has finished.
-        """
-        imgs = []
-        fidxs = []
-        vidxs = []
-        org_szs = []
-        instances = []
-        eff_scales = []
-        done = False
-
-        for _ in range(self.batch_size):
-            frame = self.pipeline.frame_buffer.get()
-            if frame["image"] is None:
-                done = True
-                break
-            frame["image"], eff_scale = apply_sizematcher(
-                frame["image"],
-                self.preprocess_config["max_height"],
-                self.preprocess_config["max_width"],
-            )
-            if self.instances_key:
-                frame["instances"] = frame["instances"] * eff_scale
-            if self.preprocess_config["ensure_rgb"] and frame["image"].shape[-3] != 3:
-                frame["image"] = frame["image"].repeat(1, 3, 1, 1)
-            elif (
-                self.preprocess_config["ensure_grayscale"]
-                and frame["image"].shape[-3] != 1
-            ):
-                frame["image"] = F.rgb_to_grayscale(
-                    frame["image"], num_output_channels=1
-                )
-
-            eff_scales.append(torch.tensor(eff_scale))
-            imgs.append(frame["image"].unsqueeze(dim=0))
-            fidxs.append(frame["frame_idx"])
-            vidxs.append(frame["video_idx"])
-            org_szs.append(frame["orig_size"].unsqueeze(dim=0))
-            if self.instances_key:
-                instances.append(frame["instances"].unsqueeze(dim=0))
-
-        return imgs, fidxs, vidxs, org_szs, instances, eff_scales, done
-
-    def _run_inference_on_batch(
-        self, imgs, fidxs, vidxs, org_szs, instances, eff_scales
-    ) -> Iterator[Dict[str, np.ndarray]]:
-        """Run inference on a prepared batch of frames.
-
-        Args:
-            imgs: List of image tensors.
-            fidxs: List of frame indices.
-            vidxs: List of video indices.
-            org_szs: List of original sizes.
-            instances: List of instance tensors.
-            eff_scales: List of effective scales.
-
-        Yields:
-            Dictionaries containing inference results for each frame.
-        """
-        # TODO: all preprocessing should be moved into InferenceModels to be exportable.
-        imgs = torch.concatenate(imgs, dim=0)
-        fidxs = torch.tensor(fidxs, dtype=torch.int32)
-        vidxs = torch.tensor(vidxs, dtype=torch.int32)
-        org_szs = torch.concatenate(org_szs, dim=0)
-        eff_scales = torch.tensor(eff_scales, dtype=torch.float32)
-        if self.instances_key:
-            instances = torch.concatenate(instances, dim=0)
-        ex = {
-            "image": imgs,
-            "frame_idx": fidxs,
-            "video_idx": vidxs,
-            "orig_size": org_szs,
-            "eff_scale": eff_scales,
-        }
-        if self.instances_key:
-            ex["instances"] = instances
-        if self.preprocess:
-            scale = self.preprocess_config["scale"]
-            if scale != 1.0:
-                if self.instances_key:
-                    ex["image"], ex["instances"] = apply_resizer(
-                        ex["image"], ex["instances"]
-                    )
-                else:
-                    ex["image"] = resize_image(ex["image"], scale)
-            ex["image"] = apply_pad_to_stride(ex["image"], self.max_stride)
-        outputs_list = self.inference_model(ex)
-        if outputs_list is not None:
-            for output in outputs_list:
-                output = self._convert_tensors_to_numpy(output)
-                yield output
-
     def _predict_generator(self) -> Iterator[Dict[str, np.ndarray]]:
         """Create a generator that yields batches of inference results.
 
@@ -501,14 +400,114 @@ class Predictor(ABC):
         # Loop over data batches.
         self.pipeline.start()
         total_frames = self.pipeline.total_len()
+        done = False
 
         try:
-            if self.gui:
-                # GUI mode: emit JSON progress lines
-                yield from self._predict_generator_gui(total_frames)
-            else:
-                # Normal mode: use Rich progress bar
-                yield from self._predict_generator_rich(total_frames)
+            with Progress(
+                "{task.description}",
+                BarColumn(),
+                "[progress.percentage]{task.percentage:>3.0f}%",
+                MofNCompleteColumn(),
+                "ETA:",
+                TimeRemainingColumn(),
+                "Elapsed:",
+                TimeElapsedColumn(),
+                RateColumn(),
+                auto_refresh=False,
+                refresh_per_second=4,  # Change to self.report_rate if needed
+                speed_estimate_period=5,
+            ) as progress:
+                task = progress.add_task("Predicting...", total=total_frames)
+                last_report = time()
+
+                done = False
+                while not done:
+                    imgs = []
+                    fidxs = []
+                    vidxs = []
+                    org_szs = []
+                    instances = []
+                    eff_scales = []
+                    for _ in range(self.batch_size):
+                        frame = self.pipeline.frame_buffer.get()
+                        if frame["image"] is None:
+                            done = True
+                            break
+                        frame["image"], eff_scale = apply_sizematcher(
+                            frame["image"],
+                            self.preprocess_config["max_height"],
+                            self.preprocess_config["max_width"],
+                        )
+                        if self.instances_key:
+                            frame["instances"] = frame["instances"] * eff_scale
+                        if (
+                            self.preprocess_config["ensure_rgb"]
+                            and frame["image"].shape[-3] != 3
+                        ):
+                            frame["image"] = frame["image"].repeat(1, 3, 1, 1)
+                        elif (
+                            self.preprocess_config["ensure_grayscale"]
+                            and frame["image"].shape[-3] != 1
+                        ):
+                            frame["image"] = F.rgb_to_grayscale(
+                                frame["image"], num_output_channels=1
+                            )
+
+                        eff_scales.append(torch.tensor(eff_scale))
+                        imgs.append(frame["image"].unsqueeze(dim=0))
+                        fidxs.append(frame["frame_idx"])
+                        vidxs.append(frame["video_idx"])
+                        org_szs.append(frame["orig_size"].unsqueeze(dim=0))
+                        if self.instances_key:
+                            instances.append(frame["instances"].unsqueeze(dim=0))
+                    if imgs:
+                        # TODO: all preprocessing should be moved into InferenceModels to be exportable.
+                        imgs = torch.concatenate(imgs, dim=0)
+                        fidxs = torch.tensor(fidxs, dtype=torch.int32)
+                        vidxs = torch.tensor(vidxs, dtype=torch.int32)
+                        org_szs = torch.concatenate(org_szs, dim=0)
+                        eff_scales = torch.tensor(eff_scales, dtype=torch.float32)
+                        if self.instances_key:
+                            instances = torch.concatenate(instances, dim=0)
+                        ex = {
+                            "image": imgs,
+                            "frame_idx": fidxs,
+                            "video_idx": vidxs,
+                            "orig_size": org_szs,
+                            "eff_scale": eff_scales,
+                        }
+                        if self.instances_key:
+                            ex["instances"] = instances
+                        if self.preprocess:
+                            scale = self.preprocess_config["scale"]
+                            if scale != 1.0:
+                                if self.instances_key:
+                                    ex["image"], ex["instances"] = apply_resizer(
+                                        ex["image"], ex["instances"]
+                                    )
+                                else:
+                                    ex["image"] = resize_image(ex["image"], scale)
+                            ex["image"] = apply_pad_to_stride(
+                                ex["image"], self.max_stride
+                            )
+                        outputs_list = self.inference_model(ex)
+                        if outputs_list is not None:
+                            for output in outputs_list:
+                                output = self._convert_tensors_to_numpy(output)
+                                yield output
+
+                        # Advance progress
+                        num_frames = (
+                            len(ex["frame_idx"])
+                            if "frame_idx" in ex
+                            else self.batch_size
+                        )
+                        progress.update(task, advance=num_frames)
+
+                    # Manually refresh progress bar
+                    if time() - last_report > 0.25:
+                        progress.refresh()
+                        last_report = time()
 
         except KeyboardInterrupt:
             logger.info("Inference interrupted by user")
@@ -521,112 +520,6 @@ class Predictor(ABC):
 
         self.pipeline.join()
 
-    def _predict_generator_gui(
-        self, total_frames: int
-    ) -> Iterator[Dict[str, np.ndarray]]:
-        """Generator for GUI mode with JSON progress output.
-
-        Args:
-            total_frames: Total number of frames to process.
-
-        Yields:
-            Dictionaries containing inference results for each frame.
-        """
-        start_time = time()
-        frames_processed = 0
-        last_report = time()
-        done = False
-
-        while not done:
-            imgs, fidxs, vidxs, org_szs, instances, eff_scales, done = (
-                self._process_batch()
-            )
-
-            if imgs:
-                yield from self._run_inference_on_batch(
-                    imgs, fidxs, vidxs, org_szs, instances, eff_scales
-                )
-
-                # Update progress
-                num_frames = len(fidxs)
-                frames_processed += num_frames
-
-                # Emit JSON progress (throttled to ~4Hz)
-                if time() - last_report > 0.25:
-                    elapsed = time() - start_time
-                    rate = frames_processed / elapsed if elapsed > 0 else 0
-                    remaining = total_frames - frames_processed
-                    eta = remaining / rate if rate > 0 else 0
-
-                    progress_data = {
-                        "n_processed": frames_processed,
-                        "n_total": total_frames,
-                        "rate": round(rate, 1),
-                        "eta": round(eta, 1),
-                    }
-                    print(json.dumps(progress_data), flush=True)
-                    last_report = time()
-
-        # Final progress emit to ensure 100% is shown
-        elapsed = time() - start_time
-        progress_data = {
-            "n_processed": total_frames,
-            "n_total": total_frames,
-            "rate": round(frames_processed / elapsed, 1) if elapsed > 0 else 0,
-            "eta": 0,
-        }
-        print(json.dumps(progress_data), flush=True)
-
-    def _predict_generator_rich(
-        self, total_frames: int
-    ) -> Iterator[Dict[str, np.ndarray]]:
-        """Generator for normal mode with Rich progress bar.
-
-        Args:
-            total_frames: Total number of frames to process.
-
-        Yields:
-            Dictionaries containing inference results for each frame.
-        """
-        with Progress(
-            "{task.description}",
-            BarColumn(),
-            "[progress.percentage]{task.percentage:>3.0f}%",
-            MofNCompleteColumn(),
-            "ETA:",
-            TimeRemainingColumn(),
-            "Elapsed:",
-            TimeElapsedColumn(),
-            RateColumn(),
-            auto_refresh=False,
-            refresh_per_second=4,
-            speed_estimate_period=5,
-        ) as progress:
-            task = progress.add_task("Predicting...", total=total_frames)
-            last_report = time()
-            done = False
-
-            while not done:
-                imgs, fidxs, vidxs, org_szs, instances, eff_scales, done = (
-                    self._process_batch()
-                )
-
-                if imgs:
-                    yield from self._run_inference_on_batch(
-                        imgs, fidxs, vidxs, org_szs, instances, eff_scales
-                    )
-
-                    # Advance progress
-                    num_frames = len(fidxs)
-                    progress.update(task, advance=num_frames)
-
-                # Manually refresh progress bar
-                if time() - last_report > 0.25:
-                    progress.refresh()
-                    last_report = time()
-
-        self.pipeline.join()
-
     def predict(
         self,
         make_labels: bool = True,
@@ -1214,7 +1107,7 @@ class TopDownPredictor(Predictor):
     def make_pipeline(
         self,
         inference_object: Union[str, Path, sio.Labels, sio.Video],
-        queue_maxsize: int = 32,
+        queue_maxsize: int = 8,
         frames: Optional[list] = None,
         only_labeled_frames: bool = False,
         only_suggested_frames: bool = False,
@@ -1228,7 +1121,7 @@ class TopDownPredictor(Predictor):
 
         Args:
             inference_object: (str) Path to `.slp` file or `.mp4` or sio.Labels or sio.Video to run inference on.
-            queue_maxsize: (int) Maximum size of the frame buffer queue. Default: 32.
+            queue_maxsize: (int) Maximum size of the frame buffer queue. Default: 8.
             frames: (list) List of frames indices. If `None`, all frames in the video are used. Default: None.
             only_labeled_frames: (bool) `True` if inference should be run only on user-labeled frames. Default: `False`.
             only_suggested_frames: (bool) `True` if inference should be run only on unlabeled suggested frames. Default: `False`.
@@ -1644,7 +1537,7 @@ class SingleInstancePredictor(Predictor):
     def make_pipeline(
         self,
         inference_object: Union[str, Path, sio.Labels, sio.Video],
-        queue_maxsize: int = 32,
+        queue_maxsize: int = 8,
         frames: Optional[list] = None,
         only_labeled_frames: bool = False,
         only_suggested_frames: bool = False,
@@ -1658,7 +1551,7 @@ class SingleInstancePredictor(Predictor):
 
         Args:
             inference_object: (str) Path to `.slp` file or `.mp4` or sio.Labels or sio.Video to run inference on.
-            queue_maxsize: (int) Maximum size of the frame buffer queue. Default: 32.
+            queue_maxsize: (int) Maximum size of the frame buffer queue. Default: 8.
             frames: List of frames indices. If `None`, all frames in the video are used. Default: None.
             only_labeled_frames: (bool) `True` if inference should be run only on user-labeled frames. Default: `False`.
             only_suggested_frames: (bool) `True` if inference should be run only on unlabeled suggested frames. Default: `False`.
@@ -2094,7 +1987,7 @@ class BottomUpPredictor(Predictor):
     def make_pipeline(
         self,
         inference_object: Union[str, Path, sio.Labels, sio.Video],
-        queue_maxsize: int = 32,
+        queue_maxsize: int = 8,
         frames: Optional[list] = None,
         only_labeled_frames: bool = False,
         only_suggested_frames: bool = False,
@@ -2108,7 +2001,7 @@ class BottomUpPredictor(Predictor):
 
         Args:
             inference_object: (str) Path to `.slp` file or `.mp4` or sio.Labels or sio.Video to run inference on.
-            queue_maxsize: (int) Maximum size of the frame buffer queue. Default: 32.
+            queue_maxsize: (int) Maximum size of the frame buffer queue. Default: 8.
             frames: List of frames indices. If `None`, all frames in the video are used. Default: None.
             only_labeled_frames: (bool) `True` if inference should be run only on user-labeled frames. Default: `False`.
             only_suggested_frames: (bool) `True` if inference should be run only on unlabeled suggested frames. Default: `False`.
@@ -2541,7 +2434,7 @@ class BottomUpMultiClassPredictor(Predictor):
     def make_pipeline(
         self,
         inference_object: Union[str, Path, sio.Labels, sio.Video],
-        queue_maxsize: int = 32,
+        queue_maxsize: int = 8,
         frames: Optional[list] = None,
         only_labeled_frames: bool = False,
         only_suggested_frames: bool = False,
@@ -2555,7 +2448,7 @@ class BottomUpMultiClassPredictor(Predictor):
 
         Args:
             inference_object: (str) Path to `.slp` file or `.mp4` or sio.Labels or sio.Video to run inference on.
-            queue_maxsize: (int) Maximum size of the frame buffer queue. Default: 32.
+            queue_maxsize: (int) Maximum size of the frame buffer queue. Default: 8.
             frames: List of frames indices. If `None`, all frames in the video are used. Default: None.
             only_labeled_frames: (bool) `True` if inference should be run only on user-labeled frames. Default: `False`.
             only_suggested_frames: (bool) `True` if inference should be run only on unlabeled suggested frames. Default: `False`.
@@ -3296,7 +3189,7 @@ class TopDownMultiClassPredictor(Predictor):
     def make_pipeline(
         self,
         inference_object: Union[str, Path, sio.Labels, sio.Video],
-        queue_maxsize: int = 32,
+        queue_maxsize: int = 8,
         frames: Optional[list] = None,
         only_labeled_frames: bool = False,
         only_suggested_frames: bool = False,
@@ -3310,7 +3203,7 @@ class TopDownMultiClassPredictor(Predictor):
 
         Args:
             inference_object: (str) Path to `.slp` file or `.mp4` or sio.Labels or sio.Video to run inference on.
-            queue_maxsize: (int) Maximum size of the frame buffer queue. Default: 32.
+            queue_maxsize: (int) Maximum size of the frame buffer queue. Default: 8.
             frames: (list) List of frames indices. If `None`, all frames in the video are used. Default: None.
             only_labeled_frames: (bool) `True` if inference should be run only on user-labeled frames. Default: `False`.
             only_suggested_frames: (bool) `True` if inference should be run only on unlabeled suggested frames. Default: `False`.

sleap_nn/legacy_models.py

@@ -7,8 +7,9 @@ TensorFlow/Keras backend to PyTorch format compatible with sleap-nn.
 import h5py
 import numpy as np
 import torch
-from typing import Dict, Any, Optional
+from typing import Dict, Tuple, Any, Optional, List
 from pathlib import Path
+from omegaconf import OmegaConf
 import re
 from loguru import logger
 
@@ -180,61 +181,18 @@ def parse_keras_layer_name(layer_path: str) -> Dict[str, Any]:
     return info
 
 
-def filter_legacy_weights_by_component(
-    legacy_weights: Dict[str, np.ndarray], component: Optional[str]
-) -> Dict[str, np.ndarray]:
-    """Filter legacy weights based on component type.
-
-    Args:
-        legacy_weights: Dictionary of legacy weights from load_keras_weights()
-        component: Component type to filter for. One of:
-            - "backbone": Keep only encoder/decoder weights (exclude heads)
-            - "head": Keep only head layer weights
-            - None: No filtering (keep all weights)
-
-    Returns:
-        Filtered dictionary of legacy weights
-    """
-    if component is None:
-        return legacy_weights
-
-    filtered = {}
-    for path, weight in legacy_weights.items():
-        # Check if this is a head layer (contains "Head" in the path)
-        is_head_layer = "Head" in path
-
-        if component == "backbone" and not is_head_layer:
-            filtered[path] = weight
-        elif component == "head" and is_head_layer:
-            filtered[path] = weight
-
-    return filtered
-
-
 def map_legacy_to_pytorch_layers(
-    legacy_weights: Dict[str, np.ndarray],
-    pytorch_model: torch.nn.Module,
-    component: Optional[str] = None,
+    legacy_weights: Dict[str, np.ndarray], pytorch_model: torch.nn.Module
 ) -> Dict[str, str]:
     """Create mapping between legacy Keras layers and PyTorch model layers.
 
     Args:
         legacy_weights: Dictionary of legacy weights from load_keras_weights()
         pytorch_model: PyTorch model instance to map to
-        component: Optional component type for filtering weights before mapping.
-            One of "backbone", "head", or None (no filtering).
 
     Returns:
         Dictionary mapping legacy layer paths to PyTorch parameter names
     """
-    # Filter weights based on component type
-    filtered_weights = filter_legacy_weights_by_component(legacy_weights, component)
-
-    if component is not None:
-        logger.info(
-            f"Filtered legacy weights for {component}: "
-            f"{len(filtered_weights)}/{len(legacy_weights)} weights"
-        )
     mapping = {}
 
     # Get all PyTorch parameters with their shapes
@@ -243,7 +201,7 @@ def map_legacy_to_pytorch_layers(
         pytorch_params[name] = param.shape
 
     # For each legacy weight, find the corresponding PyTorch parameter
-    for legacy_path, weight in filtered_weights.items():
+    for legacy_path, weight in legacy_weights.items():
         # Extract the layer name from the legacy path
         # Legacy path format: "model_weights/stack0_enc0_conv0/stack0_enc0_conv0/kernel:0"
         clean_path = legacy_path.replace("model_weights/", "")
@@ -262,6 +220,8 @@ def map_legacy_to_pytorch_layers(
         # This handles cases where Keras uses suffixes like _0, _1, etc.
         if "Head" in layer_name:
             # Remove trailing _N where N is a number
+            import re
+
             layer_name_clean = re.sub(r"_\d+$", "", layer_name)
         else:
             layer_name_clean = layer_name
@@ -306,17 +266,12 @@ def map_legacy_to_pytorch_layers(
     if not mapping:
         logger.info(
             f"No mappings could be created between legacy weights and PyTorch model. "
-            f"Legacy weights: {len(filtered_weights)}, PyTorch parameters: {len(pytorch_params)}"
+            f"Legacy weights: {len(legacy_weights)}, PyTorch parameters: {len(pytorch_params)}"
         )
     else:
         logger.info(
-            f"Successfully mapped {len(mapping)}/{len(pytorch_params)} PyTorch parameters from legacy weights"
+            f"Successfully mapped {len(mapping)}/{len(legacy_weights)} legacy weights to PyTorch parameters"
         )
-        unmatched_count = len(filtered_weights) - len(mapping)
-        if unmatched_count > 0:
-            logger.warning(
-                f"({unmatched_count} legacy weights did not match any parameters in this model component)"
-            )
 
     return mapping
 
@@ -325,7 +280,6 @@ def load_legacy_model_weights(
     pytorch_model: torch.nn.Module,
     h5_path: str,
     mapping: Optional[Dict[str, str]] = None,
-    component: Optional[str] = None,
 ) -> None:
     """Load legacy Keras weights into a PyTorch model.
 
@@ -334,10 +288,6 @@ def load_legacy_model_weights(
         h5_path: Path to the legacy .h5 model file
         mapping: Optional manual mapping of layer names. If None,
                  will attempt automatic mapping.
-        component: Optional component type for filtering weights. One of:
-            - "backbone": Only load encoder/decoder weights (exclude heads)
-            - "head": Only load head layer weights
-            - None: Load all weights (default, for full model loading)
     """
     # Load legacy weights
     legacy_weights = load_keras_weights(h5_path)
@@ -345,9 +295,7 @@ def load_legacy_model_weights(
     if mapping is None:
         # Attempt automatic mapping
         try:
-            mapping = map_legacy_to_pytorch_layers(
-                legacy_weights, pytorch_model, component=component
-            )
+            mapping = map_legacy_to_pytorch_layers(legacy_weights, pytorch_model)
         except Exception as e:
             logger.error(f"Failed to create weight mappings: {e}")
             return
@@ -469,9 +417,7 @@ def load_legacy_model_weights(
                 ).item()
                 diff = abs(keras_mean - torch_mean)
                 if diff > 1e-6:
-                    message = f"Weight verification failed for {pytorch_name} (linear): keras={keras_mean:.6f}, torch={torch_mean:.6f}, diff={diff:.6e}"
-                    logger.error(message)
-                    verification_errors.append(message)
+                    message = f"Weight verification failed for {pytorch_name} linear): keras={keras_mean:.6f}, torch={torch_mean:.6f}, diff={diff:.6e}"
             else:
                 # Bias : just compare all values
                 keras_mean = np.mean(original_weight)
@@ -480,7 +426,7 @@ def load_legacy_model_weights(
                 ).item()
                 diff = abs(keras_mean - torch_mean)
                 if diff > 1e-6:
-                    message = f"Weight verification failed for {pytorch_name} (bias): keras={keras_mean:.6f}, torch={torch_mean:.6f}, diff={diff:.6e}"
+                    message = f"Weight verification failed for {pytorch_name} bias): keras={keras_mean:.6f}, torch={torch_mean:.6f}, diff={diff:.6e}"
                     logger.error(message)
                     verification_errors.append(message)