sleap-nn 0.1.0a1__tar.gz → 0.1.0a2__tar.gz

{sleap_nn-0.1.0a1 → sleap_nn-0.1.0a2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sleap-nn
-Version: 0.1.0a1
+Version: 0.1.0a2
 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,7 +13,7 @@ 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.0
+Requires-Dist: sleap-io<0.7.0,>=0.6.2
 Requires-Dist: numpy
 Requires-Dist: lightning
 Requires-Dist: kornia

{sleap_nn-0.1.0a1 → sleap_nn-0.1.0a2}/pyproject.toml

@@ -29,7 +29,7 @@ classifiers = [
     "Programming Language :: Python :: 3.13",
 ]
 dependencies = [
-    "sleap-io>=0.6.0,<0.7.0",
+    "sleap-io>=0.6.2,<0.7.0",
     "numpy",
     "lightning",
     "kornia",

{sleap_nn-0.1.0a1 → sleap_nn-0.1.0a2}/sleap_nn/__init__.py

@@ -50,7 +50,7 @@ logger.add(
     colorize=False,
 )
 
-__version__ = "0.1.0a1"
+__version__ = "0.1.0a2"
 
 # Public API
 from sleap_nn.evaluation import load_metrics

{sleap_nn-0.1.0a1 → sleap_nn-0.1.0a2}/sleap_nn/config/trainer_config.py

@@ -208,6 +208,23 @@ class EarlyStoppingConfig:
     stop_training_on_plateau: bool = True
 
 
+@define
+class EvalConfig:
+    """Configuration for epoch-end evaluation.
+
+    Attributes:
+        enabled: (bool) Enable epoch-end evaluation metrics. *Default*: `False`.
+        frequency: (int) Evaluate every N epochs. *Default*: `1`.
+        oks_stddev: (float) OKS standard deviation for evaluation. *Default*: `0.025`.
+        oks_scale: (float) OKS scale override. If None, uses default. *Default*: `None`.
+    """
+
+    enabled: bool = False
+    frequency: int = field(default=1, validator=validators.ge(1))
+    oks_stddev: float = field(default=0.025, validator=validators.gt(0))
+    oks_scale: Optional[float] = None
+
+
 @define
 class HardKeypointMiningConfig:
     """Configuration for online hard keypoint mining.
@@ -310,6 +327,7 @@ class TrainerConfig:
         factory=HardKeypointMiningConfig
     )
     zmq: Optional[ZMQConfig] = field(factory=ZMQConfig)  # Required for SLEAP GUI
+    eval: EvalConfig = field(factory=EvalConfig)  # Epoch-end evaluation config
 
     @staticmethod
     def validate_optimizer_name(value):

{sleap_nn-0.1.0a1 → sleap_nn-0.1.0a2}/sleap_nn/evaluation.py

@@ -29,11 +29,27 @@ def get_instances(labeled_frame: sio.LabeledFrame) -> List[MatchInstance]:
     """
     instance_list = []
     frame_idx = labeled_frame.frame_idx
-    video_path = (
-        labeled_frame.video.backend.source_filename
-        if hasattr(labeled_frame.video.backend, "source_filename")
-        else labeled_frame.video.backend.filename
-    )
+
+    # Extract video path with fallbacks for embedded videos
+    video = labeled_frame.video
+    video_path = None
+    if video is not None:
+        backend = getattr(video, "backend", None)
+        if backend is not None:
+            # Try source_filename first (for embedded videos with provenance)
+            video_path = getattr(backend, "source_filename", None)
+            if video_path is None:
+                video_path = getattr(backend, "filename", None)
+        # Fallback to video.filename if backend doesn't have it
+        if video_path is None:
+            video_path = getattr(video, "filename", None)
+            # Handle list filenames (image sequences)
+            if isinstance(video_path, list) and video_path:
+                video_path = video_path[0]
+    # Final fallback: use a unique identifier
+    if video_path is None:
+        video_path = f"video_{id(video)}" if video is not None else "unknown"
+
     for instance in labeled_frame.instances:
         match_instance = MatchInstance(
             instance=instance, frame_idx=frame_idx, video_path=video_path
@@ -47,6 +63,10 @@ def find_frame_pairs(
 ) -> List[Tuple[sio.LabeledFrame, sio.LabeledFrame]]:
     """Find corresponding frames across two sets of labels.
 
+    This function uses sleap-io's robust video matching API to handle various
+    scenarios including embedded videos, cross-platform paths, and videos with
+    different metadata.
+
     Args:
         labels_gt: A `sio.Labels` instance with ground truth instances.
         labels_pr: A `sio.Labels` instance with predicted instances.
@@ -56,16 +76,15 @@ def find_frame_pairs(
     Returns:
         A list of pairs of `sio.LabeledFrame`s in the form `(frame_gt, frame_pr)`.
     """
+    # Use sleap-io's robust video matching API (added in 0.6.2)
+    # The match() method returns a MatchResult with video_map: {pred_video: gt_video}
+    match_result = labels_gt.match(labels_pr)
+
     frame_pairs = []
-    for video_gt in labels_gt.videos:
-        # Find matching video instance in predictions.
-        video_pr = None
-        for video in labels_pr.videos:
-            if video_gt.matches_content(video) and video_gt.matches_path(video):
-                video_pr = video
-                break
-
-        if video_pr is None:
+    # Iterate over matched video pairs (pred_video -> gt_video mapping)
+    for video_pr, video_gt in match_result.video_map.items():
+        if video_gt is None:
+            # No match found for this prediction video
             continue
 
         # Find labeled frames in this video.
@@ -786,11 +805,26 @@ def run_evaluation(
     """Evaluate SLEAP-NN model predictions against ground truth labels."""
     logger.info("Loading ground truth labels...")
     ground_truth_instances = sio.load_slp(ground_truth_path)
+    logger.info(
+        f"  Ground truth: {len(ground_truth_instances.videos)} videos, "
+        f"{len(ground_truth_instances.labeled_frames)} frames"
+    )
 
     logger.info("Loading predicted labels...")
     predicted_instances = sio.load_slp(predicted_path)
+    logger.info(
+        f"  Predictions: {len(predicted_instances.videos)} videos, "
+        f"{len(predicted_instances.labeled_frames)} frames"
+    )
+
+    logger.info("Matching videos and frames...")
+    # Get match stats before creating evaluator
+    match_result = ground_truth_instances.match(predicted_instances)
+    logger.info(
+        f"  Videos matched: {match_result.n_videos_matched}/{len(match_result.video_map)}"
+    )
 
-    logger.info("Creating evaluator...")
+    logger.info("Matching instances...")
     evaluator = Evaluator(
         ground_truth_instances=ground_truth_instances,
         predicted_instances=predicted_instances,
@@ -799,21 +833,38 @@ def run_evaluation(
         match_threshold=match_threshold,
         user_labels_only=user_labels_only,
     )
+    logger.info(
+        f"  Frame pairs: {len(evaluator.frame_pairs)}, "
+        f"Matched instances: {len(evaluator.positive_pairs)}, "
+        f"Unmatched GT: {len(evaluator.false_negatives)}"
+    )
 
     logger.info("Computing evaluation metrics...")
     metrics = evaluator.evaluate()
 
+    # Compute PCK at specific thresholds (5 and 10 pixels)
+    dists = metrics["distance_metrics"]["dists"]
+    dists_clean = np.copy(dists)
+    dists_clean[np.isnan(dists_clean)] = np.inf
+    pck_5 = (dists_clean < 5).mean()
+    pck_10 = (dists_clean < 10).mean()
+
     # Print key metrics
     logger.info("Evaluation Results:")
-    logger.info(f"mOKS: {metrics['mOKS']['mOKS']:.4f}")
-    logger.info(f"mAP (OKS VOC): {metrics['voc_metrics']['oks_voc.mAP']:.4f}")
-    logger.info(f"mAR (OKS VOC): {metrics['voc_metrics']['oks_voc.mAR']:.4f}")
-    logger.info(f"Average Distance: {metrics['distance_metrics']['avg']:.4f}")
-    logger.info(f"mPCK: {metrics['pck_metrics']['mPCK']:.4f}")
+    logger.info(f"  mOKS: {metrics['mOKS']['mOKS']:.4f}")
+    logger.info(f"  mAP (OKS VOC): {metrics['voc_metrics']['oks_voc.mAP']:.4f}")
+    logger.info(f"  mAR (OKS VOC): {metrics['voc_metrics']['oks_voc.mAR']:.4f}")
+    logger.info(f"  Average Distance: {metrics['distance_metrics']['avg']:.2f} px")
+    logger.info(f"  dist.p50: {metrics['distance_metrics']['p50']:.2f} px")
+    logger.info(f"  dist.p95: {metrics['distance_metrics']['p95']:.2f} px")
+    logger.info(f"  dist.p99: {metrics['distance_metrics']['p99']:.2f} px")
+    logger.info(f"  mPCK: {metrics['pck_metrics']['mPCK']:.4f}")
+    logger.info(f"  PCK@5px: {pck_5:.4f}")
+    logger.info(f"  PCK@10px: {pck_10:.4f}")
     logger.info(
-        f"Visibility Precision: {metrics['visibility_metrics']['precision']:.4f}"
+        f"  Visibility Precision: {metrics['visibility_metrics']['precision']:.4f}"
     )
-    logger.info(f"Visibility Recall: {metrics['visibility_metrics']['recall']:.4f}")
+    logger.info(f"  Visibility Recall: {metrics['visibility_metrics']['recall']:.4f}")
 
     # Save metrics if path provided
     if save_metrics:

{sleap_nn-0.1.0a1 → sleap_nn-0.1.0a2}/sleap_nn/inference/bottomup.py

@@ -1,5 +1,6 @@
 """Inference modules for BottomUp models."""
 
+import logging
 from typing import Dict, Optional
 import torch
 import lightning as L
@@ -7,6 +8,8 @@ from sleap_nn.inference.peak_finding import find_local_peaks
 from sleap_nn.inference.paf_grouping import PAFScorer
 from sleap_nn.inference.identity import classify_peaks_from_maps
 
+logger = logging.getLogger(__name__)
+
 
 class BottomUpInferenceModel(L.LightningModule):
     """BottomUp Inference model.
@@ -63,8 +66,28 @@ class BottomUpInferenceModel(L.LightningModule):
         return_pafs: Optional[bool] = False,
         return_paf_graph: Optional[bool] = False,
         input_scale: float = 1.0,
+        max_peaks_per_node: Optional[int] = None,
     ):
-        """Initialise the model attributes."""
+        """Initialise the model attributes.
+
+        Args:
+            torch_model: A `nn.Module` that accepts images and predicts confidence maps.
+            paf_scorer: A `PAFScorer` instance for grouping instances.
+            cms_output_stride: Output stride of confidence maps relative to images.
+            pafs_output_stride: Output stride of PAFs relative to images.
+            peak_threshold: Minimum confidence map value for valid peaks.
+            refinement: Peak refinement method: None, "integral", or "local".
+            integral_patch_size: Size of patches for integral refinement.
+            return_confmaps: If True, return confidence maps in output.
+            return_pafs: If True, return PAFs in output.
+            return_paf_graph: If True, return intermediate PAF graph in output.
+            input_scale: Scale factor applied to input images.
+            max_peaks_per_node: Maximum number of peaks allowed per node before
+                skipping PAF scoring. If any node has more peaks than this limit,
+                empty predictions are returned. This prevents combinatorial explosion
+                during early training when confidence maps are noisy. Set to None to
+                disable this check (default). Recommended value: 100.
+        """
         super().__init__()
         self.torch_model = torch_model
         self.paf_scorer = paf_scorer
@@ -77,6 +100,7 @@ class BottomUpInferenceModel(L.LightningModule):
         self.return_pafs = return_pafs
         self.return_paf_graph = return_paf_graph
         self.input_scale = input_scale
+        self.max_peaks_per_node = max_peaks_per_node
 
     def _generate_cms_peaks(self, cms):
         # TODO: append nans to batch them -> tensor (vectorize the initial paf grouping steps)
@@ -124,26 +148,68 @@ class BottomUpInferenceModel(L.LightningModule):
         )  # (batch, h, w, 2*edges)
         cms_peaks, cms_peak_vals, cms_peak_channel_inds = self._generate_cms_peaks(cms)
 
-        (
-            predicted_instances,
-            predicted_peak_scores,
-            predicted_instance_scores,
-            edge_inds,
-            edge_peak_inds,
-            line_scores,
-        ) = self.paf_scorer.predict(
-            pafs=pafs,
-            peaks=cms_peaks,
-            peak_vals=cms_peak_vals,
-            peak_channel_inds=cms_peak_channel_inds,
-        )
-
-        predicted_instances = [p / self.input_scale for p in predicted_instances]
-        predicted_instances_adjusted = []
-        for idx, p in enumerate(predicted_instances):
-            predicted_instances_adjusted.append(
-                p / inputs["eff_scale"][idx].to(p.device)
+        # Check if too many peaks per node (prevents combinatorial explosion)
+        skip_paf_scoring = False
+        if self.max_peaks_per_node is not None:
+            n_nodes = cms.shape[1]
+            for b in range(self.batch_size):
+                for node_idx in range(n_nodes):
+                    n_peaks = int((cms_peak_channel_inds[b] == node_idx).sum().item())
+                    if n_peaks > self.max_peaks_per_node:
+                        logger.warning(
+                            f"Skipping PAF scoring: node {node_idx} has {n_peaks} peaks "
+                            f"(max_peaks_per_node={self.max_peaks_per_node}). "
+                            f"Model may need more training."
+                        )
+                        skip_paf_scoring = True
+                        break
+                if skip_paf_scoring:
+                    break
+
+        if skip_paf_scoring:
+            # Return empty predictions for each sample
+            device = cms.device
+            n_nodes = cms.shape[1]
+            predicted_instances_adjusted = []
+            predicted_peak_scores = []
+            predicted_instance_scores = []
+            for _ in range(self.batch_size):
+                predicted_instances_adjusted.append(
+                    torch.full((0, n_nodes, 2), float("nan"), device=device)
+                )
+                predicted_peak_scores.append(
+                    torch.full((0, n_nodes), float("nan"), device=device)
+                )
+                predicted_instance_scores.append(torch.tensor([], device=device))
+            edge_inds = [
+                torch.tensor([], dtype=torch.int32, device=device)
+            ] * self.batch_size
+            edge_peak_inds = [
+                torch.tensor([], dtype=torch.int32, device=device).reshape(0, 2)
+            ] * self.batch_size
+            line_scores = [torch.tensor([], device=device)] * self.batch_size
+        else:
+            (
+                predicted_instances,
+                predicted_peak_scores,
+                predicted_instance_scores,
+                edge_inds,
+                edge_peak_inds,
+                line_scores,
+            ) = self.paf_scorer.predict(
+                pafs=pafs,
+                peaks=cms_peaks,
+                peak_vals=cms_peak_vals,
+                peak_channel_inds=cms_peak_channel_inds,
             )
+
+            predicted_instances = [p / self.input_scale for p in predicted_instances]
+            predicted_instances_adjusted = []
+            for idx, p in enumerate(predicted_instances):
+                predicted_instances_adjusted.append(
+                    p / inputs["eff_scale"][idx].to(p.device)
+                )
+
         out = {
             "pred_instance_peaks": predicted_instances_adjusted,
             "pred_peak_values": predicted_peak_scores,

{sleap_nn-0.1.0a1 → sleap_nn-0.1.0a2}/sleap_nn/training/callbacks.py

@@ -662,3 +662,277 @@ class ProgressReporterZMQ(Callback):
         return {
             k: float(v.item()) if hasattr(v, "item") else v for k, v in logs.items()
         }
+
+
+class EpochEndEvaluationCallback(Callback):
+    """Callback to run full evaluation metrics at end of validation epochs.
+
+    This callback collects predictions and ground truth during validation,
+    then runs the full evaluation pipeline (OKS, mAP, PCK, etc.) and logs
+    metrics to WandB.
+
+    Attributes:
+        skeleton: sio.Skeleton for creating instances.
+        videos: List of sio.Video objects.
+        eval_frequency: Run evaluation every N epochs (default: 1).
+        oks_stddev: OKS standard deviation (default: 0.025).
+        oks_scale: Optional OKS scale override.
+        metrics_to_log: List of metric keys to log.
+    """
+
+    def __init__(
+        self,
+        skeleton: "sio.Skeleton",
+        videos: list,
+        eval_frequency: int = 1,
+        oks_stddev: float = 0.025,
+        oks_scale: Optional[float] = None,
+        metrics_to_log: Optional[list] = None,
+    ):
+        """Initialize the callback.
+
+        Args:
+            skeleton: sio.Skeleton for creating instances.
+            videos: List of sio.Video objects.
+            eval_frequency: Run evaluation every N epochs (default: 1).
+            oks_stddev: OKS standard deviation (default: 0.025).
+            oks_scale: Optional OKS scale override.
+            metrics_to_log: List of metric keys to log. If None, logs all available.
+        """
+        super().__init__()
+        self.skeleton = skeleton
+        self.videos = videos
+        self.eval_frequency = eval_frequency
+        self.oks_stddev = oks_stddev
+        self.oks_scale = oks_scale
+        self.metrics_to_log = metrics_to_log or [
+            "mOKS",
+            "oks_voc.mAP",
+            "oks_voc.mAR",
+            "avg_distance",
+            "p50_distance",
+            "mPCK",
+            "visibility_precision",
+            "visibility_recall",
+        ]
+
+    def on_validation_epoch_start(self, trainer, pl_module):
+        """Enable prediction collection at the start of validation.
+
+        Skip during sanity check to avoid inference issues.
+        """
+        if trainer.sanity_checking:
+            return
+        pl_module._collect_val_predictions = True
+
+    def on_validation_epoch_end(self, trainer, pl_module):
+        """Run evaluation and log metrics at end of validation epoch."""
+        import sleap_io as sio
+        import numpy as np
+        from lightning.pytorch.loggers import WandbLogger
+        from sleap_nn.evaluation import Evaluator
+
+        # Check frequency (epoch is 0-indexed, so add 1)
+        if (trainer.current_epoch + 1) % self.eval_frequency != 0:
+            pl_module._collect_val_predictions = False
+            return
+
+        # Only run on rank 0 for distributed training
+        if not trainer.is_global_zero:
+            pl_module._collect_val_predictions = False
+            return
+
+        # Check if we have predictions
+        if not pl_module.val_predictions or not pl_module.val_ground_truth:
+            logger.warning("No predictions collected for epoch-end evaluation")
+            pl_module._collect_val_predictions = False
+            return
+
+        try:
+            # Build sio.Labels from accumulated predictions and ground truth
+            pred_labels = self._build_pred_labels(pl_module.val_predictions, sio, np)
+            gt_labels = self._build_gt_labels(pl_module.val_ground_truth, sio, np)
+
+            # Check if we have valid frames to evaluate
+            if len(pred_labels) == 0:
+                logger.warning(
+                    "No valid predictions for epoch-end evaluation "
+                    "(all predictions may be empty or NaN)"
+                )
+                pl_module._collect_val_predictions = False
+                pl_module.val_predictions = []
+                pl_module.val_ground_truth = []
+                return
+
+            # Run evaluation
+            evaluator = Evaluator(
+                ground_truth_instances=gt_labels,
+                predicted_instances=pred_labels,
+                oks_stddev=self.oks_stddev,
+                oks_scale=self.oks_scale,
+                user_labels_only=False,  # All validation frames are "user" frames
+            )
+            metrics = evaluator.evaluate()
+
+            # Log to WandB
+            self._log_metrics(trainer, metrics, trainer.current_epoch)
+
+            logger.info(
+                f"Epoch {trainer.current_epoch} evaluation: "
+                f"mOKS={metrics['mOKS']['mOKS']:.4f}, "
+                f"mAP={metrics['voc_metrics']['oks_voc.mAP']:.4f}"
+            )
+
+        except Exception as e:
+            logger.warning(f"Epoch-end evaluation failed: {e}")
+
+        # Cleanup
+        pl_module._collect_val_predictions = False
+        pl_module.val_predictions = []
+        pl_module.val_ground_truth = []
+
+    def _build_pred_labels(self, predictions: list, sio, np) -> "sio.Labels":
+        """Convert prediction dicts to sio.Labels."""
+        labeled_frames = []
+        for pred in predictions:
+            pred_peaks = pred["pred_peaks"]
+            pred_scores = pred["pred_scores"]
+
+            # Handle NaN/missing predictions
+            if pred_peaks is None or (
+                isinstance(pred_peaks, np.ndarray) and np.isnan(pred_peaks).all()
+            ):
+                continue
+
+            # Handle multi-instance predictions (bottomup)
+            if len(pred_peaks.shape) == 2:
+                # Single instance: (n_nodes, 2) -> (1, n_nodes, 2)
+                pred_peaks = pred_peaks.reshape(1, -1, 2)
+                pred_scores = pred_scores.reshape(1, -1)
+
+            instances = []
+            for inst_idx in range(len(pred_peaks)):
+                inst_points = pred_peaks[inst_idx]
+                inst_scores = pred_scores[inst_idx] if pred_scores is not None else None
+
+                # Skip if all NaN
+                if np.isnan(inst_points).all():
+                    continue
+
+                inst = sio.PredictedInstance.from_numpy(
+                    points_data=inst_points,
+                    skeleton=self.skeleton,
+                    point_scores=(
+                        inst_scores
+                        if inst_scores is not None
+                        else np.ones(len(inst_points))
+                    ),
+                    score=(
+                        float(np.nanmean(inst_scores))
+                        if inst_scores is not None
+                        else 1.0
+                    ),
+                )
+                instances.append(inst)
+
+            if instances:
+                lf = sio.LabeledFrame(
+                    video=self.videos[pred["video_idx"]],
+                    frame_idx=pred["frame_idx"],
+                    instances=instances,
+                )
+                labeled_frames.append(lf)
+
+        return sio.Labels(
+            videos=self.videos,
+            skeletons=[self.skeleton],
+            labeled_frames=labeled_frames,
+        )
+
+    def _build_gt_labels(self, ground_truth: list, sio, np) -> "sio.Labels":
+        """Convert ground truth dicts to sio.Labels."""
+        labeled_frames = []
+        for gt in ground_truth:
+            instances = []
+            gt_instances = gt["gt_instances"]
+
+            # Handle shape variations
+            if len(gt_instances.shape) == 2:
+                # (n_nodes, 2) -> (1, n_nodes, 2)
+                gt_instances = gt_instances.reshape(1, -1, 2)
+
+            for i in range(min(gt["num_instances"], len(gt_instances))):
+                inst_data = gt_instances[i]
+                if np.isnan(inst_data).all():
+                    continue
+                inst = sio.Instance.from_numpy(
+                    points_data=inst_data,
+                    skeleton=self.skeleton,
+                )
+                instances.append(inst)
+
+            if instances:
+                lf = sio.LabeledFrame(
+                    video=self.videos[gt["video_idx"]],
+                    frame_idx=gt["frame_idx"],
+                    instances=instances,
+                )
+                labeled_frames.append(lf)
+
+        return sio.Labels(
+            videos=self.videos,
+            skeletons=[self.skeleton],
+            labeled_frames=labeled_frames,
+        )
+
+    def _log_metrics(self, trainer, metrics: dict, epoch: int):
+        """Log evaluation metrics to WandB."""
+        import numpy as np
+        from lightning.pytorch.loggers import WandbLogger
+
+        # Get WandB logger
+        wandb_logger = None
+        for log in trainer.loggers:
+            if isinstance(log, WandbLogger):
+                wandb_logger = log
+                break
+
+        if wandb_logger is None:
+            return
+
+        log_dict = {"epoch": epoch}
+
+        # Extract key metrics with consistent naming
+        if "mOKS" in self.metrics_to_log:
+            log_dict["val_mOKS"] = metrics["mOKS"]["mOKS"]
+
+        if "oks_voc.mAP" in self.metrics_to_log:
+            log_dict["val_oks_voc_mAP"] = metrics["voc_metrics"]["oks_voc.mAP"]
+
+        if "oks_voc.mAR" in self.metrics_to_log:
+            log_dict["val_oks_voc_mAR"] = metrics["voc_metrics"]["oks_voc.mAR"]
+
+        if "avg_distance" in self.metrics_to_log:
+            val = metrics["distance_metrics"]["avg"]
+            if not np.isnan(val):
+                log_dict["val_avg_distance"] = val
+
+        if "p50_distance" in self.metrics_to_log:
+            val = metrics["distance_metrics"]["p50"]
+            if not np.isnan(val):
+                log_dict["val_p50_distance"] = val
+
+        if "mPCK" in self.metrics_to_log:
+            log_dict["val_mPCK"] = metrics["pck_metrics"]["mPCK"]
+
+        if "visibility_precision" in self.metrics_to_log:
+            val = metrics["visibility_metrics"]["precision"]
+            if not np.isnan(val):
+                log_dict["val_visibility_precision"] = val
+
+        if "visibility_recall" in self.metrics_to_log:
+            val = metrics["visibility_metrics"]["recall"]
+            if not np.isnan(val):
+                log_dict["val_visibility_recall"] = val
+
+        wandb_logger.experiment.log(log_dict, commit=False)