created basics notebook and related bug fixes

Author: Matt Sokoloff

examples/annotation_types/annotation_type_basics.ipynb

@@ -0,0 +1,28 @@
+from .geometry import Line
+from .geometry import Point
+from .geometry import Mask
+from .geometry import Polygon
+from .geometry import Rectangle
+from .geometry import Geometry
+
+from .annotation import ClassificationAnnotation
+from .annotation import VideoClassificationAnnotation
+from .annotation import ObjectAnnotation
+from .annotation import VideoObjectAnnotation
+
+from .ner import TextEntity
+
+from .classification import Checklist
+from .classification import ClassificationAnswer
+from .classification import Dropdown
+from .classification import Radio
+from .classification import Text
+
+from .data import RasterData
+from .data import TextData
+from .data import VideoData
+
+from .label import Label
+
+from .collection import LabelCollection
+from .collection import LabelGenerator

labelbox/data/annotation_types/__init__.py

@@ -3,10 +3,12 @@
 
 import numpy as np
 import requests
+from typing_extensions import Literal
 from pydantic import root_validator
 from PIL import Image
 
 from .base_data import BaseData
+from ..types import TypedArray
 
 
 class RasterData(BaseData):
@@ -16,7 +18,7 @@ class RasterData(BaseData):
     im_bytes: Optional[bytes] = None
     file_path: Optional[str] = None
     url: Optional[str] = None
-    arr: Optional[np.ndarray] = None
+    arr: Optional[TypedArray[Literal['uint8']]] = None
 
     def bytes_to_np(self, image_bytes: bytes) -> np.ndarray:
         """
@@ -64,14 +66,23 @@ def data(self) -> np.ndarray:
             self.im_bytes = im_bytes
             return self.bytes_to_np(im_bytes)
         elif self.url is not None:
-            response = requests.get(self.url)
-            response.raise_for_status()
-            im_bytes = response.content
+            im_bytes = self.fetch_remote()
             self.im_bytes = im_bytes
             return self.bytes_to_np(im_bytes)
         else:
             raise ValueError("Must set either url, file_path or im_bytes")
 
+    def fetch_remote(self) -> bytes:
+        """
+        Method for accessing url.
+
+        If url is not publicly accessible or requires another access pattern
+        simply override this function
+        """
+        response = requests.get(self.url)
+        response.raise_for_status()
+        return response.content
+
     def create_url(self, signer: Callable[[bytes], str]) -> str:
         """
         Utility for creating a url from any of the other image representations.
@@ -95,7 +106,7 @@ def create_url(self, signer: Callable[[bytes], str]) -> str:
                 "One of url, im_bytes, file_path, arr must not be None.")
         return self.url
 
-    @root_validator
+    @root_validator()
     def validate_args(cls, values):
         file_path = values.get("file_path")
         im_bytes = values.get("im_bytes")
@@ -118,8 +129,6 @@ def validate_args(cls, values):
         return values
 
     class Config:
-        # Required for numpy arrays
-        arbitrary_types_allowed = True
         # Required for sharing references
         copy_on_model_validation = False
         # Required for discriminating between data types

labelbox/data/annotation_types/data/raster.py

@@ -30,14 +30,23 @@ def data(self) -> str:
             self.text = text
             return text
         elif self.url:
-            response = requests.get(self.url)
-            response.raise_for_status()
-            text = response.text
+            text = self.fetch_remote()
             self.text = text
             return text
         else:
             raise ValueError("Must set either url, file_path or im_bytes")
 
+    def fetch_remote(self) -> str:
+        """
+        Method for accessing url.
+
+        If url is not publicly accessible or requires another access pattern
+        simply override this function
+        """
+        response = requests.get(self.url)
+        response.raise_for_status()
+        return response.text
+
     def create_url(self, signer: Callable[[bytes], str]) -> None:
         """
         Utility for creating a url from any of the other text references.

labelbox/data/annotation_types/data/text.py

@@ -2,13 +2,15 @@
 import os
 import urllib.request
 from typing import Callable, Dict, Generator, Optional, Tuple
+from typing_extensions import Literal
 from uuid import uuid4
 
 import cv2
 import numpy as np
 from pydantic import root_validator
 
 from .base_data import BaseData
+from ..types import TypedArray
 
 logger = logging.getLogger(__name__)
 
@@ -19,7 +21,7 @@ class VideoData(BaseData):
     """
     file_path: Optional[str] = None
     url: Optional[str] = None
