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

sleap_nn/training/callbacks.py

@@ -85,15 +85,10 @@ class CSVLoggerCallback(Callback):
                 if key == "epoch":
                     log_data["epoch"] = trainer.current_epoch
                 elif key == "learning_rate":
-                    # Handle multiple formats:
-                    # 1. Direct "learning_rate" key
-                    # 2. "train/lr" key (current format from lightning modules)
-                    # 3. "lr-*" keys from LearningRateMonitor (legacy)
+                    # Handle both direct logging and LearningRateMonitor format (lr-*)
                     value = metrics.get(key, None)
                     if value is None:
-                        value = metrics.get("train/lr", None)
-                    if value is None:
-                        # Look for lr-* keys from LearningRateMonitor (legacy)
+                        # Look for lr-* keys from LearningRateMonitor
                         for metric_key in metrics.keys():
                             if metric_key.startswith("lr-"):
                                 value = metrics[metric_key]
@@ -286,49 +281,45 @@ class WandBVizCallback(Callback):
 
             # Get the wandb logger to use its experiment for logging
             wandb_logger = self._get_wandb_logger(trainer)
+            if wandb_logger is None:
+                return  # No wandb logger, skip visualization logging
+
+            # Get visualization data
+            train_data = self.train_viz_fn()
+            val_data = self.val_viz_fn()
+
+            # Render and log for each enabled mode
+            # Use the logger's experiment to let Lightning manage step tracking
+            log_dict = {}
+            for mode_name, renderer in self.renderers.items():
+                suffix = "" if mode_name == "direct" else f"_{mode_name}"
+                train_img = renderer.render(train_data, caption=f"Train Epoch {epoch}")
+                val_img = renderer.render(val_data, caption=f"Val Epoch {epoch}")
+                log_dict[f"train_predictions{suffix}"] = train_img
+                log_dict[f"val_predictions{suffix}"] = val_img
+
+            if log_dict:
+                # Include epoch so wandb can use it as x-axis (via define_metric)
+                log_dict["epoch"] = epoch
+                # Use commit=False to accumulate with other metrics in this step
+                # Lightning will commit when it logs its own metrics
+                wandb_logger.experiment.log(log_dict, commit=False)
+
+            # Optionally also log to table for backwards compat
+            if self.log_table and "direct" in self.renderers:
+                train_img = self.renderers["direct"].render(
+                    train_data, caption=f"Train Epoch {epoch}"
+                )
+                val_img = self.renderers["direct"].render(
+                    val_data, caption=f"Val Epoch {epoch}"
+                )
+                table = wandb.Table(
+                    columns=["Epoch", "Train", "Validation"],
+                    data=[[epoch, train_img, val_img]],
+                )
+                wandb_logger.experiment.log({"predictions_table": table}, commit=False)
 
-            # Only do visualization work if wandb logger is available
-            if wandb_logger is not None:
-                # Get visualization data
-                train_data = self.train_viz_fn()
-                val_data = self.val_viz_fn()
-
-                # Render and log for each enabled mode
-                # Use the logger's experiment to let Lightning manage step tracking
-                log_dict = {}
-                for mode_name, renderer in self.renderers.items():
-                    suffix = "" if mode_name == "direct" else f"_{mode_name}"
-                    train_img = renderer.render(
-                        train_data, caption=f"Train Epoch {epoch}"
-                    )
-                    val_img = renderer.render(val_data, caption=f"Val Epoch {epoch}")
-                    log_dict[f"viz/train/predictions{suffix}"] = train_img
-                    log_dict[f"viz/val/predictions{suffix}"] = val_img
-
-                if log_dict:
-                    # Include epoch so wandb can use it as x-axis (via define_metric)
-                    log_dict["epoch"] = epoch
-                    # Use commit=False to accumulate with other metrics in this step
-                    # Lightning will commit when it logs its own metrics
-                    wandb_logger.experiment.log(log_dict, commit=False)
-
-                # Optionally also log to table for backwards compat
-                if self.log_table and "direct" in self.renderers:
-                    train_img = self.renderers["direct"].render(
-                        train_data, caption=f"Train Epoch {epoch}"
-                    )
-                    val_img = self.renderers["direct"].render(
-                        val_data, caption=f"Val Epoch {epoch}"
-                    )
-                    table = wandb.Table(
-                        columns=["Epoch", "Train", "Validation"],
-                        data=[[epoch, train_img, val_img]],
-                    )
-                    wandb_logger.experiment.log(
-                        {"predictions_table": table}, commit=False
-                    )
-
-        # Sync all processes - barrier must be reached by ALL ranks
+        # Sync all processes
         trainer.strategy.barrier()
 
 
@@ -387,377 +378,80 @@ class WandBVizCallbackWithPAFs(WandBVizCallback):
 
             # Get the wandb logger to use its experiment for logging
             wandb_logger = self._get_wandb_logger(trainer)
