From e14006778993f35b2883dfcdbd73cb0fde7736f6 Mon Sep 17 00:00:00 2001
From: Aaron Lelevier <aaron.lelevier@gmail.com>
Date: Mon, 21 Jan 2019 08:56:36 -0800
Subject: [PATCH] add helper function 'cv2_utils.findContours' to maintain
 version compatibility (#354)

fixes #339
---
 demo/predictor.py                             |  3 ++-
 maskrcnn_benchmark/utils/cv2_util.py          | 24 +++++++++++++++++++
 .../instances2dict_with_polygons.py           |  4 +++-
 3 files changed, 29 insertions(+), 2 deletions(-)
 create mode 100644 maskrcnn_benchmark/utils/cv2_util.py

diff --git a/demo/predictor.py b/demo/predictor.py
index c8790c255..2f27b2fb2 100644
--- a/demo/predictor.py
+++ b/demo/predictor.py
@@ -8,6 +8,7 @@
 from maskrcnn_benchmark.structures.image_list import to_image_list
 from maskrcnn_benchmark.modeling.roi_heads.mask_head.inference import Masker
 from maskrcnn_benchmark import layers as L
+from maskrcnn_benchmark.utils import cv2_util
 
 
 class COCODemo(object):
@@ -287,7 +288,7 @@ def overlay_mask(self, image, predictions):
 
         for mask, color in zip(masks, colors):
             thresh = mask[0, :, :, None]
-            _, contours, hierarchy = cv2.findContours(
+            contours, hierarchy = cv2_util.findContours(
                 thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
             )
             image = cv2.drawContours(image, contours, -1, color, 3)
diff --git a/maskrcnn_benchmark/utils/cv2_util.py b/maskrcnn_benchmark/utils/cv2_util.py
new file mode 100644
index 000000000..0bbc0fb2d
--- /dev/null
+++ b/maskrcnn_benchmark/utils/cv2_util.py
@@ -0,0 +1,24 @@
+"""
+Module for cv2 utility functions and maintaining version compatibility
+between 3.x and 4.x
+"""
+import cv2
+
+
+def findContours(*args, **kwargs):
+    """
+    Wraps cv2.findContours to maintain compatiblity between versions
+    3 and 4
+
+    Returns:
+        contours, hierarchy
+    """
+    if cv2.__version__.startswith('4'):
+        contours, hierarchy = cv2.findContours(*args, **kwargs)
+    elif cv2.__version__.startswith('3'):
+        _, contours, hierarchy = cv2.findContours(*args, **kwargs)
+    else:
+        raise AssertionError(
+            'cv2 must be either version 3 or 4 to call this method')
+
+    return contours, hierarchy
diff --git a/tools/cityscapes/instances2dict_with_polygons.py b/tools/cityscapes/instances2dict_with_polygons.py
index fbfc8d13b..0bb560497 100644
--- a/tools/cityscapes/instances2dict_with_polygons.py
+++ b/tools/cityscapes/instances2dict_with_polygons.py
@@ -13,6 +13,8 @@
 from cityscapesscripts.evaluation.instance import *
 from cityscapesscripts.helpers.csHelpers import *
 import cv2
+from maskrcnn_benchmark.utils import cv2_util
+
 
 def instances2dict_with_polygons(imageFileList, verbose=False):
     imgCount     = 0
@@ -46,7 +48,7 @@ def instances2dict_with_polygons(imageFileList, verbose=False):
             #instances[id2label[instanceObj.labelID].name].append(instanceObj.toDict())
             if id2label[instanceObj.labelID].hasInstances:
                 mask = (imgNp == instanceId).astype(np.uint8)
-                im2, contour, hier = cv2.findContours(
+                contour, hier = cv2_util.findContours(
                     mask.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
 
                 polygons = [c.reshape(-1).tolist() for c in contour]