-    frames: Optional[Dict[int, np.ndarray]] = None
+    frames: Optional[Dict[int, TypedArray[Literal['uint8']]]] = None
 
     def load_frames(self, overwrite: bool = False) -> None:
         """
@@ -33,8 +35,14 @@ def load_frames(self, overwrite: bool = False) -> None:
             return
 
         for count, frame in self.frame_generator():
+            if self.frames is None:
+                self.frames = {}
             self.frames[count] = frame
 
+    @property
+    def data(self):
+        return self.frame_generator()
+
     def frame_generator(
             self,
             cache_frames=False,
@@ -48,26 +56,27 @@ def frame_generator(
             download_dir (str): Directory to save the video to. Defaults to `/tmp` dir
         """
         if self.frames is not None:
-            for idx, img in self.frames.items():
-                yield idx, img
+            for idx, frame in self.frames.items():
+                yield idx, frame
             return
         elif self.url and not self.file_path:
             file_path = os.path.join(download_dir, f"{uuid4()}.mp4")
             logger.info("Downloading the video locally to %s", file_path)
-            urllib.request.urlretrieve(self.url, file_path)
+            self.fetch_remote(file_path)
             self.file_path = file_path
 
         vidcap = cv2.VideoCapture(self.file_path)
 
-        success, img = vidcap.read()
+        success, frame = vidcap.read()
         count = 0
-        self.frames = {}
+        if cache_frames:
+            self.frames = {}
         while success:
-            img = img[:, :, ::-1]
-            yield count, img
+            frame = frame[:, :, ::-1]
+            yield count, frame
             if cache_frames:
-                self.frames[count] = img
-            success, img = vidcap.read()
+                self.frames[count] = frame
+            success, frame = vidcap.read()
             count += 1
 
     def __getitem__(self, idx: int) -> np.ndarray:
@@ -77,6 +86,18 @@ def __getitem__(self, idx: int) -> np.ndarray:
             )
         return self.frames[idx]
 
+    def fetch_remote(self, local_path) -> None:
+        """
+        Method for downloading data from self.url
+
+        If url is not publicly accessible or requires another access pattern
+        simply override this function
+
+        Args:
+            local_path: Where to save the thing too.
+        """
+        urllib.request.urlretrieve(self.url, local_path)
+
     def create_url(self, signer: Callable[[bytes], str]) -> None:
         """
         Utility for creating a url from any of the other video references.
@@ -134,7 +155,5 @@ def validate_data(cls, values):
         return values
 
     class Config:
-        # Required for numpy arrays
-        arbitrary_types_allowed = True
         # Required for discriminating between data types
         extra = 'forbid'

labelbox/data/annotation_types/data/video.py

@@ -21,6 +21,5 @@ class FeatureSchema(BaseModel):
     def must_set_one(cls, values):
         if values['schema_id'] is None and values['name'] is None:
             raise ValueError(
-                "Must set either schema_id or display name for all feature schemas"
-            )
+                "Must set either schema_id or name for all feature schemas")
         return values

labelbox/data/annotation_types/feature.py

@@ -1,6 +1,7 @@
-from typing import Callable, Tuple
+from typing import Callable, Tuple, Union
 
 import numpy as np
+from pydantic.class_validators import validator
 from rasterio.features import shapes
 from shapely.geometry import MultiPolygon, shape
 
@@ -11,21 +12,39 @@
 class Mask(Geometry):
     # Raster data can be shared across multiple masks... or not
     mask: RasterData
-    color_rgb: Tuple[int, int, int]
+    # RGB or Grayscale
+    color: Union[int, Tuple[int, int, int]]
 
     @property
     def geometry(self):
-        mask = self.mask.data
-        mask = np.alltrue(mask == self.color_rgb, axis=2).astype(np.uint8)
+        mask = self.raster(binary=True)
         polygons = (
             shape(shp)
             for shp, val in shapes(mask, mask=None)
             # ignore if shape is area of smaller than 1 pixel
             if val >= 1)
         return MultiPolygon(polygons).__geo_interface__
 
-    def raster(self) -> np.ndarray:
-        return self.mask.data
+    def raster(self, binary=False) -> np.ndarray:
+        """
+        Removes all pixels from the segmentation mask that do not equal self.color
+
+        Returns:
+            np.ndarray representing only this object
+        """
+        mask = self.mask.data
+        if len(mask.shape) == 2:
+            mask = np.expand_dims(mask, axis=-1)
+        mask = np.alltrue(mask == self.color, axis=2).astype(np.uint8)
+        if binary:
+            return mask
+        elif isinstance(self.color, int):
+            return mask * self.color
+        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
 
     def create_url(self, signer: Callable[[bytes], str]) -> str:
         """
