Merge pull request #27 from MiXaiLL76/getCatIds_fix

Worked out the ability to work with skeletons and various key points

examples/eval_example.ipynb

@@ -10,7 +10,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "faster_coco_eval.__version__='1.5.1'\n"
+      "faster_coco_eval.__version__='1.5.3'\n"
      ]
     }
    ],

faster_coco_eval/core/coco.py

@@ -90,7 +90,7 @@ def __init__(self, annotation_file=None):
             if type(annotation_file) is str:
                 self.dataset = self.load_json(annotation_file)
             elif type(annotation_file) is dict:
-                self.dataset = annotation_file
+                self.dataset = copy.deepcopy(annotation_file)
             else:
                 self.dataset = None
 
@@ -322,7 +322,8 @@ def loadRes(self, resFile, min_score=0):
         elif type(resFile) is np.ndarray:
             anns = self.loadNumpyAnnotations(resFile)
         else:
-            anns = resFile
+            anns = copy.deepcopy(resFile)
+
         assert type(anns) is list, "results in not an array of objects"
 
         anns = [ann for ann in anns if ann.get("score", 1) >= self.score_tresh]

faster_coco_eval/core/cocoeval.py

@@ -70,6 +70,7 @@ def __init__(
         iouType="segm",
         print_function=logger.debug,
         extra_calc=False,
+        kpt_oks_sigmas=None,
     ):
         """Initialize CocoEval using coco APIs for gt and dt :param cocoGt:
 
@@ -88,7 +89,9 @@ def __init__(
         self.eval: dict = {}  # accumulated evaluation results
         self._gts = defaultdict(list)  # gt for evaluation
         self._dts = defaultdict(list)  # dt for evaluation
-        self.params = Params(iouType=iouType)  # parameters
+        self.params = Params(
+            iouType=iouType, kpt_sigmas=kpt_oks_sigmas
+        )  # parameters
         self._paramsEval: dict = {}  # parameters for evaluation
         self.stats: list = []  # result summarization
         self.ious: dict = {}  # ious between all gts and dts
@@ -100,6 +103,18 @@ def __init__(
             self.params.imgIds = sorted(cocoGt.getImgIds())
             self.params.catIds = sorted(cocoGt.getCatIds())
 
+            if iouType == "keypoints":
+                if cocoDt is not None:
+                    self.params.catIds = sorted(list(cocoDt.cat_img_map.keys()))
+                else:
+                    self.params.catIds = sorted(
+                        [
+                            category_id
+                            for category_id, category in cocoGt.cats.items()
+                            if len(category.get("keypoints", []))
+                        ]
+                    )
+
         self.print_function = print_function  # output print function
 
     def _prepare(self):
@@ -136,7 +151,7 @@ def _toMask(anns, coco):
             gt["ignore"] = gt["ignore"] if "ignore" in gt else 0
             gt["ignore"] = "iscrowd" in gt and gt["iscrowd"]
             if p.iouType == "keypoints":
-                gt["ignore"] = (gt["num_keypoints"] == 0) or gt["ignore"]
+                gt["ignore"] = (gt.get("num_keypoints") == 0) or gt["ignore"]
         self._gts = defaultdict(list)  # gt for evaluation
         self._dts = defaultdict(list)  # dt for evaluation
         for gt in gts:
@@ -466,16 +481,32 @@ def _summarizeDets():
 
         def _summarizeKps():
             stats = np.zeros((10,))
-            stats[0] = _summarize(1, maxDets=20)
-            stats[1] = _summarize(1, maxDets=20, iouThr=0.5)
-            stats[2] = _summarize(1, maxDets=20, iouThr=0.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=0.5)
-            stats[7] = _summarize(0, maxDets=20, iouThr=0.75)
-            stats[8] = _summarize(0, maxDets=20, areaRng="medium")
-            stats[9] = _summarize(0, maxDets=20, areaRng="large")
+            stats[0] = _summarize(1, maxDets=self.params.maxDets[-1])
+            stats[1] = _summarize(
+                1, maxDets=self.params.maxDets[-1], iouThr=0.5
+            )
+            stats[2] = _summarize(
+                1, maxDets=self.params.maxDets[-1], iouThr=0.75
+            )
+            stats[3] = _summarize(
+                1, maxDets=self.params.maxDets[-1], areaRng="medium"
+            )
+            stats[4] = _summarize(
+                1, maxDets=self.params.maxDets[-1], areaRng="large"
+            )
+            stats[5] = _summarize(0, maxDets=self.params.maxDets[-1])
+            stats[6] = _summarize(
+                0, maxDets=self.params.maxDets[-1], iouThr=0.5
+            )
+            stats[7] = _summarize(
+                0, maxDets=self.params.maxDets[-1], iouThr=0.75
+            )
+            stats[8] = _summarize(
+                0, maxDets=self.params.maxDets[-1], areaRng="medium"
+            )
+            stats[9] = _summarize(
+                0, maxDets=self.params.maxDets[-1], areaRng="large"
+            )
             return stats
 
         if not self.eval:
@@ -534,6 +565,7 @@ def setKpParams(self):
         ]
         self.areaRngLbl = ["all", "medium", "large"]
         self.useCats = 1
+
         self.kpt_oks_sigmas = (
             np.array(
                 [
@@ -559,11 +591,13 @@ def setKpParams(self):
             / 10.0
         )
 
-    def __init__(self, iouType="segm"):
+    def __init__(self, iouType="segm", kpt_sigmas=None):
         if iouType == "segm" or iouType == "bbox":
             self.setDetParams()
         elif iouType == "keypoints":
             self.setKpParams()
+            if kpt_sigmas is not None:
+                self.kpt_oks_sigmas = np.array(kpt_sigmas)
         else:
             raise Exception("iouType not supported")
         self.iouType = iouType

faster_coco_eval/core/faster_eval_api.py

@@ -169,7 +169,7 @@ def accumulate(self):
                         self.eval["counts"][0], self.eval["counts"][3], -1
                     )
                 )
-                assert self.ground_truth_matches.shape[1] == len(
+                assert self.ground_truth_matches.shape[1] <= len(
                     self.cocoGt.anns
                 )
 
@@ -178,7 +178,7 @@ def accumulate(self):
                         self.eval["counts"][0], self.eval["counts"][3], -1
                     )
                 )
-                assert self.ground_truth_orig_id.shape[1] == len(
+                assert self.ground_truth_orig_id.shape[1] <= len(
                     self.cocoGt.anns
                 )
                 self.math_matches()
@@ -262,9 +262,21 @@ def computeAnnIoU(self, gt_ann, dt_ann):
             g.append(gt_ann["bbox"])
             d.append(dt_ann["bbox"])
 
-        return maskUtils.iou(d, g, [0]).max()
+        iscrowd = [0 for _ in g]
+
+        _iou = maskUtils.iou(d, g, iscrowd)
+
+        if len(_iou) == 0:
+            return 0
+        else:
+            return _iou.max()
 
     def compute_mIoU(self, categories=None, raw=False):
+        if self.params.iouType == "keypoints":
+            return np.array(
+                [val.max() for val in self.ious.values() if len(val)]
+            ).mean()
+
         g = []
         d = []
         s = []
@@ -283,7 +295,7 @@ def compute_mIoU(self, categories=None, raw=False):
                     else:
                         raise Exception("unknown iouType for iou computation")
 
-        iscrowd = [0 for o in g]
+        iscrowd = [0 for _ in g]
 
         ious = maskUtils.iou(d, g, iscrowd)
         if raw:

faster_coco_eval/extra/draw.py

@@ -26,6 +26,7 @@ def generate_ann_polygon(
     iouType: str = "bbox",
     text: Optional[str] = None,
     legendgroup: Optional[str] = None,
+    category_id_to_skeleton: Optional[dict] = None,
 ):
     all_x = []
     all_y = []
@@ -34,7 +35,7 @@ def generate_ann_polygon(
         x1, y1, w, h = ann["bbox"]
         all_x = [x1, x1 + w, x1 + w, x1, x1, None]
         all_y = [y1, y1, y1 + h, y1 + h, y1, None]
-    else:
+    elif iouType == "segm":
         ann = convert_ann_rle_to_poly(ann)
 
         for poly in ann["segmentation"]:
@@ -43,6 +44,23 @@ def generate_ann_polygon(
                 poly = np.array(poly).reshape(-1, 2)
                 all_x += poly[:, 0].tolist() + [None]
                 all_y += poly[:, 1].tolist() + [None]
+    elif iouType == "keypoints":
+        skeleton = category_id_to_skeleton.get(ann.get("category_id"))
+        keypoints = ann.get("keypoints")
+
+        if (skeleton is None) or (keypoints is None):
+            return
+
+        xyz = np.int0(keypoints).reshape(-1, 3)
+        ready_bones = {i: True for i in range(xyz.shape[0])}
+        for p1, p2 in skeleton:
+            if ready_bones.get(p1 - 1, False) and ready_bones.get(
+                p2 - 1, False
+            ):
+                all_x += [xyz[int(p1 - 1), 0], xyz[int(p2 - 1), 0], None]
+                all_y += [xyz[int(p1 - 1), 1], xyz[int(p2 - 1), 1], None]
+    else:
+        raise ValueError()
 
     return go.Scatter(
         x=all_x,
@@ -103,6 +121,10 @@ def display_image(
     categories_labels = {
         category["id"]: category["name"] for _, category in cocoGt.cats.items()
     }
+    category_id_to_skeleton = {
+        category["id"]: category.get("skeleton")
+        for _, category in cocoGt.cats.items()
+    }
 
     if len(gt_anns) > 0:
         for ann in gt_anns.values():
@@ -118,8 +140,10 @@ def display_image(
                                 categories_labels[ann["category_id"]],
                             ),
                             legendgroup="fn",
+                            category_id_to_skeleton=category_id_to_skeleton,
                         )
-                        polygons.append(poly)
+                        if poly is not None:
+                            polygons.append(poly)
                 else:
                     if display_gt:
                         poly = generate_ann_polygon(
@@ -131,8 +155,10 @@ def display_image(
                                 categories_labels[ann["category_id"]],
                             ),
                             legendgroup="gt",
+                            category_id_to_skeleton=category_id_to_skeleton,
                         )
-                        polygons.append(poly)
+                        if poly is not None:
+                            polygons.append(poly)
 
     if len(dt_anns) > 0:
         for ann in dt_anns.values():
@@ -156,8 +182,10 @@ def display_image(
                                 ann["iou"],
                             ),
                             legendgroup="tp",
+                            category_id_to_skeleton=category_id_to_skeleton,
                         )
-                        polygons.append(poly)
+                        if poly is not None:
+                            polygons.append(poly)
                 else:
                     if display_fp:
                         poly = generate_ann_polygon(
@@ -175,8 +203,10 @@ def display_image(
                                 ann["score"],
                             ),
                             legendgroup="fp",
+                            category_id_to_skeleton=category_id_to_skeleton,
                         )
-                        polygons.append(poly)
+                        if poly is not None:
+                            polygons.append(poly)
 
     fig = px.imshow(
         im,

faster_coco_eval/extra/extra.py

@@ -18,6 +18,7 @@ def __init__(
         iou_tresh: float = 0.0,
         recall_count: int = 100,
         useCats: bool = False,
+        kpt_oks_sigmas: list = None,
     ):
         self.iouType = iouType
         self.min_score = min_score
@@ -28,6 +29,12 @@ def __init__(
         self.cocoDt = copy.deepcopy(cocoDt)
         self.eval = None
 
+        if iouType == "keypoints":
+            self.useCats = True
+            self.kpt_oks_sigmas = np.array(kpt_oks_sigmas)
+        else:
+            self.kpt_oks_sigmas = None
+
         assert self.cocoGt is not None, "cocoGt is empty"
 
         if (self.cocoGt is not None) and (self.cocoDt is not None):
@@ -38,7 +45,11 @@ def evaluate(self):
         assert self.cocoDt is not None, "cocoDt is empty"
 
         cocoEval = COCOeval_faster(
-            self.cocoGt, self.cocoDt, self.iouType, extra_calc=True
+            self.cocoGt,
+            self.cocoDt,
+            self.iouType,
+            extra_calc=True,
+            kpt_oks_sigmas=self.kpt_oks_sigmas,
         )
         cocoEval.params.maxDets = [len(self.cocoGt.anns)]
 

faster_coco_eval/version.py

@@ -1,2 +1,2 @@
-__version__ = "1.5.2"
+__version__ = "1.5.3"
 __author__ = "MiXaiLL76"