update to faster and simple code (#607)

Co-authored-by: 吕文玉 <[email protected]>

Author: MiXaiLL76

rtdetrv2_pytorch/requirements.txt

@@ -1,4 +1,6 @@
-faster-coco-eval>=1.6.5
+torch>=2.0.1
+torchvision>=0.15.2
+faster-coco-eval>=1.6.6
 PyYAML
 tensorboard
 scipy

rtdetrv2_pytorch/src/data/dataset/coco_dataset.py

@@ -6,10 +6,8 @@
 """
 
 import torch
-import torch.utils.data
-
+from faster_coco_eval.utils.pytorch import FasterCocoDetection
 import torchvision
-torchvision.disable_beta_transforms_warning()
 
 from PIL import Image 
 from faster_coco_eval.core import mask as coco_mask
@@ -20,14 +18,15 @@
 
 __all__ = ['CocoDetection']
 
+torchvision.disable_beta_transforms_warning()
 
 @register()
-class CocoDetection(torchvision.datasets.CocoDetection, DetDataset):
+class CocoDetection(FasterCocoDetection, DetDataset):
     __inject__ = ['transforms', ]
     __share__ = ['remap_mscoco_category']
 
     def __init__(self, img_folder, ann_file, transforms, return_masks=False, remap_mscoco_category=False):
-        super(CocoDetection, self).__init__(img_folder, ann_file)
+        super(FasterCocoDetection, self).__init__(img_folder, ann_file)
         self._transforms = transforms
         self.prepare = ConvertCocoPolysToMask(return_masks)
         self.img_folder = img_folder
@@ -42,7 +41,7 @@ def __getitem__(self, idx):
         return img, target
 
     def load_item(self, idx):
-        image, target = super(CocoDetection, self).__getitem__(idx)
+        image, target = super(FasterCocoDetection, self).__getitem__(idx)
         image_id = self.ids[idx]
         target = {'image_id': image_id, 'annotations': target}
 

rtdetrv2_pytorch/src/data/dataset/coco_eval.py

@@ -8,195 +8,9 @@
 # MiXaiLL76 replacing pycocotools with faster-coco-eval for better performance and support.
 """
 
-import os
-import contextlib
-import copy
-import numpy as np
-import torch
-
-from faster_coco_eval import COCO, COCOeval_faster
-import faster_coco_eval.core.mask as mask_util
 from ...core import register
-from ...misc import dist_utils
-__all__ = ['CocoEvaluator',]
-
+from faster_coco_eval.utils.pytorch import FasterCocoEvaluator
 
 @register()
-class CocoEvaluator(object):
-    def __init__(self, coco_gt, iou_types):
-        assert isinstance(iou_types, (list, tuple))
-        coco_gt = copy.deepcopy(coco_gt)
-        self.coco_gt : COCO = coco_gt
-        self.iou_types = iou_types
-
-        self.coco_eval = {}
-        for iou_type in iou_types:
-            self.coco_eval[iou_type] = COCOeval_faster(coco_gt, iouType=iou_type, print_function=print, separate_eval=True)
-
-        self.img_ids = []
-        self.eval_imgs = {k: [] for k in iou_types}
-
-    def cleanup(self):
-        self.coco_eval = {}
-        for iou_type in self.iou_types:
-            self.coco_eval[iou_type] = COCOeval_faster(self.coco_gt, iouType=iou_type, print_function=print, separate_eval=True)
-        self.img_ids = []
-        self.eval_imgs = {k: [] for k in self.iou_types}
-
-
-    def update(self, predictions):
-        img_ids = list(np.unique(list(predictions.keys())))
-        self.img_ids.extend(img_ids)
-
-        for iou_type in self.iou_types:
-            results = self.prepare(predictions, iou_type)
-            coco_eval = self.coco_eval[iou_type]
-
-            with open(os.devnull, 'w') as devnull:
-                with contextlib.redirect_stdout(devnull):
-                    coco_dt = self.coco_gt.loadRes(results) if results else COCO()
-                    coco_eval.cocoDt = coco_dt
-                    coco_eval.params.imgIds = list(img_ids)
-                    coco_eval.evaluate()
-
-            self.eval_imgs[iou_type].append(np.array(coco_eval._evalImgs_cpp).reshape(len(coco_eval.params.catIds), len(coco_eval.params.areaRng), len(coco_eval.params.imgIds)))
-
-    def synchronize_between_processes(self):
-        for iou_type in self.iou_types:
-            img_ids, eval_imgs = merge(self.img_ids, self.eval_imgs[iou_type])
-
-            coco_eval = self.coco_eval[iou_type]
-            coco_eval.params.imgIds = img_ids
-            coco_eval._paramsEval = copy.deepcopy(coco_eval.params)
-            coco_eval._evalImgs_cpp = eval_imgs
-
-    def accumulate(self):
-        for coco_eval in self.coco_eval.values():
-            coco_eval.accumulate()
-
-    def summarize(self):
-        for iou_type, coco_eval in self.coco_eval.items():
-            print("IoU metric: {}".format(iou_type))
-            coco_eval.summarize()
-
-    def prepare(self, predictions, iou_type):
-        if iou_type == "bbox":
-            return self.prepare_for_coco_detection(predictions)
-        elif iou_type == "segm":
-            return self.prepare_for_coco_segmentation(predictions)
-        elif iou_type == "keypoints":
-            return self.prepare_for_coco_keypoint(predictions)
-        else:
-            raise ValueError("Unknown iou type {}".format(iou_type))
-
-    def prepare_for_coco_detection(self, predictions):
-        coco_results = []
-        for original_id, prediction in predictions.items():
-            if len(prediction) == 0:
-                continue
-
-            boxes = prediction["boxes"]
-            boxes = convert_to_xywh(boxes).tolist()
-            scores = prediction["scores"].tolist()
-            labels = prediction["labels"].tolist()
-
-            coco_results.extend(
-                [
-                    {
-                        "image_id": original_id,
-                        "category_id": labels[k],
-                        "bbox": box,
-                        "score": scores[k],
-                    }
-                    for k, box in enumerate(boxes)
-                ]
-            )
-        return coco_results
-
-    def prepare_for_coco_segmentation(self, predictions):
-        coco_results = []
-        for original_id, prediction in predictions.items():
-            if len(prediction) == 0:
-                continue
-
-            scores = prediction["scores"]
-            labels = prediction["labels"]
-            masks = prediction["masks"]
-
-            masks = masks > 0.5
-
-            scores = prediction["scores"].tolist()
-            labels = prediction["labels"].tolist()
-
-            rles = [
-                mask_util.encode(np.array(mask[0, :, :, np.newaxis], dtype=np.uint8, order="F"))[0]
-                for mask in masks
-            ]
-            for rle in rles:
-                rle["counts"] = rle["counts"].decode("utf-8")
-
-            coco_results.extend(
-                [
-                    {
-                        "image_id": original_id,
-                        "category_id": labels[k],
-                        "segmentation": rle,
-                        "score": scores[k],
-                    }
-                    for k, rle in enumerate(rles)
-                ]
-            )
-        return coco_results
-
-    def prepare_for_coco_keypoint(self, predictions):
-        coco_results = []
-        for original_id, prediction in predictions.items():
-            if len(prediction) == 0:
-                continue
-
-            boxes = prediction["boxes"]
-            boxes = convert_to_xywh(boxes).tolist()
-            scores = prediction["scores"].tolist()
-            labels = prediction["labels"].tolist()
-            keypoints = prediction["keypoints"]
-            keypoints = keypoints.flatten(start_dim=1).tolist()
-
-            coco_results.extend(
-                [
-                    {
-                        "image_id": original_id,
-                        "category_id": labels[k],
-                        'keypoints': keypoint,
-                        "score": scores[k],
-                    }
-                    for k, keypoint in enumerate(keypoints)
-                ]
-            )
-        return coco_results
-
-
-def convert_to_xywh(boxes):
-    xmin, ymin, xmax, ymax = boxes.unbind(1)
-    return torch.stack((xmin, ymin, xmax - xmin, ymax - ymin), dim=1)
-
-def merge(img_ids, eval_imgs):
-    all_img_ids = dist_utils.all_gather(img_ids)
-    all_eval_imgs = dist_utils.all_gather(eval_imgs)
-
-    merged_img_ids = []
-    for p in all_img_ids:
-        merged_img_ids.extend(p)
-
-    merged_eval_imgs = []
-    for p in all_eval_imgs:
-        merged_eval_imgs.extend(p)
-
-
-    merged_img_ids = np.array(merged_img_ids)
-    merged_eval_imgs = np.concatenate(merged_eval_imgs, axis=2).ravel()
-    # merged_eval_imgs = np.array(merged_eval_imgs).T.ravel()
-
-    # keep only unique (and in sorted order) images
-    merged_img_ids, idx = np.unique(merged_img_ids, return_index=True)
-
-    return merged_img_ids.tolist(), merged_eval_imgs.tolist()
+class CocoEvaluator(FasterCocoEvaluator):
+    pass