@@ -38,3 +57,21 @@ def create_url(self, signer: Callable[[bytes], str]) -> str:
             the url for the mask
         """
         return self.mask.create_url(signer)
+
+    @validator('color')
+    def is_valid_color(cls, color):
+        #Does the dtype matter? Can it be a float?
+        if isinstance(color, (tuple, list)):
+            if len(color) != 3:
+                raise ValueError(
+                    "Segmentation colors must be either a (r,g,b) tuple or a single grayscale value"
+                )
+            elif not all([0 <= c <= 255 for c in color]):
+                raise ValueError(
+                    f"All rgb colors must be between 0 and 255. Found : {color}"
+                )
+        elif not (0 <= color <= 255):
+            raise ValueError(
+                f"All rgb colors must be between 0 and 255. Found : {color}")
+
+        return color

labelbox/data/annotation_types/geometry/mask.py

@@ -9,6 +9,9 @@
 class Rectangle(Geometry):
     """
     Represents a 2d rectangle. Also known as a bounding box.
+
+    start: Top left coordinate of the rectangle
+    end: Bottom right coordinate of the rectangle
     """
     start: Point
     end: Point

labelbox/data/annotation_types/geometry/rectangle.py

@@ -11,6 +11,8 @@
 from .annotation import (ClassificationAnnotation, ObjectAnnotation,
                          VideoClassificationAnnotation, VideoObjectAnnotation)
 
+from pydantic import validator
+
 
 class Label(BaseModel):
     data: Union[VideoData, RasterData, TextData]
@@ -65,7 +67,7 @@ def create_data_row(self, dataset: "Entity.Dataset",
         Returns:
             Label with updated references to new data row
         """
-        args = {'row_data': self.add_url_to_data(signer)}
+        args = {'row_data': self.data.create_url(signer)}
         if self.data.external_id is not None:
             args.update({'external_id': self.data.external_id})
 
@@ -151,3 +153,19 @@ def _assign_option(self, classification: ClassificationAnnotation,
             raise TypeError(
                 f"Unexpected type for answer found. {type(classification.value.answer)}"
             )
+
+    @validator("annotations", pre=True)
+    def validate_union(cls, value):
+        supported = tuple([
+            field.type_
+            for field in cls.__fields__['annotations'].sub_fields[0].sub_fields
+        ])
+        if not isinstance(value, list):
+            raise TypeError(f"Annotations must be a list. Found {type(value)}")
+
+        for v in value:
+            if not isinstance(v, supported):
+                raise TypeError(
+                    f"Annotations should be a list containing the following classes : {supported}. Found {type(v)}"
+                )
+        return value

labelbox/data/annotation_types/label.py

@@ -1,4 +1,29 @@
-from pydantic import Field
+from typing import Generic, TypeVar
 from typing_extensions import Annotated
 
+from pydantic import Field
+from pydantic.fields import ModelField
+import numpy as np
+
 Cuid = Annotated[str, Field(min_length=25, max_length=25)]
+
+DType = TypeVar('DType')
+
+
+class TypedArray(np.ndarray, Generic[DType]):
+
+    @classmethod
+    def __get_validators__(cls):
+        yield cls.validate
+
+    @classmethod
+    def validate(cls, val, field: ModelField):
+        if not isinstance(val, np.ndarray):
+            raise TypeError(f"Expected numpy array. Found {type(val)}")
+        dtype_field = field.sub_fields[0]
+        actual_dtype = dtype_field.type_.__args__[0]
+        if val.dtype != actual_dtype:
+            raise TypeError(
+                f"Expected numpy array have type {actual_dtype}. Found {val.dtype}"
+            )
+        return val