-
-            # Only do visualization work if wandb logger is available
-            if wandb_logger is not None:
-                # Get visualization data
-                train_data = self.train_viz_fn()
-                val_data = self.val_viz_fn()
-                train_pafs_data = self.train_pafs_viz_fn()
-                val_pafs_data = self.val_pafs_viz_fn()
-
-                # Render and log for each enabled mode
-                # Use the logger's experiment to let Lightning manage step tracking
-                log_dict = {}
-                for mode_name, renderer in self.renderers.items():
-                    suffix = "" if mode_name == "direct" else f"_{mode_name}"
-                    train_img = renderer.render(
-                        train_data, caption=f"Train Epoch {epoch}"
-                    )
-                    val_img = renderer.render(val_data, caption=f"Val Epoch {epoch}")
-                    log_dict[f"viz/train/predictions{suffix}"] = train_img
-                    log_dict[f"viz/val/predictions{suffix}"] = val_img
-
-                # Render PAFs (always use matplotlib/direct for PAFs)
-                from io import BytesIO
-                import matplotlib.pyplot as plt
-                from PIL import Image
-
-                train_pafs_fig = self._mpl_renderer.render_pafs(train_pafs_data)
-                buf = BytesIO()
-                train_pafs_fig.savefig(
-                    buf, format="png", bbox_inches="tight", pad_inches=0
-                )
-                buf.seek(0)
-                plt.close(train_pafs_fig)
-                train_pafs_pil = Image.open(buf)
-                log_dict["viz/train/pafs"] = wandb.Image(
-                    train_pafs_pil, caption=f"Train PAFs Epoch {epoch}"
-                )
-
-                val_pafs_fig = self._mpl_renderer.render_pafs(val_pafs_data)
-                buf = BytesIO()
-                val_pafs_fig.savefig(
-                    buf, format="png", bbox_inches="tight", pad_inches=0
-                )
-                buf.seek(0)
-                plt.close(val_pafs_fig)
-                val_pafs_pil = Image.open(buf)
-                log_dict["viz/val/pafs"] = wandb.Image(
-                    val_pafs_pil, caption=f"Val PAFs Epoch {epoch}"
-                )
-
-                if log_dict:
-                    # Include epoch so wandb can use it as x-axis (via define_metric)
-                    log_dict["epoch"] = epoch
-                    # Use commit=False to accumulate with other metrics in this step
-                    # Lightning will commit when it logs its own metrics
-                    wandb_logger.experiment.log(log_dict, commit=False)
-
-                # Optionally also log to table
-                if self.log_table and "direct" in self.renderers:
-                    train_img = self.renderers["direct"].render(
-                        train_data, caption=f"Train Epoch {epoch}"
-                    )
-                    val_img = self.renderers["direct"].render(
-                        val_data, caption=f"Val Epoch {epoch}"
-                    )
-                    table = wandb.Table(
-                        columns=[
-                            "Epoch",
-                            "Train",
-                            "Validation",
-                            "Train PAFs",
-                            "Val PAFs",
-                        ],
-                        data=[
-                            [
-                                epoch,
-                                train_img,
-                                val_img,
-                                log_dict["viz/train/pafs"],
-                                log_dict["viz/val/pafs"],
-                            ]
-                        ],
-                    )
-                    wandb_logger.experiment.log(
-                        {"predictions_table": table}, commit=False
-                    )
-
-        # Sync all processes - barrier must be reached by ALL ranks
-        trainer.strategy.barrier()
-
-
-class UnifiedVizCallback(Callback):
-    """Unified callback for all visualization outputs during training.
-
-    This callback consolidates all visualization functionality into a single callback,
-    eliminating redundant dataset copies and inference runs. It handles:
-    - Local disk saving (matplotlib figures)
-    - WandB logging (multiple modes: direct, boxes, masks)
-    - Model-specific visualizations (PAFs for bottomup, class maps for multi_class_bottomup)
-
-    Benefits over separate callbacks:
-    - Uses ONE sample per epoch for all visualizations (no dataset deepcopy)
-    - Runs inference ONCE per sample (vs 4-8x in previous implementation)
-    - Outputs to multiple destinations from the same data
-    - Simpler code with less duplication
-
-    Attributes:
-        model_trainer: Reference to the ModelTrainer (for lazy access to lightning_model).
-        train_pipeline: Iterator over training visualization dataset.
-        val_pipeline: Iterator over validation visualization dataset.
-        model_type: Type of model (affects which visualizations are enabled).
-        save_local: Whether to save matplotlib figures to disk.
-        local_save_dir: Directory for local visualization saves.
-        log_wandb: Whether to log visualizations to wandb.
-        wandb_modes: List of wandb rendering modes ("direct", "boxes", "masks").
-        wandb_box_size: Size of keypoint boxes in pixels (for "boxes" mode).
-        wandb_confmap_threshold: Threshold for confmap masks (for "masks" mode).
-        log_wandb_table: Whether to also log to a wandb.Table.
-    """
-
-    def __init__(
-        self,
-        model_trainer,
-        train_dataset,
-        val_dataset,
-        model_type: str,
-        save_local: bool = True,
-        local_save_dir: Optional[Path] = None,
-        log_wandb: bool = False,
-        wandb_modes: Optional[list] = None,
-        wandb_box_size: float = 5.0,
-        wandb_confmap_threshold: float = 0.1,
-        log_wandb_table: bool = False,
-    ):
-        """Initialize the unified visualization callback.
-
-        Args:
-            model_trainer: ModelTrainer instance (lightning_model accessed lazily).
-            train_dataset: Training visualization dataset (will be cycled).
-            val_dataset: Validation visualization dataset (will be cycled).
-            model_type: Model type string (e.g., "bottomup", "multi_class_bottomup").
-            save_local: If True, save matplotlib figures to local_save_dir.
-            local_save_dir: Path to directory for saving visualization images.
-            log_wandb: If True, log visualizations to wandb.
-            wandb_modes: List of wandb rendering modes. Defaults to ["direct"].
-            wandb_box_size: Size of keypoint boxes in pixels.
-            wandb_confmap_threshold: Threshold for confidence map masks.
-            log_wandb_table: If True, also log to a wandb.Table.
-        """
-        super().__init__()
-        from itertools import cycle
-
-        self.model_trainer = model_trainer
-        self.train_pipeline = cycle(train_dataset)
-        self.val_pipeline = cycle(val_dataset)
-        self.model_type = model_type
-
-        # Local disk config
-        self.save_local = save_local
-        self.local_save_dir = local_save_dir
-
-        # WandB config
-        self.log_wandb = log_wandb
-        self.wandb_modes = wandb_modes or ["direct"]
-        self.wandb_box_size = wandb_box_size
-        self.wandb_confmap_threshold = wandb_confmap_threshold
-        self.log_wandb_table = log_wandb_table
-
-        # Auto-enable model-specific visualizations
-        self.viz_pafs = model_type == "bottomup"
-        self.viz_class_maps = model_type == "multi_class_bottomup"
-
-        # Initialize renderers
-        from sleap_nn.training.utils import MatplotlibRenderer, WandBRenderer
-
-        self._mpl_renderer = MatplotlibRenderer()
-
-        # Create wandb renderers for each enabled mode
-        self._wandb_renderers = {}
-        if log_wandb:
-            for mode in self.wandb_modes:
-                self._wandb_renderers[mode] = WandBRenderer(
-                    mode=mode,
-                    box_size=wandb_box_size,
-                    confmap_threshold=wandb_confmap_threshold,
-                )
-
-    def _get_wandb_logger(self, trainer):
-        """Get the WandbLogger from trainer's loggers."""
-        from lightning.pytorch.loggers import WandbLogger
-
-        for log in trainer.loggers:
-            if isinstance(log, WandbLogger):
-                return log
-        return None
-
-    def _get_viz_data(self, sample):
-        """Get visualization data with all needed fields based on model type.
-
-        Args:
-            sample: A sample from the visualization dataset.
-
-        Returns:
-            VisualizationData with appropriate fields populated.
-        """
-        # Build kwargs based on model type
-        kwargs = {}
-        if self.viz_pafs:
-            kwargs["include_pafs"] = True
-        if self.viz_class_maps:
-            kwargs["include_class_maps"] = True
-
-        # Access lightning_model lazily from model_trainer
-        return self.model_trainer.lightning_model.get_visualization_data(
-            sample, **kwargs
-        )
-
-    def _save_local_viz(self, data, prefix: str, epoch: int):
-        """Save visualization to local disk.
-
-        Args:
-            data: VisualizationData object.
-            prefix: Filename prefix (e.g., "train", "validation").
-            epoch: Current epoch number.
-        """
-        if not self.save_local or self.local_save_dir is None:
-            return
-
-        # Confmaps visualization
-        fig = self._mpl_renderer.render(data)
-        fig_path = self.local_save_dir / f"{prefix}.{epoch:04d}.png"
-        fig.savefig(fig_path, format="png")
-        plt.close(fig)
-
-        # PAFs visualization (for bottomup models)
-        if self.viz_pafs and data.pred_pafs is not None:
-            fig = self._mpl_renderer.render_pafs(data)
-            fig_path = self.local_save_dir / f"{prefix}.pafs_magnitude.{epoch:04d}.png"
-            fig.savefig(fig_path, format="png")
-            plt.close(fig)
-
-        # Class maps visualization (for multi_class_bottomup models)
-        if self.viz_class_maps and data.pred_class_maps is not None:
-            fig = self._render_class_maps(data)
-            fig_path = self.local_save_dir / f"{prefix}.class_maps.{epoch:04d}.png"
-            fig.savefig(fig_path, format="png")
-            plt.close(fig)
-
-    def _render_class_maps(self, data):
-        """Render class maps visualization.
-
-        Args:
-            data: VisualizationData with pred_class_maps populated.
-
-        Returns:
-            A matplotlib Figure object.
-        """
-        from sleap_nn.training.utils import plot_img, plot_confmaps
-
-        img = data.image
-        scale = 1.0
-        if img.shape[0] < 512:
-            scale = 2.0
-        if img.shape[0] < 256:
-            scale = 4.0
-
-        fig = plot_img(img, dpi=72 * scale, scale=scale)
-        plot_confmaps(
-            data.pred_class_maps,
-            output_scale=data.pred_class_maps.shape[0] / img.shape[0],
-        )
-        return fig
-
-    def _log_wandb_viz(self, data, prefix: str, epoch: int, wandb_logger):
-        """Log visualization to wandb.
-
-        Args:
-            data: VisualizationData object.
-            prefix: Log prefix (e.g., "train", "val").
-            epoch: Current epoch number.
-            wandb_logger: WandbLogger instance.
-        """
-        if not self.log_wandb or wandb_logger is None:
-            return
-
-        from io import BytesIO
-        from PIL import Image as PILImage
-
-        log_dict = {}
-
-        # Render confmaps for each enabled mode
-        for mode_name, renderer in self._wandb_renderers.items():
-            suffix = "" if mode_name == "direct" else f"_{mode_name}"
-            img = renderer.render(data, caption=f"{prefix.title()} Epoch {epoch}")
-            log_dict[f"viz/{prefix}/predictions{suffix}"] = img
-
-        # PAFs visualization (for bottomup models)
-        if self.viz_pafs and data.pred_pafs is not None:
-            pafs_fig = self._mpl_renderer.render_pafs(data)
+            if wandb_logger is None:
+                return  # No wandb logger, skip visualization logging
+
+            # Get visualization data
+            train_data = self.train_viz_fn()
+            val_data = self.val_viz_fn()
+            train_pafs_data = self.train_pafs_viz_fn()
+            val_pafs_data = self.val_pafs_viz_fn()
+
+            # Render and log for each enabled mode
+            # Use the logger's experiment to let Lightning manage step tracking
+            log_dict = {}
+            for mode_name, renderer in self.renderers.items():
+                suffix = "" if mode_name == "direct" else f"_{mode_name}"
+                train_img = renderer.render(train_data, caption=f"Train Epoch {epoch}")
+                val_img = renderer.render(val_data, caption=f"Val Epoch {epoch}")
+                log_dict[f"train_predictions{suffix}"] = train_img
+                log_dict[f"val_predictions{suffix}"] = val_img
+
+            # Render PAFs (always use matplotlib/direct for PAFs)
+            from io import BytesIO
+            import matplotlib.pyplot as plt
+            from PIL import Image
+
+            train_pafs_fig = self._mpl_renderer.render_pafs(train_pafs_data)
             buf = BytesIO()
-            pafs_fig.savefig(buf, format="png", bbox_inches="tight", pad_inches=0)
+            train_pafs_fig.savefig(buf, format="png", bbox_inches="tight", pad_inches=0)
             buf.seek(0)
-            plt.close(pafs_fig)
-            pafs_pil = PILImage.open(buf)
-            log_dict[f"viz/{prefix}/pafs"] = wandb.Image(
-                pafs_pil, caption=f"{prefix.title()} PAFs Epoch {epoch}"
+            plt.close(train_pafs_fig)
+            train_pafs_pil = Image.open(buf)
+            log_dict["train_pafs"] = wandb.Image(
+                train_pafs_pil, caption=f"Train PAFs Epoch {epoch}"
             )
 
-        # Class maps visualization (for multi_class_bottomup models)
-        if self.viz_class_maps and data.pred_class_maps is not None:
-            class_fig = self._render_class_maps(data)
+            val_pafs_fig = self._mpl_renderer.render_pafs(val_pafs_data)
             buf = BytesIO()
-            class_fig.savefig(buf, format="png", bbox_inches="tight", pad_inches=0)
+            val_pafs_fig.savefig(buf, format="png", bbox_inches="tight", pad_inches=0)
             buf.seek(0)
-            plt.close(class_fig)
-            class_pil = PILImage.open(buf)
-            log_dict[f"viz/{prefix}/class_maps"] = wandb.Image(
-                class_pil, caption=f"{prefix.title()} Class Maps Epoch {epoch}"
-            )
-
-        if log_dict:
-            log_dict["epoch"] = epoch
-            wandb_logger.experiment.log(log_dict, commit=False)
-
-        # Optionally log to table for backwards compatibility
-        if self.log_wandb_table and "direct" in self._wandb_renderers:
-            train_img = self._wandb_renderers["direct"].render(
-                data, caption=f"{prefix.title()} Epoch {epoch}"
-            )
-            table_data = [[epoch, train_img]]
-            columns = ["Epoch", prefix.title()]
-
-            if self.viz_pafs and data.pred_pafs is not None:
-                columns.append(f"{prefix.title()} PAFs")
-                table_data[0].append(log_dict.get(f"viz/{prefix}/pafs"))
-
-            if self.viz_class_maps and data.pred_class_maps is not None:
-                columns.append(f"{prefix.title()} Class Maps")
-                table_data[0].append(log_dict.get(f"viz/{prefix}/class_maps"))
-
-            table = wandb.Table(columns=columns, data=table_data)
-            wandb_logger.experiment.log(
-                {f"predictions_table_{prefix}": table}, commit=False
+            plt.close(val_pafs_fig)
+            val_pafs_pil = Image.open(buf)
+            log_dict["val_pafs"] = wandb.Image(
+                val_pafs_pil, caption=f"Val PAFs Epoch {epoch}"
             )
 
-    def on_train_epoch_end(self, trainer, pl_module):
-        """Generate and output all visualizations at epoch end.
+            if log_dict:
+                # Include epoch so wandb can use it as x-axis (via define_metric)
+                log_dict["epoch"] = epoch
+                # Use commit=False to accumulate with other metrics in this step
+                # Lightning will commit when it logs its own metrics
+                wandb_logger.experiment.log(log_dict, commit=False)
+
+            # Optionally also log to table
+            if self.log_table and "direct" in self.renderers:
+                train_img = self.renderers["direct"].render(
+                    train_data, caption=f"Train Epoch {epoch}"
+                )
+                val_img = self.renderers["direct"].render(
+                    val_data, caption=f"Val Epoch {epoch}"
+                )
+                table = wandb.Table(
+                    columns=["Epoch", "Train", "Validation", "Train PAFs", "Val PAFs"],
+                    data=[
+                        [
+                            epoch,
+                            train_img,
+                            val_img,
+                            log_dict["train_pafs"],
+                            log_dict["val_pafs"],
+                        ]
+                    ],
+                )
+                wandb_logger.experiment.log({"predictions_table": table}, commit=False)
 
-        Args:
-            trainer: PyTorch Lightning trainer.
-            pl_module: Lightning module (not used, we use self.lightning_module).
-        """
-        if trainer.is_global_zero:
-            epoch = trainer.current_epoch
-            wandb_logger = self._get_wandb_logger(trainer) if self.log_wandb else None
-
-            # Get ONE sample for train visualization
-            train_sample = next(self.train_pipeline)
-            # Run inference ONCE with all needed data
-            train_data = self._get_viz_data(train_sample)
-            # Output to all destinations
-            self._save_local_viz(train_data, "train", epoch)
-            self._log_wandb_viz(train_data, "train", epoch, wandb_logger)
-
-            # Same for validation
-            val_sample = next(self.val_pipeline)
-            val_data = self._get_viz_data(val_sample)
-            self._save_local_viz(val_data, "validation", epoch)
-            self._log_wandb_viz(val_data, "val", epoch, wandb_logger)
-
-        # Sync all processes - barrier must be reached by ALL ranks
+        # Sync all processes
         trainer.strategy.barrier()
 
 
@@ -968,638 +662,3 @@ 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",
-            "distance/avg",
-            "distance/p50",
-            "distance/p95",
-            "distance/p99",
-            "mPCK",
-            "PCK@5",
-            "PCK@10",
-            "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
-
-        # Determine if we should run evaluation this epoch (only on rank 0)
-        should_evaluate = (
-            trainer.current_epoch + 1
-        ) % self.eval_frequency == 0 and trainer.is_global_zero
-
-        if should_evaluate:
-            # 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")
-            else:
-                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)"
-                        )
-                    else:
-                        # 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"PCK@5={metrics['pck_metrics']['PCK@5']:.4f}, "
-                            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 - all ranks reset the flag, rank 0 clears the lists
-        pl_module._collect_val_predictions = False
-        if trainer.is_global_zero:
-            pl_module.val_predictions = []
-            pl_module.val_ground_truth = []
-
-        # Sync all processes - barrier must be reached by ALL ranks
-        trainer.strategy.barrier()
-
-    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
-        # All eval metrics use eval/val/ prefix since they're computed on validation data
-        if "mOKS" in self.metrics_to_log:
-            log_dict["eval/val/mOKS"] = metrics["mOKS"]["mOKS"]
-
-        if "oks_voc.mAP" in self.metrics_to_log:
-            log_dict["eval/val/oks_voc_mAP"] = metrics["voc_metrics"]["oks_voc.mAP"]
-
-        if "oks_voc.mAR" in self.metrics_to_log:
-            log_dict["eval/val/oks_voc_mAR"] = metrics["voc_metrics"]["oks_voc.mAR"]
-
-        # Distance metrics grouped under eval/val/distance/
-        if "distance/avg" in self.metrics_to_log:
-            val = metrics["distance_metrics"]["avg"]
-            if not np.isnan(val):
-                log_dict["eval/val/distance/avg"] = val
-
-        if "distance/p50" in self.metrics_to_log:
-            val = metrics["distance_metrics"]["p50"]
-            if not np.isnan(val):
-                log_dict["eval/val/distance/p50"] = val
-
-        if "distance/p95" in self.metrics_to_log:
-            val = metrics["distance_metrics"]["p95"]
-            if not np.isnan(val):
-                log_dict["eval/val/distance/p95"] = val
-
-        if "distance/p99" in self.metrics_to_log:
-            val = metrics["distance_metrics"]["p99"]
-            if not np.isnan(val):
-                log_dict["eval/val/distance/p99"] = val
-
-        # PCK metrics
-        if "mPCK" in self.metrics_to_log:
-            log_dict["eval/val/mPCK"] = metrics["pck_metrics"]["mPCK"]
-
-        # PCK at specific thresholds (precomputed in evaluation.py)
-        if "PCK@5" in self.metrics_to_log:
-            log_dict["eval/val/PCK_5"] = metrics["pck_metrics"]["PCK@5"]
-
-        if "PCK@10" in self.metrics_to_log:
-            log_dict["eval/val/PCK_10"] = metrics["pck_metrics"]["PCK@10"]
-
-        # Visibility metrics
-        if "visibility_precision" in self.metrics_to_log:
-            val = metrics["visibility_metrics"]["precision"]
-            if not np.isnan(val):
-                log_dict["eval/val/visibility_precision"] = val
-
-        if "visibility_recall" in self.metrics_to_log:
-            val = metrics["visibility_metrics"]["recall"]
-            if not np.isnan(val):
-                log_dict["eval/val/visibility_recall"] = val
-
-        wandb_logger.experiment.log(log_dict, commit=False)
-
-        # Update best metrics in summary (excluding epoch)
-        for key, value in log_dict.items():
-            if key == "epoch":
-                continue
-            # Create summary key like "best/eval/val/mOKS"
-            summary_key = f"best/{key}"
-            current_best = wandb_logger.experiment.summary.get(summary_key)
-            # For distance metrics, lower is better; for others, higher is better
-            is_distance = "distance" in key
-            if current_best is None:
-                wandb_logger.experiment.summary[summary_key] = value
-            elif is_distance and value < current_best:
-                wandb_logger.experiment.summary[summary_key] = value
-            elif not is_distance and value > current_best:
-                wandb_logger.experiment.summary[summary_key] = value
-
-
-def match_centroids(
-    pred_centroids: "np.ndarray",
-    gt_centroids: "np.ndarray",
-    max_distance: float = 50.0,
-) -> tuple:
-    """Match predicted centroids to ground truth using Hungarian algorithm.
-
-    Args:
-        pred_centroids: Predicted centroid locations, shape (n_pred, 2).
-        gt_centroids: Ground truth centroid locations, shape (n_gt, 2).
-        max_distance: Maximum distance threshold for valid matches (in pixels).
-
-    Returns:
-        Tuple of:
-            - matched_pred_indices: Indices of matched predictions
-            - matched_gt_indices: Indices of matched ground truth
-            - unmatched_pred_indices: Indices of unmatched predictions (false positives)
-            - unmatched_gt_indices: Indices of unmatched ground truth (false negatives)
-    """
-    import numpy as np
-    from scipy.optimize import linear_sum_assignment
-    from scipy.spatial.distance import cdist
-
-    n_pred = len(pred_centroids)
-    n_gt = len(gt_centroids)
-
-    # Handle edge cases
-    if n_pred == 0 and n_gt == 0:
-        return np.array([]), np.array([]), np.array([]), np.array([])
-    if n_pred == 0:
-        return np.array([]), np.array([]), np.array([]), np.arange(n_gt)
-    if n_gt == 0:
-        return np.array([]), np.array([]), np.arange(n_pred), np.array([])
-
-    # Compute pairwise distances
-    cost_matrix = cdist(pred_centroids, gt_centroids)
-
-    # Run Hungarian algorithm for optimal matching
-    pred_indices, gt_indices = linear_sum_assignment(cost_matrix)
-
-    # Filter matches that exceed max_distance
-    matched_pred = []
-    matched_gt = []
-    for p_idx, g_idx in zip(pred_indices, gt_indices):
-        if cost_matrix[p_idx, g_idx] <= max_distance:
-            matched_pred.append(p_idx)
-            matched_gt.append(g_idx)
-
-    matched_pred = np.array(matched_pred)
-    matched_gt = np.array(matched_gt)
-
-    # Find unmatched indices
-    all_pred = set(range(n_pred))
-    all_gt = set(range(n_gt))
-    unmatched_pred = np.array(list(all_pred - set(matched_pred)))
-    unmatched_gt = np.array(list(all_gt - set(matched_gt)))
-
-    return matched_pred, matched_gt, unmatched_pred, unmatched_gt
-
-
-class CentroidEvaluationCallback(Callback):
-    """Callback to run centroid-specific evaluation metrics at end of validation epochs.
-
-    This callback is designed specifically for centroid models, which predict a single
-    point (centroid) per instance rather than full pose skeletons. It computes
-    distance-based metrics and detection metrics that are more appropriate for
-    point detection tasks than OKS/PCK metrics.
-
-    Metrics computed:
-        - Distance metrics: mean, median, p90, p95, max Euclidean distance
-        - Detection metrics: precision, recall, F1 score
-        - Counts: true positives, false positives, false negatives
-
-    Attributes:
-        videos: List of sio.Video objects.
-        eval_frequency: Run evaluation every N epochs (default: 1).
-        match_threshold: Maximum distance (pixels) for matching pred to GT (default: 50.0).
-    """
-
-    def __init__(
-        self,
-        videos: list,
-        eval_frequency: int = 1,
-        match_threshold: float = 50.0,
-    ):
-        """Initialize the callback.
-
-        Args:
-            videos: List of sio.Video objects.
-            eval_frequency: Run evaluation every N epochs (default: 1).
-            match_threshold: Maximum distance in pixels for a prediction to be
-                considered a match to a ground truth centroid (default: 50.0).
-        """
-        super().__init__()
-        self.videos = videos
-        self.eval_frequency = eval_frequency
-        self.match_threshold = match_threshold
-
-    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 centroid evaluation and log metrics at end of validation epoch."""
-        import numpy as np
-        from lightning.pytorch.loggers import WandbLogger
-
-        # Determine if we should run evaluation this epoch (only on rank 0)
-        should_evaluate = (
-            trainer.current_epoch + 1
-        ) % self.eval_frequency == 0 and trainer.is_global_zero
-
-        if should_evaluate:
-            # Check if we have predictions
-            if not pl_module.val_predictions or not pl_module.val_ground_truth:
-                logger.warning(
-                    "No predictions collected for centroid epoch-end evaluation"
-                )
-            else:
-                try:
-                    metrics = self._compute_metrics(
-                        pl_module.val_predictions, pl_module.val_ground_truth, np
-                    )
-
-                    # Log to WandB
-                    self._log_metrics(trainer, metrics, trainer.current_epoch)
-
-                    logger.info(
-                        f"Epoch {trainer.current_epoch} centroid evaluation: "
-                        f"precision={metrics['precision']:.4f}, "
-                        f"recall={metrics['recall']:.4f}, "
-                        f"dist_avg={metrics['dist_avg']:.2f}px"
-                    )
-
-                except Exception as e:
-                    logger.warning(f"Centroid epoch-end evaluation failed: {e}")
-
-        # Cleanup - all ranks reset the flag, rank 0 clears the lists
-        pl_module._collect_val_predictions = False
-        if trainer.is_global_zero:
-            pl_module.val_predictions = []
-            pl_module.val_ground_truth = []
-
-        # Sync all processes - barrier must be reached by ALL ranks
-        trainer.strategy.barrier()
-
-    def _compute_metrics(self, predictions: list, ground_truth: list, np) -> dict:
-        """Compute centroid-specific metrics.
-
-        Args:
-            predictions: List of prediction dicts with "pred_peaks" key.
-            ground_truth: List of ground truth dicts with "gt_instances" key.
-            np: NumPy module.
-
-        Returns:
-            Dictionary of computed metrics.
-        """
-        all_distances = []
-        total_tp = 0
-        total_fp = 0
-        total_fn = 0
-
-        # Group predictions and GT by frame
-        pred_by_frame = {}
-        for pred in predictions:
-            key = (pred["video_idx"], pred["frame_idx"])
-            if key not in pred_by_frame:
-                pred_by_frame[key] = []
-            # pred_peaks shape: (n_inst, 1, 2) -> extract centroids as (n_inst, 2)
-            centroids = pred["pred_peaks"].reshape(-1, 2)
-            # Filter out NaN centroids
-            valid_mask = ~np.isnan(centroids).any(axis=1)
-            pred_by_frame[key].append(centroids[valid_mask])
-
-        gt_by_frame = {}
-        for gt in ground_truth:
-            key = (gt["video_idx"], gt["frame_idx"])
-            if key not in gt_by_frame:
-                gt_by_frame[key] = []
-            # gt_instances shape: (n_inst, 1, 2) -> extract centroids as (n_inst, 2)
-            centroids = gt["gt_instances"].reshape(-1, 2)
-            # Filter out NaN centroids
-            valid_mask = ~np.isnan(centroids).any(axis=1)
-            gt_by_frame[key].append(centroids[valid_mask])
-
-        # Process each frame
-        all_frames = set(pred_by_frame.keys()) | set(gt_by_frame.keys())
-        for frame_key in all_frames:
-            # Concatenate all predictions for this frame
-            if frame_key in pred_by_frame:
-                frame_preds = np.concatenate(pred_by_frame[frame_key], axis=0)
-            else:
-                frame_preds = np.zeros((0, 2))
-
-            # Concatenate all GT for this frame
-            if frame_key in gt_by_frame:
-                frame_gt = np.concatenate(gt_by_frame[frame_key], axis=0)
-            else:
-                frame_gt = np.zeros((0, 2))
-
-            # Match predictions to ground truth
-            matched_pred, matched_gt, unmatched_pred, unmatched_gt = match_centroids(
-                frame_preds, frame_gt, max_distance=self.match_threshold
-            )
-
-            # Compute distances for matched pairs
-            if len(matched_pred) > 0:
-                matched_pred_points = frame_preds[matched_pred]
-                matched_gt_points = frame_gt[matched_gt]
-                distances = np.linalg.norm(
-                    matched_pred_points - matched_gt_points, axis=1
-                )
-                all_distances.extend(distances.tolist())
-
-            # Update counts
-            total_tp += len(matched_pred)
-            total_fp += len(unmatched_pred)
-            total_fn += len(unmatched_gt)
-
-        # Compute aggregate metrics
-        all_distances = np.array(all_distances)
-
-        # Distance metrics (only if we have matches)
-        if len(all_distances) > 0:
-            dist_avg = float(np.mean(all_distances))
-            dist_median = float(np.median(all_distances))
-            dist_p90 = float(np.percentile(all_distances, 90))
-            dist_p95 = float(np.percentile(all_distances, 95))
-            dist_max = float(np.max(all_distances))
-        else:
-            dist_avg = dist_median = dist_p90 = dist_p95 = dist_max = float("nan")
-
-        # Detection metrics
-        precision = (
-            total_tp / (total_tp + total_fp) if (total_tp + total_fp) > 0 else 0.0
-        )
-        recall = total_tp / (total_tp + total_fn) if (total_tp + total_fn) > 0 else 0.0
-        f1 = (
-            2 * precision * recall / (precision + recall)
-            if (precision + recall) > 0
-            else 0.0
-        )
-
-        return {
-            "dist_avg": dist_avg,
-            "dist_median": dist_median,
-            "dist_p90": dist_p90,
-            "dist_p95": dist_p95,
-            "dist_max": dist_max,
-            "precision": precision,
-            "recall": recall,
-            "f1": f1,
-            "n_true_positives": total_tp,
-            "n_false_positives": total_fp,
-            "n_false_negatives": total_fn,
-            "n_total_predictions": total_tp + total_fp,
-            "n_total_ground_truth": total_tp + total_fn,
-        }
-
-    def _log_metrics(self, trainer, metrics: dict, epoch: int):
-        """Log centroid 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}
-
-        # Distance metrics (with NaN handling)
-        if not np.isnan(metrics["dist_avg"]):
-            log_dict["eval/val/centroid_dist_avg"] = metrics["dist_avg"]
-        if not np.isnan(metrics["dist_median"]):
-            log_dict["eval/val/centroid_dist_median"] = metrics["dist_median"]
-        if not np.isnan(metrics["dist_p90"]):
-            log_dict["eval/val/centroid_dist_p90"] = metrics["dist_p90"]
-        if not np.isnan(metrics["dist_p95"]):
-            log_dict["eval/val/centroid_dist_p95"] = metrics["dist_p95"]
-        if not np.isnan(metrics["dist_max"]):
-            log_dict["eval/val/centroid_dist_max"] = metrics["dist_max"]
-
-        # Detection metrics
-        log_dict["eval/val/centroid_precision"] = metrics["precision"]
-        log_dict["eval/val/centroid_recall"] = metrics["recall"]
-        log_dict["eval/val/centroid_f1"] = metrics["f1"]
-
-        # Counts
-        log_dict["eval/val/centroid_n_tp"] = metrics["n_true_positives"]
-        log_dict["eval/val/centroid_n_fp"] = metrics["n_false_positives"]
-        log_dict["eval/val/centroid_n_fn"] = metrics["n_false_negatives"]
-
-        wandb_logger.experiment.log(log_dict, commit=False)
-
-        # Update best metrics in summary
-        for key, value in log_dict.items():
-            if key == "epoch":
-                continue
-            summary_key = f"best/{key}"
-            current_best = wandb_logger.experiment.summary.get(summary_key)
-            # For distance metrics, lower is better; for others, higher is better
-            is_distance = "dist" in key
-            if current_best is None:
-                wandb_logger.experiment.summary[summary_key] = value
-            elif is_distance and value < current_best:
-                wandb_logger.experiment.summary[summary_key] = value
-            elif not is_distance and value > current_best:
-                wandb_logger.experiment.summary[summary_key] = value