fix off by one error in mask rcnn (cvat-ai#801)

cleemesser · Mar 5, 2020 · db62c3a · db62c3a
1 parent 83e0f79
commit db62c3a
Showing 1 changed file with 8 additions and 8 deletions.
diff --git a/utils/open_model_zoo/mask_rcnn_inception_resnet_v2_atrous_coco/interp.py b/utils/open_model_zoo/mask_rcnn_inception_resnet_v2_atrous_coco/interp.py
@@ -1,5 +1,6 @@
 import numpy as np
 import cv2
+from skimage.measure import approximate_polygon, find_contours
 
 
 MASK_THRESHOLD = .5
@@ -43,22 +44,21 @@ def segm_postprocess(box: list, raw_cls_mask, im_h, im_w, threshold):
             y = box[4] * height
             right = box[5] * width
             bottom = box[6] * height
-            mask = masks[index][label]
+            mask = masks[index][label - 1]
 
             mask = segm_postprocess((x, y, right, bottom),
                                     mask,
                                     height,
                                     width,
                                     MASK_THRESHOLD)
 
-            contour, _ = cv2.findContours(mask,
-                                          cv2.RETR_EXTERNAL,
-                                          cv2.CHAIN_APPROX_TC89_KCOS)
+            contours = find_contours(mask, MASK_THRESHOLD)
+            contour = contours[0]
+            contour = np.flip(contour, axis=1)
+            contour = approximate_polygon(contour, tolerance=2.5)
+            segmentation = contour.tolist()
 
-            contour = contour[0]
-            contour = contour.tolist()
-            contour = [x[0] for x in contour]
 
             # NOTE: if you want to see the boxes, uncomment next line
             # results.add_box(x, y, right, bottom, label, frame_number)
-            results.add_polygon(contour, label, frame_number)
+            results.add_polygon(segmentation, label, frame_number)