Use mask AP in compute_AP instead

kunthet · Feb 12, 2018 · 4b36a80 · 4b36a80
1 parent a208fa6
commit 4b36a80
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 10 deletions.
diff --git a/.gitignore b/.gitignore
@@ -1,3 +1,5 @@
 .ipynb_checkpoints/
 __pycache__/
 mask_rcnn_coco.h5
+.idea/
+
diff --git a/inspect_model.ipynb b/inspect_model.ipynb
@@ -358,8 +358,8 @@
    ],
    "source": [
     "# Draw precision-recall curve\n",
-    "AP, precisions, recalls, overlaps = utils.compute_ap(gt_bbox, gt_class_id, \n",
-    "                                          r['rois'], r['class_ids'], r['scores'])\n",
+    "AP, precisions, recalls, overlaps = utils.compute_ap(gt_bbox, gt_class_id, gt_mask,\n",
+    "                                          r['rois'], r['class_ids'], r['scores'], r['masks'])\n",
     "visualize.plot_precision_recall(AP, precisions, recalls)"
    ]
   },
@@ -427,8 +427,8 @@
     "        # Compute AP\n",
     "        r = results[0]\n",
     "        AP, precisions, recalls, overlaps =\\\n",
-    "            utils.compute_ap(gt_bbox, gt_class_id,\n",
-    "                              r['rois'], r['class_ids'], r['scores'])\n",
+    "            utils.compute_ap(gt_bbox, gt_class_id, gt_mask,\n",
+    "                              r['rois'], r['class_ids'], r['scores'], r['masks'])\n",
     "        APs.append(AP)\n",
     "    return APs\n",
     "\n",

diff --git a/train_shapes.ipynb b/train_shapes.ipynb
@@ -994,8 +994,8 @@
     "    r = results[0]\n",
     "    # Compute AP\n",
     "    AP, precisions, recalls, overlaps =\\\n",
-    "        utils.compute_ap(gt_bbox, gt_class_id,\n",
-    "                         r[\"rois\"], r[\"class_ids\"], r[\"scores\"])\n",
+    "        utils.compute_ap(gt_bbox, gt_class_id, gt_mask,\n",
+    "                         r[\"rois\"], r[\"class_ids\"], r[\"scores\"], r['masks'])\n",
     "    APs.append(AP)\n",
     "    \n",
     "print(\"mAP: \", np.mean(APs))"

diff --git a/utils.py b/utils.py
@@ -92,6 +92,7 @@ def compute_overlaps(boxes1, boxes2):
         overlaps[:, i] = compute_iou(box2, boxes1, area2[i], area1)
     return overlaps
 
+
 def compute_overlaps_masks(masks1, masks2):
     '''Computes IoU overlaps between two sets of masks.
     masks1, masks2: [Height, Width, instances]
@@ -288,6 +289,7 @@ def prepare(self, class_map=None):
         TODO: class map is not supported yet. When done, it should handle mapping
               classes from different datasets to the same class ID.
         """
+
         def clean_name(name):
             """Returns a shorter version of object names for cleaner display."""
             return ",".join(name.split(",")[:1])
@@ -582,8 +584,8 @@ def trim_zeros(x):
     return x[~np.all(x == 0, axis=1)]
 
 
-def compute_ap(gt_boxes, gt_class_ids,
-               pred_boxes, pred_class_ids, pred_scores,
+def compute_ap(gt_boxes, gt_class_ids, gt_masks,
+               pred_boxes, pred_class_ids, pred_scores, pred_masks,
                iou_threshold=0.5):
     """Compute Average Precision at a set IoU threshold (default 0.5).
 
@@ -596,15 +598,17 @@ def compute_ap(gt_boxes, gt_class_ids,
     # Trim zero padding and sort predictions by score from high to low
     # TODO: cleaner to do zero unpadding upstream
     gt_boxes = trim_zeros(gt_boxes)
+    gt_masks = gt_masks[..., :gt_boxes.shape[0]]
     pred_boxes = trim_zeros(pred_boxes)
     pred_scores = pred_scores[:pred_boxes.shape[0]]
     indices = np.argsort(pred_scores)[::-1]
     pred_boxes = pred_boxes[indices]
     pred_class_ids = pred_class_ids[indices]
     pred_scores = pred_scores[indices]
+    pred_masks = pred_masks[..., indices]
 
-    # Compute IoU overlaps [pred_boxes, gt_boxes]
-    overlaps = compute_overlaps(pred_boxes, gt_boxes)
+    # Compute IoU overlaps [pred_masks, gt_masks]
+    overlaps = compute_overlaps_masks(pred_masks, gt_masks)
 
     # Loop through ground truth boxes and find matching predictions
     match_count = 0