From 3017ae5fd78fd0602e80fa1e0f0ad74ab7ff8e1b Mon Sep 17 00:00:00 2001
From: Andrew Beveridge <andrew@beveridge.uk>
Date: Mon, 16 Dec 2024 20:21:40 -0500
Subject: [PATCH] Added loading of existing raw scores

---
 tests/model-metrics/test-all-models.py | 75 ++++++++++++++------------
 1 file changed, 42 insertions(+), 33 deletions(-)

diff --git a/tests/model-metrics/test-all-models.py b/tests/model-metrics/test-all-models.py
index 922a25f..cc7cb7b 100644
--- a/tests/model-metrics/test-all-models.py
+++ b/tests/model-metrics/test-all-models.py
@@ -28,39 +28,48 @@ def evaluate_track(track_name, track_path, test_model, mus_db):
     output_dir = os.path.join(RESULTS_PATH, test_model, track_name)
     os.makedirs(output_dir, exist_ok=True)
 
-    # Check if separated files already exist
-    vocals_path = os.path.join(output_dir, "vocals.wav")
-    instrumental_path = os.path.join(output_dir, "instrumental.wav")
-
-    if not (os.path.exists(vocals_path) and os.path.exists(instrumental_path)):
-        logger.info("Performing separation...")
-        separator = Separator(output_dir=output_dir)
-        separator.load_model(model_filename=test_model)
-        separator.separate(os.path.join(track_path, "mixture.wav"), custom_output_names={"Vocals": "vocals", "Instrumental": "instrumental"})
-
-    # Get track from MUSDB
-    track = next((t for t in mus_db if t.name == track_name), None)
-    if track is None:
-        raise ValueError(f"Track {track_name} not found in MUSDB18")
-
-    # Load estimated stems
-    estimates = {}
-    for stem_name in ["vocals", "accompaniment"]:
-        stem_path = vocals_path if stem_name == "vocals" else instrumental_path
-        audio, _ = sf.read(stem_path)
-        if len(audio.shape) == 1:
-            audio = np.expand_dims(audio, axis=1)
-        estimates[stem_name] = audio
-
-    # Evaluate using museval
-    scores = museval.eval_mus_track(track, estimates, output_dir=output_dir, mode="v4")
-
-    # Move and rename the results file
-    test_results = os.path.join(output_dir, "test", f"{track_name}.json")
-    new_results = os.path.join(output_dir, "museval-results.json")
-    if os.path.exists(test_results):
-        os.rename(test_results, new_results)
-        os.rmdir(os.path.join(output_dir, "test"))
+    # Check if evaluation results already exist
+    results_file = os.path.join(output_dir, "museval-results.json")
+    if os.path.exists(results_file):
+        logger.info("Found existing evaluation results, loading from file...")
+        with open(results_file) as f:
+            json_data = json.load(f)
+        scores = museval.TrackStore(track_name)
+        scores.scores = json_data
+    else:
+        # Check if separated files already exist
+        vocals_path = os.path.join(output_dir, "vocals.wav")
+        instrumental_path = os.path.join(output_dir, "instrumental.wav")
+
+        if not (os.path.exists(vocals_path) and os.path.exists(instrumental_path)):
+            logger.info("Performing separation...")
+            separator = Separator(output_dir=output_dir)
+            separator.load_model(model_filename=test_model)
+            separator.separate(os.path.join(track_path, "mixture.wav"), custom_output_names={"Vocals": "vocals", "Instrumental": "instrumental"})
+
+        # Get track from MUSDB
+        track = next((t for t in mus_db if t.name == track_name), None)
+        if track is None:
+            raise ValueError(f"Track {track_name} not found in MUSDB18")
+
+        # Load estimated stems
+        estimates = {}
+        for stem_name in ["vocals", "accompaniment"]:
+            stem_path = vocals_path if stem_name == "vocals" else instrumental_path
+            audio, _ = sf.read(stem_path)
+            if len(audio.shape) == 1:
+                audio = np.expand_dims(audio, axis=1)
+            estimates[stem_name] = audio
+
+        # Evaluate using museval
+        logger.info("Performing evaluation...")
+        scores = museval.eval_mus_track(track, estimates, output_dir=output_dir, mode="v4")
+
+        # Move and rename the results file
+        test_results = os.path.join(output_dir, "test", f"{track_name}.json")
+        if os.path.exists(test_results):
+            os.rename(test_results, results_file)
+            os.rmdir(os.path.join(output_dir, "test"))
 
     # Calculate aggregate scores
     results_store = museval.EvalStore()