add files

coco_utils.py

@@ -27,6 +27,24 @@ def is_dist_avail_and_initialized():
         return False
     return True
 
+class CocoDetection(torchvision.datasets.CocoDetection):
+    def __init__(self, img_folder, feature_extractor, ann_file):
+        super(CocoDetection, self).__init__(img_folder, ann_file)
+        self.feature_extractor = feature_extractor
+
+    def __getitem__(self, idx):
+        # read in PIL image and target in COCO format
+        img, target = super(CocoDetection, self).__getitem__(idx)
+        
+        # preprocess image and target (converting target to DETR format, resizing + normalization of both image and target)
+        image_id = self.ids[idx]
+        target = {'image_id': image_id, 'annotations': target}
+        encoding = self.feature_extractor(images=img, annotations=target, return_tensors="pt")
+        pixel_values = encoding["pixel_values"].squeeze() # remove batch dimension
+        target = encoding["labels"][0] # remove batch dimension
+
+        return pixel_values, target
+
 
 def get_world_size():
     if not is_dist_avail_and_initialized():

cocoevaluate.py

@@ -17,7 +17,7 @@ import evaluate
 import datasets
 import pyarrow as pa
 
-from .coco_utils import CocoEvaluator, get_coco_api_from_dataset
+from .coco_utils import CocoEvaluator, get_coco_api_from_dataset, CocoDetection
 
 # 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_dataset, iou_types=['bbox'], **kwargs):
+    def __init__(self, coco_path, feature_extractor, annotation_path, iou_types=['bbox'], **kwargs):
         super().__init__(**kwargs)
-        base_ds = get_coco_api_from_dataset(coco_dataset)
+        self.coco_dataset = CocoDetection(coco_path, feature_extractor, annotation_path)
+        base_ds = get_coco_api_from_dataset(self.coco_dataset)
         self.coco_evaluator = CocoEvaluator(base_ds, iou_types)