Merge pull request #235 from Labelbox/ms/consistent-drawing-logic

update drawing logic

Author: msokoloff1

examples/label_export/images.ipynb

@@ -38,7 +38,10 @@ def bytes_to_np(self, image_bytes: bytes) -> np.ndarray:
         Returns:
             numpy array representing the image
         """
-        return np.array(Image.open(BytesIO(image_bytes)))[:, :, :3]
+        arr = np.array(Image.open(BytesIO(image_bytes)))
+        if len(arr.shape) == 2:
+            arr = np.stack((arr,) * 3, axis=-1)
+        return arr[:, :, :3]
 
     def np_to_bytes(self, arr: np.ndarray) -> bytes:
         """

labelbox/data/annotation_types/data/raster.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/geometry.py

@@ -19,9 +19,9 @@ def geometry(self) -> geojson.MultiLineString:
     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 = 1) -> np.ndarray:
         """
         Draw the line onto a 3d mask
         Args:

labelbox/data/annotation_types/geometry/line.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,45 @@ 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/mask.py

@@ -18,9 +18,9 @@ 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:
@@ -37,4 +37,4 @@ def raster(self,
         return cv2.circle(canvas, (int(self.x), int(self.y)),
                           radius=thickness,
                           color=color,
-                          thickness=-1)
+                          thickness=1)

labelbox/data/annotation_types/geometry/point.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/polygon.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/annotation_types/geometry/rectangle.py

@@ -128,12 +128,11 @@ def mask_miou(ground_truths: List[ObjectAnnotation],
     Returns:
         float representing the iou score for the masks
     """
-    prediction_np = np.max(
-        [pred.value.raster(binary=True) for pred in predictions], axis=0)
-    ground_truth_np = np.max([
-        ground_truth.value.raster(binary=True) for ground_truth in ground_truths
-    ],
-                             axis=0)
+    prediction_np = np.max([pred.value.raster(color=1) for pred in predictions],
+                           axis=0)
+    ground_truth_np = np.max(
+        [ground_truth.value.raster(color=1) for ground_truth in ground_truths],
+        axis=0)
     if prediction_np.shape != ground_truth_np.shape:
         raise ValueError(
             "Prediction and mask must have the same shape."

labelbox/data/metrics/iou.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()