Update landmarkdiff/metrics_agg.py to v0.3.2

landmarkdiff/metrics_agg.py

@@ -41,12 +41,8 @@ class MetricsAggregator:
     """
 
     HIGHER_BETTER = {
-        "ssim": True,
-        "psnr": True,
-        "identity_sim": True,
-        "lpips": False,
-        "fid": False,
-        "nme": False,
+        "ssim": True, "psnr": True, "identity_sim": True,
+        "lpips": False, "fid": False, "nme": False,
     }
 
     def __init__(self) -> None:
@@ -61,15 +57,13 @@ class MetricsAggregator:
         **metadata: Any,
     ) -> None:
         """Add a single evaluation record."""
-        self.records.append(
-            MetricRecord(
-                experiment=experiment,
-                procedure=procedure,
-                metrics=metrics,
-                checkpoint_step=checkpoint_step,
-                metadata=metadata,
-            )
-        )
+        self.records.append(MetricRecord(
+            experiment=experiment,
+            procedure=procedure,
+            metrics=metrics,
+            checkpoint_step=checkpoint_step,
+            metadata=metadata,
+        ))
 
     def add_batch(
         self,
@@ -82,9 +76,7 @@ class MetricsAggregator:
         """
         for rec in records:
             proc = rec.get("procedure", "all")
-            metrics = {
-                k: v for k, v in rec.items() if k != "procedure" and isinstance(v, (int, float))
-            }
+            metrics = {k: v for k, v in rec.items() if k != "procedure" and isinstance(v, (int, float))}
             self.add(experiment, proc, metrics)
 
     @property
@@ -219,7 +211,10 @@ class MetricsAggregator:
             val = self.mean(exp, metric, procedure)
             if math.isnan(val):
                 continue
-            if (higher_better and val > best_val) or (not higher_better and val < best_val):
+            if higher_better and val > best_val:
+                best_val = val
+                best_exp = exp
+            elif not higher_better and val < best_val:
                 best_val = val
                 best_exp = exp
 
@@ -309,5 +304,6 @@ class MetricsAggregator:
                 procedure=rec["procedure"],
                 metrics=rec["metrics"],
                 checkpoint_step=rec.get("checkpoint_step"),
+                **rec.get("metadata", {}),
             )
         return agg