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

sleap_nn/inference/provenance.py

@@ -0,0 +1,292 @@
+"""Provenance metadata utilities for inference outputs.
+
+This module provides utilities for building and managing provenance metadata
+that is stored in SLP files produced during inference. Provenance metadata
+helps track where predictions came from and how they were generated.
+"""
+
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Optional, Union
+
+import sleap_io as sio
+
+import sleap_nn
+from sleap_nn.system_info import get_system_info_dict
+
+
+def build_inference_provenance(
+    model_paths: Optional[list[str]] = None,
+    model_type: Optional[str] = None,
+    start_time: Optional[datetime] = None,
+    end_time: Optional[datetime] = None,
+    input_labels: Optional[sio.Labels] = None,
+    input_path: Optional[Union[str, Path]] = None,
+    frames_processed: Optional[int] = None,
+    frames_total: Optional[int] = None,
+    frame_selection_method: Optional[str] = None,
+    inference_params: Optional[dict[str, Any]] = None,
+    tracking_params: Optional[dict[str, Any]] = None,
+    device: Optional[str] = None,
+    cli_args: Optional[dict[str, Any]] = None,
+    include_system_info: bool = True,
+) -> dict[str, Any]:
+    """Build provenance metadata dictionary for inference output.
+
+    This function creates a comprehensive provenance dictionary that captures
+    all relevant metadata about an inference run, enabling reproducibility
+    and tracking of prediction origins.
+
+    Args:
+        model_paths: List of paths to model checkpoints used for inference.
+        model_type: Type of model used (e.g., "top_down", "bottom_up",
+            "single_instance").
+        start_time: Datetime when inference started.
+        end_time: Datetime when inference finished.
+        input_labels: Input Labels object if inference was run on an SLP file.
+            The provenance from this object will be preserved.
+        input_path: Path to input file (SLP or video).
+        frames_processed: Number of frames that were processed.
+        frames_total: Total number of frames in the input.
+        frame_selection_method: Method used to select frames (e.g., "all",
+            "labeled", "suggested", "range").
+        inference_params: Dictionary of inference parameters (peak_threshold,
+            integral_refinement, batch_size, etc.).
+        tracking_params: Dictionary of tracking parameters if tracking was run.
+        device: Device used for inference (e.g., "cuda:0", "cpu", "mps").
+        cli_args: Command-line arguments if available.
+        include_system_info: If True, include detailed system information.
+            Set to False for lighter-weight provenance.
+
+    Returns:
+        Dictionary containing provenance metadata suitable for storing in
+        Labels.provenance.
+
+    Example:
+        >>> from datetime import datetime
+        >>> provenance = build_inference_provenance(
+        ...     model_paths=["/path/to/model.ckpt"],
+        ...     model_type="top_down",
+        ...     start_time=datetime.now(),
+        ...     end_time=datetime.now(),
+        ...     device="cuda:0",
+        ... )
+        >>> labels.provenance = provenance
+        >>> labels.save("predictions.slp")
+    """
+    provenance: dict[str, Any] = {}
+
+    # Timestamps
+    if start_time is not None:
+        provenance["inference_start_timestamp"] = start_time.isoformat()
+    if end_time is not None:
+        provenance["inference_end_timestamp"] = end_time.isoformat()
+    if start_time is not None and end_time is not None:
+        runtime_seconds = (end_time - start_time).total_seconds()
+        provenance["inference_runtime_seconds"] = runtime_seconds
+
+    # Version information
+    provenance["sleap_nn_version"] = sleap_nn.__version__
+    provenance["sleap_io_version"] = sio.__version__
+
+    # Model information
+    if model_paths is not None:
+        # Store as absolute POSIX paths for cross-platform compatibility
+        provenance["model_paths"] = [
+            Path(p).resolve().as_posix() if isinstance(p, (str, Path)) else str(p)
+            for p in model_paths
+        ]
+    if model_type is not None:
+        provenance["model_type"] = model_type
+
+    # Input data lineage
+    if input_path is not None:
+        provenance["source_file"] = (
+            Path(input_path).resolve().as_posix()
+            if isinstance(input_path, (str, Path))
+            else str(input_path)
+        )
+
+    # Preserve input provenance if available
+    if input_labels is not None and hasattr(input_labels, "provenance"):
+        input_prov = dict(input_labels.provenance)
+        if input_prov:
+            provenance["input_provenance"] = input_prov
+            # Also set source_labels for compatibility with sleap-io conventions
+            if "filename" in input_prov:
+                provenance["source_labels"] = input_prov["filename"]
+
+    # Frame selection information
+    if frames_processed is not None or frames_total is not None:
+        frame_info: dict[str, Any] = {}
+        if frame_selection_method is not None:
+            frame_info["method"] = frame_selection_method
+        if frames_processed is not None:
+            frame_info["frames_processed"] = frames_processed
+        if frames_total is not None:
+            frame_info["frames_total"] = frames_total
+        if frame_info:
+            provenance["frame_selection"] = frame_info
+
+    # Inference parameters
+    if inference_params is not None:
+        # Filter out None values and convert paths
+        clean_params = {}
+        for key, value in inference_params.items():
+            if value is not None:
+                if isinstance(value, Path):
+                    clean_params[key] = value.as_posix()
+                else:
+                    clean_params[key] = value
+        if clean_params:
+            provenance["inference_config"] = clean_params
+
+    # Tracking parameters
+    if tracking_params is not None:
+        clean_tracking = {k: v for k, v in tracking_params.items() if v is not None}
+        if clean_tracking:
+            provenance["tracking_config"] = clean_tracking
+
+    # Device information
+    if device is not None:
+        provenance["device"] = device
+
+    # CLI arguments
+    if cli_args is not None:
+        # Filter out None values
+        clean_cli = {k: v for k, v in cli_args.items() if v is not None}
+        if clean_cli:
+            provenance["cli_args"] = clean_cli
+
+    # System information (can be disabled for lighter provenance)
+    if include_system_info:
+        try:
+            system_info = get_system_info_dict()
+            # Extract key fields for provenance (avoid excessive nesting)
+            provenance["system_info"] = {
+                "python_version": system_info.get("python_version"),
+                "platform": system_info.get("platform"),
+                "pytorch_version": system_info.get("pytorch_version"),
+                "cuda_version": system_info.get("cuda_version"),
+                "accelerator": system_info.get("accelerator"),
+                "gpu_count": system_info.get("gpu_count"),
+            }
+            # Include GPU names if available
+            if system_info.get("gpus"):
+                provenance["system_info"]["gpus"] = [
+                    gpu.get("name") for gpu in system_info["gpus"]
+                ]
+        except Exception:
+            # Don't fail inference if system info collection fails
+            pass
+
+    return provenance
+
+
+def build_tracking_only_provenance(
+    input_labels: Optional[sio.Labels] = None,
+    input_path: Optional[Union[str, Path]] = None,
+    start_time: Optional[datetime] = None,
+    end_time: Optional[datetime] = None,
+    tracking_params: Optional[dict[str, Any]] = None,
+    frames_processed: Optional[int] = None,
+    include_system_info: bool = True,
+) -> dict[str, Any]:
+    """Build provenance metadata for tracking-only pipeline.
+
+    This is a simplified version of build_inference_provenance for when
+    only tracking is run without model inference.
+
+    Args:
+        input_labels: Input Labels object with existing predictions.
+        input_path: Path to input SLP file.
+        start_time: Datetime when tracking started.
+        end_time: Datetime when tracking finished.
+        tracking_params: Dictionary of tracking parameters.
+        frames_processed: Number of frames that were tracked.
+        include_system_info: If True, include system information.
+
+    Returns:
+        Dictionary containing provenance metadata.
+    """
+    provenance: dict[str, Any] = {}
+
+    # Timestamps
+    if start_time is not None:
+        provenance["tracking_start_timestamp"] = start_time.isoformat()
+    if end_time is not None:
+        provenance["tracking_end_timestamp"] = end_time.isoformat()
+    if start_time is not None and end_time is not None:
+        runtime_seconds = (end_time - start_time).total_seconds()
+        provenance["tracking_runtime_seconds"] = runtime_seconds
+
+    # Version information
+    provenance["sleap_nn_version"] = sleap_nn.__version__
+    provenance["sleap_io_version"] = sio.__version__
+
+    # Note that this is tracking-only
+    provenance["pipeline_type"] = "tracking_only"
+
+    # Input data lineage
+    if input_path is not None:
+        provenance["source_file"] = (
+            Path(input_path).resolve().as_posix()
+            if isinstance(input_path, (str, Path))
+            else str(input_path)
+        )
+
+    # Preserve input provenance if available
+    if input_labels is not None and hasattr(input_labels, "provenance"):
+        input_prov = dict(input_labels.provenance)
+        if input_prov:
+            provenance["input_provenance"] = input_prov
+            if "filename" in input_prov:
+                provenance["source_labels"] = input_prov["filename"]
+
+    # Frame information
+    if frames_processed is not None:
+        provenance["frames_processed"] = frames_processed
+
+    # Tracking parameters
+    if tracking_params is not None:
+        clean_tracking = {k: v for k, v in tracking_params.items() if v is not None}
+        if clean_tracking:
+            provenance["tracking_config"] = clean_tracking
+
+    # System information
+    if include_system_info:
+        try:
+            system_info = get_system_info_dict()
+            provenance["system_info"] = {
+                "python_version": system_info.get("python_version"),
+                "platform": system_info.get("platform"),
+                "pytorch_version": system_info.get("pytorch_version"),
+                "accelerator": system_info.get("accelerator"),
+            }
+        except Exception:
+            pass
+
+    return provenance
+
+
+def merge_provenance(
+    base_provenance: dict[str, Any],
+    additional: dict[str, Any],
+    overwrite: bool = True,
+) -> dict[str, Any]:
+    """Merge additional provenance fields into base provenance.
+
+    Args:
+        base_provenance: Base provenance dictionary.
+        additional: Additional fields to merge.
+        overwrite: If True, additional fields overwrite base fields.
+            If False, base fields take precedence.
+
+    Returns:
+        Merged provenance dictionary.
+    """
+    result = dict(base_provenance)
+    for key, value in additional.items():
+        if key not in result or overwrite:
+            result[key] = value
+    return result

