Modified the MetricWrappers arguments being passed on

This will hopefully simplify the arguments to each metric slightly.

Author: c-salomonsen

utils/load_metric.py

@@ -45,11 +45,13 @@ class MetricWrapper(nn.Module):
     {'entropy': [], 'f1': [], 'precision': []}
     """
 
-    def __init__(self, *metrics, num_classes, macro_averaging=False):
+    def __init__(self, *metrics, num_classes, macro_averaging=False, **kwargs):
         super().__init__()
         self.metrics = {}
-        self.num_classes = num_classes
-        self.macro_averaging = macro_averaging
+        self.params = {
+            "num_classes": num_classes,
+            "macro_averaging": macro_averaging,
+        } | kwargs
 
         for metric in metrics:
             self.metrics[metric] = self._get_metric(metric)
@@ -73,23 +75,15 @@ def _get_metric(self, key):
 
         match key.lower():
             case "entropy":
-                return EntropyPrediction(num_classes=self.num_classes)
+                return EntropyPrediction(**self.params)
             case "f1":
-                return F1Score(
-                    num_classes=self.num_classes, macro_averaging=self.macro_averaging
-                )
+                return F1Score(**self.params)
             case "recall":
-                return Recall(
-                    num_classes=self.num_classes, macro_averaging=self.macro_averaging
-                )
+                return Recall(**self.params)
             case "precision":
-                return Precision(
-                    num_classes=self.num_classes, macro_averaging=self.macro_averaging
-                )
+                return Precision(**self.params)
             case "accuracy":
-                return Accuracy(
-                    num_classes=self.num_classes, macro_averaging=self.macro_averaging
-                )
+                return Accuracy(**self.params)
             case _:
                 raise ValueError(f"Metric {key} not supported")