Update cocoevaluate.py

cocoevaluate.py

@@ -17,7 +17,7 @@ import evaluate
 import datasets
 import pyarrow as pa
 
-from .coco_utils import CocoEvaluator, get_coco_api_from_dataset, CocoDetection
+from .coco_utils import CocoEvaluator, get_coco_api_from_dataset
 
 # TODO: Add BibTeX citation
 _CITATION = """\
@@ -72,9 +72,10 @@ def summarize_if_long_list(obj):
 @evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
 class COCOEvaluate(evaluate.Metric):
     """TODO: Short description of my evaluation module."""
-    def __init__(self, coco_path, feature_extractor, annotation_path, iou_types=['bbox'], **kwargs):
+    def __init__(self, coco_dataset, iou_types=['bbox'], **kwargs):
         super().__init__(**kwargs)
-        self.coco_dataset = CocoDetection(coco_path, feature_extractor, annotation_path)
+        # self.coco_dataset = CocoDetection(coco_path, feature_extractor, annotation_path)
+        self.coco_dataset = coco_dataset
         base_ds = get_coco_api_from_dataset(self.coco_dataset)
         self.coco_evaluator = CocoEvaluator(base_ds, iou_types)
 
@@ -92,22 +93,23 @@ class COCOEvaluate(evaluate.Metric):
                 'predictions': [
                         datasets.Features(
                             {
-                                'scores': datasets.Array2D(shape=(100,), dtype='float32'),
-                                'labels': datasets.Array2D(shape=(100,), dtype='int64'),
-                                'boxes': datasets.Array2D(shape=(100, 4), dtype='float32'),
+                                'scores': datasets.Sequence(datasets.Value("float")),
+                                'labels': datasets.Sequence(datasets.Value("int64")),
+                                'boxes': datasets.Sequence(datasets.Sequence(datasets.Value("float"))),
                             })
                     ]
                     ,
                 'references': [
                         datasets.Features(
                         {
-                            'size': datasets.Value('int64'),
-                            'image_id': datasets.Value('int64'),
-                            'boxes': datasets.Array2D(shape=(20, 4), dtype='float32'),
-                            'class_labels': datasets.Value('int64'),
-                            'iscrowd': datasets.Value('int64'),
-                            'orig_size': datasets.Value('int64'),
-                            'area': datasets.Value('float32'),
+                            'size': datasets.Sequence(datasets.Value("int64")),
+                            'image_id': datasets.Sequence(datasets.Value("int64")),
+                            'boxes': datasets.Sequence(datasets.Sequence(datasets.Value("float"))),
+                            'class_labels': datasets.Sequence(datasets.Value("int64")),
+                            'iscrowd': datasets.Sequence(datasets.Value("int64")),
+                            'orig_size': datasets.Sequence(datasets.Value("int64")),
+                            'area': datasets.Sequence(datasets.Value("int64")),
+                            
                         }
                     )
                 ],
@@ -125,16 +127,36 @@ class COCOEvaluate(evaluate.Metric):
         # TODO: Download external resources if needed
         pass
 
-    def _compute(self, predictions, references,):
+    def _preprocess(self, predictions):
+        """Optional: preprocess the predictions and references before computing the scores"""
+        processed_predictions = []
+        for pred in predictions:
+            processed_pred = {}
+            for key in pred.keys():
+                processed_pred[key] = pred[key].detach().cpu().tolist()
+            processed_predictions.append(processed_pred)
+        return processed_predictions
+    
+    def add(self, *, prediction=None, reference=None, **kwargs):
+        """Preprocesses the predictions and references and calls the function of the parent class."""
+        if prediction is not None:
+            prediction = self._preprocess(prediction)
+        if reference is not None:
+            reference = self._preprocess(reference)
+        super().add(prediction=prediction, references=reference, **kwargs)
+
+    def _compute(self, predictions, references):
         """Returns the scores"""
         for pred, ref in zip(predictions, references):
-            res = {target['image_id'].item(): output for target, output in zip(ref, pred)}
+            res = {}
+            for target, output in zip(ref, pred):
+                res[target["image_id"][0]] = output
             self.coco_evaluator.update(res)
         self.coco_evaluator.synchronize_between_processes()
         self.coco_evaluator.accumulate()
+
         self.coco_evaluator.summarize()
 
-        accuracy = None
-        return {
-            "accuracy": accuracy,
-        }
+        stats = self.coco_evaluator.get_results()
+
+        return stats