Metrics - Fixed Missing Reset

luxonis_train/attached_modules/metrics/common.py

@@ -1,6 +1,7 @@
 import logging
 
 import torchmetrics
+from torch import Tensor
 
 from .base_metric import BaseMetric
 
@@ -47,14 +48,17 @@ def __init__(self, **kwargs):
 
         self.metric = self.Metric(**kwargs)
 
-    def update(self, preds, target, *args, **kwargs):
+    def update(self, preds, target, *args, **kwargs) -> None:
         if self.task in ["multiclass"]:
             target = target.argmax(dim=1)
         self.metric.update(preds, target, *args, **kwargs)
 
-    def compute(self):
+    def compute(self) -> Tensor:
         return self.metric.compute()
 
+    def reset(self) -> None:
+        self.metric.reset()
+
 
 class Accuracy(TorchMetricWrapper):
     Metric = torchmetrics.Accuracy

luxonis_train/attached_modules/metrics/mean_average_precision.py

@@ -12,7 +12,7 @@
 from .base_metric import BaseMetric
 
 
-class MeanAveragePrecision(BaseMetric, detection.MeanAveragePrecision):
+class MeanAveragePrecision(BaseMetric):
     """Compute the Mean-Average-Precision (mAP) and Mean-Average-Recall (mAR) for object
     detection predictions.
 
@@ -62,6 +62,9 @@ def prepare(
 
         return output_list, label_list
 
+    def reset(self) -> None:
+        self.metric.reset()
+
     def compute(self) -> tuple[Tensor, dict[str, Tensor]]:
         metric_dict = self.metric.compute()