Update coco_utils.py

coco_utils.py

@@ -27,24 +27,6 @@ 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():
@@ -150,6 +132,83 @@ class CocoEvaluator(object):
         for iou_type, coco_eval in self.coco_eval.items():
             print("IoU metric: {}".format(iou_type))
             coco_eval.summarize()
+    
+    def _post_process_stats(self, stats, coco_eval_object, iou_type='bbox'):
+        # bbox & segm:
+        # stats[0] = _summarize(1)
+        # stats[1] = _summarize(1, iouThr=.5, maxDets=self.params.maxDets[2])
+        # stats[2] = _summarize(1, iouThr=.75, maxDets=self.params.maxDets[2])
+        # stats[3] = _summarize(1, areaRng='small', maxDets=self.params.maxDets[2])
+        # stats[4] = _summarize(1, areaRng='medium', maxDets=self.params.maxDets[2])
+        # stats[5] = _summarize(1, areaRng='large', maxDets=self.params.maxDets[2])
+        # stats[6] = _summarize(0, maxDets=self.params.maxDets[0])
+        # stats[7] = _summarize(0, maxDets=self.params.maxDets[1])
+        # stats[8] = _summarize(0, maxDets=self.params.maxDets[2])
+        # stats[9] = _summarize(0, areaRng='small', maxDets=self.params.maxDets[2])
+        # stats[10] = _summarize(0, areaRng='medium', maxDets=self.params.maxDets[2])
+        # stats[11] = _summarize(0, areaRng='large', maxDets=self.params.maxDets[2])
+
+        # keypoints:
+        # stats[0] = _summarize(1, maxDets=20)
+        # stats[1] = _summarize(1, maxDets=20, iouThr=.5)
+        # stats[2] = _summarize(1, maxDets=20, iouThr=.75)
+        # stats[3] = _summarize(1, maxDets=20, areaRng='medium')
+        # stats[4] = _summarize(1, maxDets=20, areaRng='large')
+        # stats[5] = _summarize(0, maxDets=20)
+        # stats[6] = _summarize(0, maxDets=20, iouThr=.5)
+        # stats[7] = _summarize(0, maxDets=20, iouThr=.75)
+        # stats[8] = _summarize(0, maxDets=20, areaRng='medium')
+        # stats[9] = _summarize(0, maxDets=20, areaRng='large')
+        if iou_type not in ['bbox', 'segm', 'keypoints']:
+            raise ValueError(f"iou_type '{iou_type}' not supported")
+
+        current_max_dets = coco_eval_object.params.maxDets
+
+        index_to_title = {
+            "bbox": {
+                0: f"AP-IoU=0.50:0.95-area=all-maxDets={current_max_dets[2]}",
+                1: f"AP-IoU=0.50-area=all-maxDets={current_max_dets[2]}",
+                2: f"AP-IoU=0.75-area=all-maxDets={current_max_dets[2]}",
+                3: f"AP-IoU=0.50:0.95-area=small-maxDets={current_max_dets[2]}",
+                4: f"AP-IoU=0.50:0.95-area=medium-maxDets={current_max_dets[2]}",
+                5: f"AP-IoU=0.50:0.95-area=large-maxDets={current_max_dets[2]}",
+                6: f"AR-IoU=0.50:0.95-area=all-maxDets={current_max_dets[0]}",
+                7: f"AR-IoU=0.50:0.95-area=all-maxDets={current_max_dets[1]}",
+                8: f"AR-IoU=0.50:0.95-area=all-maxDets={current_max_dets[2]}",
+                9: f"AR-IoU=0.50:0.95-area=small-maxDets={current_max_dets[2]}",
+                10: f"AR-IoU=0.50:0.95-area=medium-maxDets={current_max_dets[2]}",
+                11: f"AR-IoU=0.50:0.95-area=large-maxDets={current_max_dets[2]}",
+            },
+            "keypoints":
+                {
+                    0: "AP-IoU=0.50:0.95-area=all-maxDets=20",
+                    1: "AP-IoU=0.50-area=all-maxDets=20",
+                    2: "AP-IoU=0.75-area=all-maxDets=20",
+                    3: "AP-IoU=0.50:0.95-area=medium-maxDets=20",
+                    4: "AP-IoU=0.50:0.95-area=large-maxDets=20",
+                    5: "AR-IoU=0.50:0.95-area=all-maxDets=20",
+                    6: "AR-IoU=0.50-area=all-maxDets=20",
+                    7: "AR-IoU=0.75-area=all-maxDets=20",
+                    8: "AR-IoU=0.50:0.95-area=medium-maxDets=20",
+                    9: "AR-IoU=0.50:0.95-area=large-maxDets=20",
+                },
+        }
+
+        output_dict = {}
+        for index, stat in enumerate(stats):
+            output_dict[index_to_title[iou_type][index]] = stat
+        
+        return output_dict
+
+
+    def get_results(self):
+        output_dict = {}
+
+        for iou_type, coco_eval in self.coco_eval.items():
+            if iou_type == 'segm':
+                iou_type = 'bbox'
+            output_dict[f"iou_{iou_type}"] = self._post_process_stats(coco_eval.stats, coco_eval, iou_type)
+        return output_dict
 
     def prepare(self, predictions, iou_type):
         if iou_type == "bbox":
@@ -168,9 +227,17 @@ class CocoEvaluator(object):
                 continue
 
             boxes = prediction["boxes"]
+            if not isinstance(boxes, torch.Tensor):
+                boxes = torch.as_tensor(boxes)
             boxes = convert_to_xywh(boxes).tolist()
-            scores = prediction["scores"].tolist()
-            labels = prediction["labels"].tolist()
+
+            scores = prediction["scores"]
+            if not isinstance(scores, list):
+                scores = scores.tolist()
+
+            labels = prediction["labels"]
+            if not isinstance(labels, list):
+                labels = prediction["labels"].tolist()
 
             coco_results.extend(
                 [