Consolidate metrics saving to use SLEAP 1.4 format and eliminate code duplication

- Updated run_evaluation() to save metrics using SLEAP 1.4 format (single "metrics" key)
- Refactored train.py to use run_evaluation() instead of duplicating evaluation code
- Removed unused imports (numpy, sleap_io) from train.py
- Updated test_evaluation.py to load metrics in SLEAP 1.4 format
- Ensured load_metrics() function is compatible with the save format

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <[email protected]>

Author: gitttt-1234

sleap_nn/evaluation.py

@@ -745,22 +745,8 @@ def run_evaluation(
         save_path = Path(save_metrics)
         save_path.parent.mkdir(parents=True, exist_ok=True)
 
-        # Convert metrics to numpy arrays for saving
-        np.savez(
-            save_path,
-            mOKS=metrics["mOKS"]["mOKS"],
-            mAP=metrics["voc_metrics"]["oks_voc.mAP"],
-            mAR=metrics["voc_metrics"]["oks_voc.mAR"],
-            avg_distance=metrics["distance_metrics"]["avg"],
-            mPCK=metrics["pck_metrics"]["mPCK"],
-            visibility_precision=metrics["visibility_metrics"]["precision"],
-            visibility_recall=metrics["visibility_metrics"]["recall"],
-            # Save full metrics dict as well
-            voc_metrics=metrics["voc_metrics"],
-            distance_metrics=metrics["distance_metrics"],
-            pck_metrics=metrics["pck_metrics"],
-            visibility_metrics=metrics["visibility_metrics"],
-        )
+        # Save metrics in SLEAP 1.4 format (single "metrics" key)
+        np.savez_compressed(save_path, **{"metrics": metrics})
         logger.info(f"Metrics saved successfully to {save_path}")
 
     return metrics

sleap_nn/train.py

@@ -2,22 +2,19 @@
 
 from loguru import logger
 from pathlib import Path
-import numpy as np
 from datetime import datetime
 from time import time
 from omegaconf import DictConfig, OmegaConf
 from typing import Any, Dict, Optional, List, Tuple, Union
-import sleap_io as sio
 from sleap_nn.config.training_job_config import TrainingJobConfig
 from sleap_nn.training.model_trainer import ModelTrainer
 from sleap_nn.predict import run_inference as predict
-from sleap_nn.evaluation import Evaluator
+from sleap_nn.evaluation import run_evaluation
 from sleap_nn.config.get_config import (
     get_trainer_config,
     get_model_config,
     get_data_config,
 )
-from typing import Any, Dict, Optional, List, Tuple, Union
 
 
 def run_training(config: DictConfig):
@@ -64,7 +61,16 @@ def run_training(config: DictConfig):
                 data_paths["test"] = config.data_config.test_file_path
 
             for d_name, path in data_paths.items():
-                labels = sio.load_slp(path)
+                pred_path = (
+                    Path(trainer.config.trainer_config.ckpt_dir)
+                    / trainer.config.trainer_config.run_name
+                    / f"pred_{d_name}.slp"
+                )
+                metrics_path = (
+                    Path(trainer.config.trainer_config.ckpt_dir)
+                    / trainer.config.trainer_config.run_name
+                    / f"{d_name}_pred_metrics.npz"
+                )
 
                 pred_labels = predict(
                     data_path=path,
@@ -75,9 +81,7 @@ def run_training(config: DictConfig):
                     peak_threshold=0.2,
                     make_labels=True,
                     device=trainer.trainer.strategy.root_device,
-                    output_path=Path(trainer.config.trainer_config.ckpt_dir)
-                    / trainer.config.trainer_config.run_name
-                    / f"pred_{d_name}.slp",
+                    output_path=pred_path,
                     ensure_rgb=config.data_config.preprocessing.ensure_rgb,
                     ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
                 )
@@ -88,27 +92,18 @@ def run_training(config: DictConfig):
                     )
                     continue  # skip if there are no labeled frames
 
-                evaluator = Evaluator(
-                    ground_truth_instances=labels, predicted_instances=pred_labels
-                )
-                metrics = evaluator.evaluate()
-                np.savez_compressed(
-                    (
-                        Path(trainer.config.trainer_config.ckpt_dir)
-                        / trainer.config.trainer_config.run_name
-                        / f"{d_name}_pred_metrics.npz"
-                    ).as_posix(),
-                    **{"metrics": metrics},
+                # Run evaluation and save metrics
+                metrics = run_evaluation(
+                    ground_truth_path=path,
+                    predicted_path=pred_path.as_posix(),
+                    save_metrics=metrics_path.as_posix(),
                 )
 
                 logger.info(f"---------Evaluation on `{d_name}` dataset---------")
                 logger.info(f"OKS mAP: {metrics['voc_metrics']['oks_voc.mAP']}")
                 logger.info(f"Average distance: {metrics['distance_metrics']['avg']}")
                 logger.info(f"p90 dist: {metrics['distance_metrics']['p90']}")
                 logger.info(f"p50 dist: {metrics['distance_metrics']['p50']}")
-                logger.info(
-                    f"metrics saved to {Path(trainer.config.trainer_config.ckpt_dir) / trainer.config.trainer_config.run_name / (d_name + '_pred_metrics.npz')}"
-                )
 
 
 def train(

tests/test_evaluation.py

@@ -526,13 +526,16 @@ def test_evaluator_main(
     result = subprocess.run(cmd, check=True, capture_output=True, text=True)
     assert Path(f"{tmp_path}/metrics_test.npz").exists()
 
-    metrics = np.load(f"{tmp_path}/metrics_test.npz", allow_pickle=True)
+    # Load metrics in SLEAP 1.4 format (single "metrics" key)
+    metrics_npz = np.load(f"{tmp_path}/metrics_test.npz", allow_pickle=True)
+    assert "metrics" in metrics_npz
+    metrics = metrics_npz["metrics"].item()
     assert "voc_metrics" in metrics
     assert "mOKS" in metrics
     assert "distance_metrics" in metrics
     assert "pck_metrics" in metrics
     assert "visibility_metrics" in metrics
-    voc_metrics = metrics["voc_metrics"].item()
+    voc_metrics = metrics["voc_metrics"]
     assert "pck_voc.mAP" in voc_metrics
     assert "pck_voc.mAR" in voc_metrics
     assert "oks_voc.mAP" in voc_metrics