fix reference tests

aminemindee · Oct 12, 2022 · 1b5c181 · 1b5c181
1 parent 5da5b09
commit 1b5c181
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 6 deletions.
diff --git a/references/detection/evaluate_pytorch.py b/references/detection/evaluate_pytorch.py
@@ -5,6 +5,10 @@
 
 import os
 
+import numpy as np
+
+from doctr.file_utils import CLASS_NAME
+
 os.environ["USE_TORCH"] = "1"
 
 import logging
@@ -44,9 +48,12 @@ def evaluate(model, val_loader, batch_transforms, val_metric, amp=False):
             out = model(images, targets, return_preds=True)
         # Compute metric
         loc_preds = out["preds"]
-        for boxes_gt, boxes_pred in zip(targets, loc_preds):
-            # Remove scores
-            val_metric.update(gts=boxes_gt, preds=boxes_pred[:, :-1])
+        for target, loc_pred in zip(targets, loc_preds):
+            if isinstance(target, np.ndarray):
+                target = {CLASS_NAME: target}
+            for boxes_gt, boxes_pred in zip(target.values(), loc_pred.values()):
+                # Remove scores
+                val_metric.update(gts=boxes_gt, preds=boxes_pred[:, :-1])
 
         val_loss += out["loss"].item()
         batch_cnt += 1

diff --git a/references/detection/evaluate_tensorflow.py b/references/detection/evaluate_tensorflow.py
@@ -5,6 +5,10 @@
 
 import os
 
+import numpy as np
+
+from doctr.file_utils import CLASS_NAME
+
 os.environ["USE_TF"] = "1"
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
 
@@ -39,9 +43,12 @@ def evaluate(model, val_loader, batch_transforms, val_metric):
         out = model(images, targets, training=False, return_preds=True)
         # Compute metric
         loc_preds = out["preds"]
-        for boxes_gt, boxes_pred in zip(targets, loc_preds):
-            # Remove scores
-            val_metric.update(gts=boxes_gt, preds=boxes_pred[:, :-1])
+        for target, loc_pred in zip(targets, loc_preds):
+            if isinstance(target, np.ndarray):
+                target = {CLASS_NAME: target}
+            for boxes_gt, boxes_pred in zip(target.values(), loc_pred.values()):
+                # Remove scores
+                val_metric.update(gts=boxes_gt, preds=boxes_pred[:, :-1])
 
         val_loss += out["loss"].numpy()
         batch_cnt += 1