update drawing logic

Author: Matt Sokoloff

labelbox/data/annotation_types/geometry/geometry.py

@@ -1,4 +1,4 @@
-from typing import Dict, Any, Optional, Union
+from typing import Dict, Any, Optional, Union, Tuple
 
 import geojson
 import numpy as np
@@ -23,7 +23,12 @@ def shapely(
                geom.MultiLineString, geom.MultiPolygon]:
         return geom.shape(self.geometry)
 
-    def raster(self, *args, **kwargs) -> np.ndarray:
+    def raster(self,
+               height: Optional[int] = None,
+               width: Optional[int] = None,
+               canvas: Optional[np.ndarray] = None,
+               color: Optional[Union[int, Tuple[int,int,int]]] = None,
+               thickness: Optional[int] = 1) -> np.ndarray:
         raise NotImplementedError("Subclass must override this")
 
     def get_or_create_canvas(self, height: Optional[int], width: Optional[int],

labelbox/data/annotation_types/geometry/line.py

@@ -17,11 +17,11 @@ def geometry(self) -> geojson.MultiLineString:
             [[[point.x, point.y] for point in self.points]])
 
     def raster(self,
-               height: Optional[int] = None,
+              height: Optional[int] = None,
                width: Optional[int] = None,
-               thickness: int = 1,
-               color: Union[Tuple, int] = (255, 255, 255),
-               canvas: Optional[np.ndarray] = None) -> np.ndarray:
+               canvas: Optional[np.ndarray] = None,
+               color: Union[int, Tuple[int,int,int]] = (255, 255, 255),
+               thickness: int = 5) -> np.ndarray:
         """
         Draw the line onto a 3d mask
         Args:

labelbox/data/annotation_types/geometry/mask.py

@@ -18,7 +18,7 @@ class Mask(Geometry):
 
     @property
     def geometry(self):
-        mask = self.raster(binary=True)
+        mask = self.raster(color = 1)
         polygons = (
             shape(shp)
             for shp, val in shapes(mask, mask=None)
@@ -29,28 +29,48 @@ def geometry(self):
     def raster(self,
                height: Optional[int] = None,
                width: Optional[int] = None,
-               binary=False) -> np.ndarray:
+               canvas: Optional[np.ndarray] = None,
+               color: Optional[Union[int, Tuple[int,int,int]]] = None,
+               thickness = None
+               ) -> np.ndarray:
         """
+        # TODO: Optionally use the color. a color of 1 will result in a binary canvas
+
+
         Removes all pixels from the segmentation mask that do not equal self.color
 
         Args:
-            height:
-
+            height (int): Optionally resize mask height before drawing.
+            width (int): Optionally resize mask width before drawing.
+            canvas (np.ndarray): Optionall provide a canvas to draw on
+            color (Union[int, Tuple[int,int,int]]): Color to draw the canvas.
+                Defaults to using the encoded color in the mask.
+                int will return the mask as a 1d array
+                tuple[int,int,int] will return the mask as a 3d array
+            thickness (None): Unused, exists for a consistent interface.
         Returns:
             np.ndarray representing only this object
+                as opposed to the mask that this object references which might have multiple objects determined by colors
         """
+
         mask = self.mask.value
         mask = np.alltrue(mask == self.color, axis=2).astype(np.uint8)
+
         if height is not None or width is not None:
             mask = cv2.resize(mask,
                               (width or mask.shape[1], height or mask.shape[0]))
-        if binary:
-            return mask
-        else:
-            color_image = np.zeros((mask.shape[0], mask.shape[1], 3),
-                                   dtype=np.uint8)
-            color_image[mask.astype(np.bool)] = self.color
-            return color_image
+
+        dims = [mask.shape[0], mask.shape[1]]
+        color = color or self.color
+        if isinstance(color, (tuple, list)):
+            dims = dims + [len(color)]
+
+        canvas = canvas if canvas is not None else np.zeros(tuple(dims), dtype=np.uint8)
+        canvas[mask.astype(np.bool)] = color
+        return canvas
+
+
+
 
     def create_url(self, signer: Callable[[bytes], str]) -> str:
         """

labelbox/data/annotation_types/geometry/point.py

@@ -18,9 +18,10 @@ def geometry(self) -> geojson.Point:
     def raster(self,
                height: Optional[int] = None,
                width: Optional[int] = None,
-               thickness: int = 1,
-               color: Union[Tuple, int] = (255, 255, 255),
-               canvas: Optional[np.ndarray] = None) -> np.ndarray:
+               canvas: Optional[np.ndarray] = None,
+               color: Union[int, Tuple[int,int,int]] = (255, 255, 255),
+               thickness: int = 10
+               ) -> np.ndarray:
         """
         Draw the point onto a 3d mask
         Args:

labelbox/data/annotation_types/geometry/polygon.py

@@ -21,9 +21,9 @@ def geometry(self) -> geojson.MultiPolygon:
     def raster(self,
                height: Optional[int] = None,
                width: Optional[int] = None,
-               color: Union[int, Tuple] = (255, 255, 255),
-               thickness: int = -1,
-               canvas: Optional[np.ndarray] = None) -> np.ndarray:
+               canvas: Optional[np.ndarray] = None,
+               color: Union[int, Tuple[int,int,int]] = (255, 255, 255),
+               thickness: int = -1) -> np.ndarray:
         """
         Draw the polygon onto a 3d mask
         Args:

labelbox/data/annotation_types/geometry/rectangle.py

@@ -31,9 +31,9 @@ def geometry(self) -> geojson.geometry.Geometry:
     def raster(self,
                height: Optional[int] = None,
                width: Optional[int] = None,
-               color: Union[int, Tuple] = (255, 255, 255),
-               thickness: int = -1,
-               canvas: Optional[np.ndarray] = None) -> np.ndarray:
+               canvas: Optional[np.ndarray] = None,
+               color: Union[int, Tuple[int,int,int]] = (255, 255, 255),
+               thickness: int = -1) -> np.ndarray:
         """
         Draw the rectangle onto a 3d mask
         Args:

labelbox/data/metrics/iou.py

@@ -129,9 +129,9 @@ def mask_miou(ground_truths: List[ObjectAnnotation],
         float representing the iou score for the masks
     """
     prediction_np = np.max(
-        [pred.value.raster(binary=True) for pred in predictions], axis=0)
+        [pred.value.raster(color=1) for pred in predictions], axis=0)
     ground_truth_np = np.max([
-        ground_truth.value.raster(binary=True) for ground_truth in ground_truths
+        ground_truth.value.raster(color=1) for ground_truth in ground_truths
     ],
                              axis=0)
     if prediction_np.shape != ground_truth_np.shape:

tests/data/annotation_types/geometry/test_line.py

@@ -19,5 +19,5 @@ def test_line():
     expected['coordinates'] = tuple([tuple([tuple(x) for x in points])])
     assert line.shapely.__geo_interface__ == expected
 
-    raster = line.raster(height=32, width=32)
+    raster = line.raster(height=32, width=32, thickness = 1)
     assert (cv2.imread("tests/data/assets/line.png") == raster).all()

tests/data/annotation_types/geometry/test_point.py

@@ -18,5 +18,5 @@ def test_point():
     expected['coordinates'] = tuple(expected['coordinates'])
     assert point.shapely.__geo_interface__ == expected
 
-    raster = point.raster(height=32, width=32)
+    raster = point.raster(height=32, width=32, thickness = 1)
     assert (cv2.imread("tests/data/assets/point.png") == raster).all()