recommended changes

Author: Matt Sokoloff

labelbox/data/annotation_types/geometry/geometry.py

@@ -1,4 +1,4 @@
-from typing import Dict, Any, Union
+from typing import Dict, Any, Optional, Union
 
 import geojson
 import numpy as np
@@ -25,3 +25,12 @@ def shapely(
 
     def raster(self, *args, **kwargs) -> np.ndarray:
         raise NotImplementedError("Subclass must override this")
+
+    def get_or_create_canvas(self, height: Optional[int], width: Optional[int],
+                             canvas: Optional[np.ndarray]):
+        if canvas is None:
+            if height is None or width is None:
+                raise ValueError(
+                    "Must either provide canvas or height and width")
+            canvas = np.zeros((height, width, 3), dtype=np.uint8)
+        return canvas

labelbox/data/annotation_types/geometry/line.py

@@ -1,4 +1,4 @@
-from typing import List, Optional
+from typing import List, Optional, Union, Tuple
 
 import geojson
 import numpy as np
@@ -19,25 +19,22 @@ def geometry(self) -> geojson.MultiLineString:
     def raster(self,
                height: Optional[int] = None,
                width: Optional[int] = None,
-               thickness=1,
-               color=(255, 255, 255),
-               canvas=None) -> np.ndarray:
+               thickness: int = 1,
+               color: Union[Tuple, int] = (255, 255, 255),
+               canvas: Optional[np.ndarray] = None) -> np.ndarray:
         """
         Draw the line onto a 3d mask
         Args:
             height (int): height of the mask
             width (int): width of the mask
             thickness (int): How thick to draw the line
-            color (int): color for the line. Only a single int since this is a grayscale mask.
+            color (int): color for the line.
+                  RGB values by default but if a 2D canvas is provided this can set this to an int.
+            canvas (np.ndarry): Canvas for drawing line on.
         Returns:
             numpy array representing the mask with the line drawn on it.
         """
-        if canvas is None:
-            if height is None or width is None:
-                raise ValueError(
-                    "Must either provide canvas or height and width")
-            canvas = np.zeros((height, width, 3), dtype=np.uint8)
-
+        canvas = self.get_or_create_canvas(height, width, canvas)
         pts = np.array(self.geometry['coordinates']).astype(np.int32)
         return cv2.polylines(canvas,
                              pts,

labelbox/data/annotation_types/geometry/point.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, Tuple, Union
 
 import geojson
 import numpy as np
@@ -19,23 +19,21 @@ def raster(self,
                height: Optional[int] = None,
                width: Optional[int] = None,
                thickness: int = 1,
-               color=(255, 255, 255),
-               canvas=None) -> np.ndarray:
+               color: Union[Tuple, int] = (255, 255, 255),
+               canvas: Optional[np.ndarray] = None) -> np.ndarray:
         """
         Draw the point onto a 3d mask
         Args:
             height (int): height of the mask
             width (int): width of the mask
             thickness (int): pixel radius of the point
-            color (int): color for the point. Only a single int since this is a grayscale mask.
+            color (int): color for the point.
+                  RGB values by default but if a 2D canvas is provided this can set this to an int.
+            canvas (np.ndarray): Canvas to draw the point on
         Returns:
             numpy array representing the mask with the point drawn on it.
         """
-        if canvas is None:
-            if height is None or width is None:
-                raise ValueError(
-                    "Must either provide canvas or height and width")
-            canvas = np.zeros((height, width, 3), dtype=np.uint8)
+        canvas = self.get_or_create_canvas(height, width, canvas)
         return cv2.circle(canvas, (int(self.x), int(self.y)),
                           radius=thickness,
                           color=color,

labelbox/data/annotation_types/geometry/polygon.py

@@ -1,4 +1,4 @@
-from typing import List, Optional
+from typing import List, Optional, Union, Tuple
 
 import numpy as np
 import geojson
@@ -21,23 +21,22 @@ def geometry(self) -> geojson.MultiPolygon:
     def raster(self,
                height: Optional[int] = None,
                width: Optional[int] = None,
-               color=(255, 255, 255),
-               thickness=-1,
-               canvas=None) -> np.ndarray:
+               color: Union[int, Tuple] = (255, 255, 255),
+               thickness: int = -1,
+               canvas: Optional[np.ndarray] = None) -> np.ndarray:
         """
         Draw the polygon onto a 3d mask
         Args:
             height (int): height of the mask
             width (int): width of the mask
-            color (int): color for the polygon. Only a single int since this is a grayscale mask.
+            color (int): color for the polygon.
+                  RGB values by default but if a 2D canvas is provided this can set this to an int.
+            thickness (int): How thick to make the polygon border. -1 fills in the polygon
+            canvas (np.ndarray): Canvas to draw the polygon on
         Returns:
             numpy array representing the mask with the polygon drawn on it.
         """
-        if canvas is None:
-            if height is None or width is None:
-                raise ValueError(
-                    "Must either provide canvas or height and width")
-            canvas = np.zeros((height, width, 3), dtype=np.uint8)
+        canvas = self.get_or_create_canvas(height, width, canvas)
         pts = np.array(self.geometry['coordinates']).astype(np.int32)
         if thickness == -1:
             return cv2.fillPoly(canvas, pts, color)

labelbox/data/annotation_types/geometry/rectangle.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Optional, Union, Tuple
 
 import cv2
 import geojson
@@ -31,23 +31,22 @@ def geometry(self) -> geojson.geometry.Geometry:
     def raster(self,
                height: Optional[int] = None,
                width: Optional[int] = None,
-               color=(255, 255, 255),
-               thickness=-1,
-               canvas=None) -> np.ndarray:
+               color: Union[int, Tuple] = (255, 255, 255),
+               thickness: int = -1,
+               canvas: Optional[np.ndarray] = None) -> np.ndarray:
         """
         Draw the rectangle onto a 3d mask
         Args:
             height (int): height of the mask
             width (int): width of the mask
-            color (int): color for the rectangle. Only a single int since this is a grayscale mask.
+            color (int): color for the polygon.
+                  RGB values by default but if a 2D canvas is provided this can set this to an int.
+            thickness (int): How thick to make the rectangle border. -1 fills in the rectangle
+            canvas (np.ndarray): Canvas to draw rectangle on
         Returns:
             numpy array representing the mask with the rectangle drawn on it.
         """
-        if canvas is None:
-            if height is None or width is None:
-                raise ValueError(
-                    "Must either provide canvas or height and width")
-            canvas = np.zeros((height, width, 3), dtype=np.uint8)
+        canvas = self.get_or_create_canvas(height, width, canvas)
         pts = np.array(self.geometry['coordinates']).astype(np.int32)
         if thickness == -1:
             return cv2.fillPoly(canvas, pts, color)

labelbox/data/serialization/labelbox_v1/converter.py

@@ -29,13 +29,10 @@ def deserialize_video(json_data: Iterable[Dict[str, Any]],
         Returns:
             LabelGenerator containing the video data.
         """
-
-        def label_generator():
-            for example in LBV1VideoIterator(json_data, client):
-                if example['Label']:
-                    yield LBV1Label(**example).to_common()
-
-        return LabelGenerator(data=label_generator())
+        label_generator = (LBV1Label(**example).to_common()
+                           for example in LBV1VideoIterator(json_data, client)
+                           if example['Label'])
+        return LabelGenerator(data=label_generator)
 
     @staticmethod
     def deserialize(json_data: Iterable[Dict[str, Any]]) -> LabelGenerator: