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

sleap_nn/training/callbacks.py

@@ -2,10 +2,15 @@
 
 import zmq
 import jsonpickle
-from typing import Callable, Optional
+from typing import Callable, Optional, Union
 from lightning.pytorch.callbacks import Callback
+from lightning.pytorch.callbacks.progress import TQDMProgressBar
 from loguru import logger
 import matplotlib
+
+matplotlib.use(
+    "Agg"
+)  # Use non-interactive backend to avoid tkinter issues on Windows CI
 import matplotlib.pyplot as plt
 from PIL import Image
 from pathlib import Path
@@ -14,6 +19,32 @@ import csv
 from sleap_nn import RANK
 
 
+class SleapProgressBar(TQDMProgressBar):
+    """Custom progress bar with better formatting for small metric values.
+
+    The default TQDMProgressBar truncates small floats like 1e-5 to "0.000".
+    This subclass formats metrics using scientific notation when appropriate.
+    """
+
+    def get_metrics(
+        self, trainer, pl_module
+    ) -> dict[str, Union[int, str, float, dict[str, float]]]:
+        """Override to format metrics with scientific notation for small values."""
+        items = super().get_metrics(trainer, pl_module)
+        formatted = {}
+        for k, v in items.items():
+            if isinstance(v, float):
+                # Use scientific notation for very small values
+                if v != 0 and abs(v) < 0.001:
+                    formatted[k] = f"{v:.2e}"
+                else:
+                    # Use 4 decimal places for normal values
+                    formatted[k] = f"{v:.4f}"
+            else:
+                formatted[k] = v
+        return formatted
+
+
 class CSVLoggerCallback(Callback):
     """Callback for logging metrics to csv.
 
@@ -53,6 +84,21 @@ class CSVLoggerCallback(Callback):
             for key in self.keys:
                 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)
+                    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)
+                        for metric_key in metrics.keys():
+                            if metric_key.startswith("lr-"):
+                                value = metrics[metric_key]
+                                break
+                    log_data[key] = value.item() if value is not None else None
                 else:
                     value = metrics.get(key, None)
                     log_data[key] = value.item() if value is not None else None
@@ -66,7 +112,11 @@ class CSVLoggerCallback(Callback):
 
 
 class WandBPredImageLogger(Callback):
-    """Callback for writing image predictions to wandb.
+    """Callback for writing image predictions to wandb as a Table.
+
+    .. deprecated::
+        This callback logs images to a wandb.Table which doesn't support
+        step sliders. Use WandBVizCallback instead for better UX.
 
     Attributes:
         viz_folder: Path to viz directory.
@@ -141,12 +191,576 @@ class WandBPredImageLogger(Callback):
                     ]
                 ]
             table = wandb.Table(columns=column_names, data=data)
-            wandb.log({f"{self.wandb_run_name}": table})
+            # Use commit=False to accumulate with other metrics in this step
+            wandb.log({f"{self.wandb_run_name}": table}, commit=False)
 
         # Sync all processes after wandb logging
         trainer.strategy.barrier()
 
 
+class WandBVizCallback(Callback):
+    """Callback for logging visualization images directly to wandb with slider support.
+
+    This callback logs images using wandb.log() which enables step slider navigation
+    in the wandb UI. Multiple visualization modes can be enabled simultaneously:
+    - viz_enabled: Pre-render with matplotlib (same as disk viz)
+    - viz_boxes: Interactive keypoint boxes with filtering
+    - viz_masks: Confidence map overlay with per-node toggling
+
+    Attributes:
+        train_viz_fn: Function that returns VisualizationData for training sample.
+        val_viz_fn: Function that returns VisualizationData for validation sample.
+        viz_enabled: Whether to log pre-rendered matplotlib images.
+        viz_boxes: Whether to log interactive keypoint boxes.
+        viz_masks: Whether to log confidence map overlay masks.
+        box_size: Size of keypoint boxes in pixels (for viz_boxes).
+        confmap_threshold: Threshold for confmap masks (for viz_masks).
+        log_table: Whether to also log to a wandb.Table (backwards compat).
+    """
+
+    def __init__(
+        self,
+        train_viz_fn: Callable,
+        val_viz_fn: Callable,
+        viz_enabled: bool = True,
+        viz_boxes: bool = False,
+        viz_masks: bool = False,
+        box_size: float = 5.0,
+        confmap_threshold: float = 0.1,
+        log_table: bool = False,
+    ):
+        """Initialize the callback.
+
+        Args:
+            train_viz_fn: Callable that returns VisualizationData for a training sample.
+            val_viz_fn: Callable that returns VisualizationData for a validation sample.
+            viz_enabled: If True, log pre-rendered matplotlib images.
+            viz_boxes: If True, log interactive keypoint boxes.
+            viz_masks: If True, log confidence map overlay masks.
+            box_size: Size of keypoint boxes in pixels (for viz_boxes).
+            confmap_threshold: Threshold for confmap mask generation (for viz_masks).
+            log_table: If True, also log images to a wandb.Table (for backwards compat).
+        """
+        super().__init__()
+        self.train_viz_fn = train_viz_fn
+        self.val_viz_fn = val_viz_fn
+        self.viz_enabled = viz_enabled
+        self.viz_boxes = viz_boxes
+        self.viz_masks = viz_masks
+        self.log_table = log_table
+
+        # Import here to avoid circular imports
+        from sleap_nn.training.utils import WandBRenderer
+
+        self.box_size = box_size
+        self.confmap_threshold = confmap_threshold
+
+        # Create renderers for each enabled mode
+        self.renderers = {}
+        if viz_enabled:
+            self.renderers["direct"] = WandBRenderer(
+                mode="direct", box_size=box_size, confmap_threshold=confmap_threshold
+            )
+        if viz_boxes:
+            self.renderers["boxes"] = WandBRenderer(
+                mode="boxes", box_size=box_size, confmap_threshold=confmap_threshold
+            )
+        if viz_masks:
+            self.renderers["masks"] = WandBRenderer(
+                mode="masks", box_size=box_size, confmap_threshold=confmap_threshold
+            )
+
+    def _get_wandb_logger(self, trainer):
+        """Get the WandbLogger from trainer's loggers."""
+        from lightning.pytorch.loggers import WandbLogger
+
+        for logger in trainer.loggers:
+            if isinstance(logger, WandbLogger):
+                return logger
+        return None
+
+    def on_train_epoch_end(self, trainer, pl_module):
+        """Log visualization images at end of each epoch."""
+        if trainer.is_global_zero:
+            epoch = trainer.current_epoch
+
+            # 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()
+
+                # 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
+        trainer.strategy.barrier()
+
+
+class WandBVizCallbackWithPAFs(WandBVizCallback):
+    """Extended WandBVizCallback that also logs PAF visualizations for bottom-up models."""
+
+    def __init__(
+        self,
+        train_viz_fn: Callable,
+        val_viz_fn: Callable,
+        train_pafs_viz_fn: Callable,
+        val_pafs_viz_fn: Callable,
+        viz_enabled: bool = True,
+        viz_boxes: bool = False,
+        viz_masks: bool = False,
+        box_size: float = 5.0,
+        confmap_threshold: float = 0.1,
+        log_table: bool = False,
+    ):
+        """Initialize the callback.
+
+        Args:
+            train_viz_fn: Callable returning VisualizationData for training sample.
+            val_viz_fn: Callable returning VisualizationData for validation sample.
+            train_pafs_viz_fn: Callable returning VisualizationData with PAFs for training.
+            val_pafs_viz_fn: Callable returning VisualizationData with PAFs for validation.
+            viz_enabled: If True, log pre-rendered matplotlib images.
+            viz_boxes: If True, log interactive keypoint boxes.
+            viz_masks: If True, log confidence map overlay masks.
+            box_size: Size of keypoint boxes in pixels.
+            confmap_threshold: Threshold for confmap mask generation.
+            log_table: If True, also log images to a wandb.Table.
+        """
+        super().__init__(
+            train_viz_fn=train_viz_fn,
+            val_viz_fn=val_viz_fn,
+            viz_enabled=viz_enabled,
+            viz_boxes=viz_boxes,
+            viz_masks=viz_masks,
+            box_size=box_size,
+            confmap_threshold=confmap_threshold,
+            log_table=log_table,
+        )
+        self.train_pafs_viz_fn = train_pafs_viz_fn
+        self.val_pafs_viz_fn = val_pafs_viz_fn
+
+        # Import here to avoid circular imports
+        from sleap_nn.training.utils import MatplotlibRenderer
+
+        self._mpl_renderer = MatplotlibRenderer()
+
+    def on_train_epoch_end(self, trainer, pl_module):
+        """Log visualization images including PAFs at end of each epoch."""
+        if trainer.is_global_zero:
+            epoch = trainer.current_epoch
+
+            # 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)
+            buf = BytesIO()
+            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}"
+            )
+
+        # 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)
+            buf = BytesIO()
+            class_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
+            )
+
+    def on_train_epoch_end(self, trainer, pl_module):
+        """Generate and output all visualizations at epoch end.
+
+        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
+        trainer.strategy.barrier()
+
+
 class MatplotlibSaver(Callback):
     """Callback for saving images rendered with matplotlib during training.
 
@@ -194,7 +808,7 @@ class MatplotlibSaver(Callback):
             ).as_posix()
 
             # Save rendered figure.
-            figure.savefig(figure_path, format="png", pad_inches=0)
+            figure.savefig(figure_path, format="png")
             plt.close(figure)
 
         # Sync all processes after file I/O
@@ -303,7 +917,11 @@ class ProgressReporterZMQ(Callback):
     def on_train_start(self, trainer, pl_module):
         """Called at the beginning of training process."""
         if trainer.is_global_zero:
-            self.send("train_begin")
+            # Include WandB URL if available
+            wandb_url = None
+            if wandb.run is not None:
+                wandb_url = wandb.run.url
+            self.send("train_begin", wandb_url=wandb_url)
         trainer.strategy.barrier()
 
     def on_train_end(self, trainer, pl_module):
@@ -350,3 +968,638 @@ 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