sleap_nn/inference/topdown.py

@@ -47,9 +47,6 @@ class CentroidCrop(L.LightningModule):
         crop_hw: Tuple (height, width) representing the crop size.
         input_scale: Float indicating if the images should be resized before being
             passed to the model.
-        precrop_resize: Float indicating the factor by which the original images
-            (not images resized for centroid model) should be resized before cropping.
-            Note: This resize happens only after getting the predictions for centroid model.
         max_stride: Maximum stride in a model that the images must be divisible by.
             If > 1, this will pad the bottom and right of the images to ensure they meet
             this divisibility criteria. Padding is applied after the scaling specified
@@ -74,7 +71,6 @@ class CentroidCrop(L.LightningModule):
         return_crops: bool = False,
         crop_hw: Optional[List[int]] = None,
         input_scale: float = 1.0,
-        precrop_resize: float = 1.0,
         max_stride: int = 1,
         use_gt_centroids: bool = False,
         anchor_ind: Optional[int] = None,
@@ -92,22 +88,25 @@ class CentroidCrop(L.LightningModule):
         self.return_crops = return_crops
         self.crop_hw = crop_hw
         self.input_scale = input_scale
-        self.precrop_resize = precrop_resize
         self.max_stride = max_stride
         self.use_gt_centroids = use_gt_centroids
         self.anchor_ind = anchor_ind
 
