support polygon and rlepolygon convert to mask

dsdl/geometry/__init__.py

@@ -1,7 +1,7 @@
 from .box import BBox
 from .label import Label, LabelList
 from .media import Image
-from .polygon import Polygon, PolygonItem
+from .polygon import Polygon, PolygonItem, RLEPolygon
 from .segmap import SegmentationMap
 from .insmap import InstanceMap
 from .keypoint import Coord2D, KeyPoints
@@ -24,6 +24,7 @@
     "LabelList",
     "Polygon",
     "PolygonItem",
+    "RLEPolygon",
     "SegmentationMap",
     "InstanceMap",
     "Coord2D",

dsdl/geometry/polygon.py

@@ -1,8 +1,12 @@
 from typing import List
 import numpy as np
+import cv2
 from PIL import ImageDraw, Image
 from .base_geometry import BaseGeometry
-
+try:
+    import pycocotools.mask as mask_util
+except:
+    mask_util = None
 
 class PolygonItem(BaseGeometry):
 
@@ -108,6 +112,20 @@ def polygons(self):
             A list of all the PolygonItem objects in the current polygon object.
         """
         return self._data
+
+    def to_mask(self, imsize: List[int] = [1000, 1000]) -> np.array:
+        """
+        generate mask from polygon.
+        """
+        mask = np.zeros(imsize).astype(np.uint8)
+        points = []
+        for p in self.polygons:
+            pts = np.array(p.points, np.int32)
+            pts = pts.reshape((-1, 1, 2))
+            points.append(pts)
+
+        mask = cv2.fillPoly(mask, points, 1)
+        return mask
 
     @property
     def openmmlabformat(self) -> List[List[float]]:
@@ -173,7 +191,7 @@ def __repr__(self):
 class RLEPolygon(BaseGeometry):
     def __init__(self, rle_data, image_shape):
         """
-        rle_data: rle list
+        rle_data: rle list or bytes
         image_shape: [H, W]
         """
         self._rle_data = rle_data
@@ -187,23 +205,29 @@ def rle_data(self):
     def image_shape(self):
         return self._image_shape
 
-    @property
-    def mask(self):
+    def to_mask(self) -> np.array:
         '''
         # ref: https://www.kaggle.com/paulorzp/run-length-encode-and-decode
         mask_rle: run-length as string formated (start length)
         shape: (height,width) of array to return
         Returns numpy array, 1 - mask, 0 - background
         '''
-        s = self._rle_data
-        shape = self._image_shape
-        starts, lengths = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]
-        starts -= 1
-        ends = starts + lengths
-        img = np.zeros(shape[0] * shape[1], dtype=np.uint8)
-        for lo, hi in zip(starts, ends):
-            img[lo:hi] = 1
-        return img.reshape(shape)
+        if mask_util is not None:
+            coco_rle = {'size': self._image_shape,
+                        'counts': self._rle_data}
+            return mask_util.decode([coco_rle])[:,:,0]
+        else:
+            shape = self._image_shape
+            lens = list(self._rle_data)
+            starts = [sum(lens[0:i]) for i in range(len(lens))]
+            img = np.zeros(shape[0]*shape[1], dtype=np.uint8)
+            flag = 1
+            for begin, length in zip(starts, lens):
+                flag = 1 - flag
+                if not flag:
+                    continue
+                img[begin: begin+length] = 1
+            return img.reshape(shape)
 
     @property
     def openmmlabformat(self):