-    def _generate_crops(self, inputs):
+    def _generate_crops(self, inputs, cms: Optional[torch.Tensor] = None):
         """Generate Crops from the predicted centroids."""
         crops_dict = []
-        for centroid, centroid_val, image, fidx, vidx, sz, eff_sc in zip(
-            self.refined_peaks_batched,
-            self.peak_vals_batched,
-            inputs["image"],
-            inputs["frame_idx"],
-            inputs["video_idx"],
-            inputs["orig_size"],
-            inputs["eff_scale"],
+        if cms is not None:
+            cms = cms.detach()
+        for idx, (centroid, centroid_val, image, fidx, vidx, sz, eff_sc) in enumerate(
+            zip(
+                self.refined_peaks_batched,
+                self.peak_vals_batched,
+                inputs["image"],
+                inputs["frame_idx"],
+                inputs["video_idx"],
+                inputs["orig_size"],
+                inputs["eff_scale"],
+            )
         ):
             if torch.any(torch.isnan(centroid)):
                 if torch.all(torch.isnan(centroid)):
@@ -149,6 +148,11 @@ class CentroidCrop(L.LightningModule):
             ex["instance_image"] = instance_image.unsqueeze(dim=1)
             ex["orig_size"] = torch.cat([torch.Tensor(sz)] * n)
             ex["eff_scale"] = torch.Tensor([eff_sc] * n)
+            ex["pred_centroids"] = centroid
+            if self.return_confmaps:
+                ex["pred_centroid_confmaps"] = torch.cat(
+                    [cms[idx].unsqueeze(dim=0)] * n
+                )
             crops_dict.append(ex)
 
         return crops_dict
@@ -204,12 +208,6 @@ class CentroidCrop(L.LightningModule):
 
             if self.return_crops:
                 crops_dict = self._generate_crops(inputs)
-                inputs["image"] = resize_image(inputs["image"], self.precrop_resize)
-                inputs["centroids"] *= self.precrop_resize
-                scaled_refined_peaks = []
-                for ref_peak in self.refined_peaks_batched:
-                    scaled_refined_peaks.append(ref_peak * self.precrop_resize)
-                self.refined_peaks_batched = scaled_refined_peaks
                 return crops_dict
             else:
                 return inputs
@@ -274,19 +272,13 @@ class CentroidCrop(L.LightningModule):
 
             # Generate crops if return_crops=True to pass the crops to CenteredInstance model.
             if self.return_crops:
-                inputs["image"] = resize_image(inputs["image"], self.precrop_resize)
-                scaled_refined_peaks = []
-                for ref_peak in self.refined_peaks_batched:
-                    scaled_refined_peaks.append(ref_peak * self.precrop_resize)
-                self.refined_peaks_batched = scaled_refined_peaks
-
                 inputs.update(
                     {
                         "centroids": self.refined_peaks_batched,
                         "centroid_vals": self.peak_vals_batched,
                     }
                 )
-                crops_dict = self._generate_crops(inputs)
+                crops_dict = self._generate_crops(inputs, cms)
                 return crops_dict
             else:
                 # batch the peaks to pass it to FindInstancePeaksGroundTruth class.
@@ -359,7 +351,11 @@ class FindInstancePeaksGroundTruth(L.LightningModule):
 
     def forward(self, batch: Dict[str, torch.Tensor]) -> Dict[str, np.array]:
         """Return the ground truth instance peaks given a set of crops."""
-        b, _, max_inst, nodes, _ = batch["instances"].shape
+        b, _, _, nodes, _ = batch["instances"].shape
+        # Use number of centroids as max_inst to ensure consistent output shape
+        # This handles the case where max_instances limits centroids but instances
+        # tensor has a different (global) max_instances from the labels file
+        num_centroids = batch["centroids"].shape[2]
         inst = (
             batch["instances"].unsqueeze(dim=-4).float()
         )  # (batch, 1, 1, n_inst, nodes, 2)
@@ -389,26 +385,26 @@ class FindInstancePeaksGroundTruth(L.LightningModule):
         parsed = 0
         for i in range(b):
             if i not in matched_batch_inds:
-                batch_peaks = torch.full((max_inst, nodes, 2), torch.nan)
-                vals = torch.full((max_inst, nodes), torch.nan)
+                batch_peaks = torch.full((num_centroids, nodes, 2), torch.nan)
+                vals = torch.full((num_centroids, nodes), torch.nan)
             else:
                 c = counts[i]
                 batch_peaks = peaks_list[parsed : parsed + c]
                 num_inst = len(batch_peaks)
                 vals = torch.ones((num_inst, nodes))
-                if c < max_inst:
+                if c < num_centroids:
                     batch_peaks = torch.cat(
                         [
                             batch_peaks,
-                            torch.full((max_inst - num_inst, nodes, 2), torch.nan),
+                            torch.full((num_centroids - num_inst, nodes, 2), torch.nan),
                         ]
                     )
                     vals = torch.cat(
-                        [vals, torch.full((max_inst - num_inst, nodes), torch.nan)]
+                        [vals, torch.full((num_centroids - num_inst, nodes), torch.nan)]
                     )
                 else:
-                    batch_peaks = batch_peaks[:max_inst]
-                    vals = vals[:max_inst]
+                    batch_peaks = batch_peaks[:num_centroids]
+                    vals = vals[:num_centroids]
                 parsed += c
 
             batch_peaks = batch_peaks.unsqueeze(dim=0)
@@ -432,33 +428,45 @@ class FindInstancePeaksGroundTruth(L.LightningModule):
         peaks_output["pred_instance_peaks"] = peaks
         peaks_output["pred_peak_values"] = peaks_vals
 
-        batch_size, num_centroids = (
-            batch["centroids"].shape[0],
-            batch["centroids"].shape[2],
-        )
+        batch_size = batch["centroids"].shape[0]
         output_dict = {}
         output_dict["centroid"] = batch["centroids"].squeeze(dim=1).reshape(-1, 1, 2)
         output_dict["centroid_val"] = batch["centroid_vals"].reshape(-1)
-        output_dict["pred_instance_peaks"] = batch["pred_instance_peaks"].reshape(
-            -1, nodes, 2
+        output_dict["pred_instance_peaks"] = peaks_output[
+            "pred_instance_peaks"
+        ].reshape(-1, nodes, 2)
+        output_dict["pred_peak_values"] = peaks_output["pred_peak_values"].reshape(
+            -1, nodes
         )
-        output_dict["pred_peak_values"] = batch["pred_peak_values"].reshape(-1, nodes)
         output_dict["instance_bbox"] = torch.zeros(
             (batch_size * num_centroids, 1, 4, 2)
         )
         frame_inds = []
         video_inds = []
         orig_szs = []
+        images = []
+        centroid_confmaps = []
         for b_idx in range(b):
             curr_batch_size = len(batch["centroids"][b_idx][0])
             frame_inds.extend([batch["frame_idx"][b_idx]] * curr_batch_size)
             video_inds.extend([batch["video_idx"][b_idx]] * curr_batch_size)
             orig_szs.append(torch.cat([batch["orig_size"][b_idx]] * curr_batch_size))
+            images.append(
+                batch["image"][b_idx].unsqueeze(0).repeat(curr_batch_size, 1, 1, 1, 1)
+            )
+            if "pred_centroid_confmaps" in batch:
+                centroid_confmaps.append(
+                    batch["pred_centroid_confmaps"][b_idx]
+                    .unsqueeze(0)
+                    .repeat(curr_batch_size, 1, 1, 1)
+                )
 
         output_dict["frame_idx"] = torch.tensor(frame_inds)
         output_dict["video_idx"] = torch.tensor(video_inds)
         output_dict["orig_size"] = torch.concatenate(orig_szs, dim=0)
-
+        output_dict["image"] = torch.cat(images, dim=0)
+        if centroid_confmaps:
+            output_dict["pred_centroid_confmaps"] = torch.cat(centroid_confmaps, dim=0)
         return output_dict
 
 
@@ -548,6 +556,8 @@ class FindInstancePeaks(L.LightningModule):
         # Network forward pass.
         # resize and pad the input image
         input_image = inputs["instance_image"]
+        # resize the crop image
+        input_image = resize_image(input_image, self.input_scale)
         if self.max_stride != 1:
             input_image = apply_pad_to_stride(input_image, self.max_stride)
 
@@ -569,8 +579,6 @@ class FindInstancePeaks(L.LightningModule):
             inputs["eff_scale"].unsqueeze(dim=1).unsqueeze(dim=2).to(peak_points.device)
         )
 
-        inputs["instance_bbox"] = inputs["instance_bbox"] / self.input_scale
-
         inputs["instance_bbox"] = inputs["instance_bbox"] / (
             inputs["eff_scale"]
             .unsqueeze(dim=1)
@@ -679,6 +687,8 @@ class TopDownMultiClassFindInstancePeaks(L.LightningModule):
         # Network forward pass.
         # resize and pad the input image
         input_image = inputs["instance_image"]
+        # resize the crop image
+        input_image = resize_image(input_image, self.input_scale)
         if self.max_stride != 1:
             input_image = apply_pad_to_stride(input_image, self.max_stride)
 
@@ -702,8 +712,6 @@ class TopDownMultiClassFindInstancePeaks(L.LightningModule):
             inputs["eff_scale"].unsqueeze(dim=1).unsqueeze(dim=2).to(peak_points.device)
         )
 
-        inputs["instance_bbox"] = inputs["instance_bbox"] / self.input_scale
-
         inputs["instance_bbox"] = inputs["instance_bbox"] / (
             inputs["eff_scale"]
             .unsqueeze(dim=1)

sleap_nn/legacy_models.py

@@ -7,9 +7,8 @@ TensorFlow/Keras backend to PyTorch format compatible with sleap-nn.
 import h5py
 import numpy as np
 import torch
-from typing import Dict, Tuple, Any, Optional, List
+from typing import Dict, Any, Optional
 from pathlib import Path
-from omegaconf import OmegaConf
 import re
 from loguru import logger
 
@@ -181,18 +180,61 @@ 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
+    legacy_weights: Dict[str, np.ndarray],
+    pytorch_model: torch.nn.Module,
+    component: Optional[str] = None,
 ) -> 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
@@ -201,7 +243,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 legacy_weights.items():
+    for legacy_path, weight in filtered_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/", "")
@@ -220,8 +262,6 @@ 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
@@ -266,12 +306,17 @@ 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(legacy_weights)}, PyTorch parameters: {len(pytorch_params)}"
+            f"Legacy weights: {len(filtered_weights)}, PyTorch parameters: {len(pytorch_params)}"
         )
     else:
         logger.info(
-            f"Successfully mapped {len(mapping)}/{len(legacy_weights)} legacy weights to PyTorch parameters"
+            f"Successfully mapped {len(mapping)}/{len(pytorch_params)} PyTorch parameters from legacy weights"
         )
+        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
 
@@ -280,6 +325,7 @@ 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.
 
@@ -288,6 +334,10 @@ 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)
@@ -295,7 +345,9 @@ def load_legacy_model_weights(
     if mapping is None:
         # Attempt automatic mapping
         try:
-            mapping = map_legacy_to_pytorch_layers(legacy_weights, pytorch_model)
+            mapping = map_legacy_to_pytorch_layers(
+                legacy_weights, pytorch_model, component=component
+            )
         except Exception as e:
             logger.error(f"Failed to create weight mappings: {e}")
             return
@@ -417,7 +469,9 @@ 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}"
+                    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)
             else:
                 # Bias : just compare all values
                 keras_mean = np.mean(original_weight)
@@ -426,7 +480,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)