Resolve comments and handle uint8 conversion

Author: gdj0nes

labelbox/data/annotation_types/data/raster.py

@@ -1,45 +1,51 @@
-from typing import Callable, Optional
-from io import BytesIO
 from abc import ABC
+from io import BytesIO
+from typing import Callable, Optional
 
 import numpy as np
-from pydantic import BaseModel
 import requests
+from PIL import Image
 from google.api_core import retry
-from typing_extensions import Literal
+from pydantic import BaseModel
 from pydantic import root_validator
-from PIL import Image
+from typing_extensions import Literal
 
 from .base_data import BaseData
 from ..types import TypedArray
 
 
 class RasterData(BaseModel, ABC):
     """Represents an image or segmentation mask.
-
     """
     im_bytes: Optional[bytes] = None
     file_path: Optional[str] = None
     url: Optional[str] = None
     arr: Optional[TypedArray[Literal['uint8']]] = None
 
-
     @classmethod
-    def from_2D_arr(cls, arr: TypedArray[Literal['uint8']], **kwargs):
-        """Construct
+    def from_2D_arr(cls, arr: Union[TypedArray[Literal['uint8']], TypedArray[Literal['int']]], **kwargs):
+        """Construct from a 2D numpy array
 
         Args:
-            arr:
-            **kwargs:
+            arr: uint8 compatible numpy array
 
         Returns:
-
+            RasterData
         """
 
         if len(arr.shape) != 2:
             raise ValueError(
                 f"Found array with shape {arr.shape}. Expected two dimensions [H, W]"
             )
+
+        if not np.issubdtype(arr.dtype, np.integer):
+            raise ValueError("Array must be an integer subtype")
+
+        if np.can_cast(arr, np.uint8):
+            arr = arr.astype(np.uint8)
+        else:
+            raise ValueError("Could not cast array to uint8, check that values are between 0 and 255")
+
         arr = np.stack((arr,) * 3, axis=-1)
         return cls(arr=arr, **kwargs)
 
@@ -164,10 +170,10 @@ def validate_args(cls, values):
 
     def __repr__(self) -> str:
         symbol_or_none = lambda data: '...' if data is not None else None
-        return  f"{self.__class__.__name__}(im_bytes={symbol_or_none(self.im_bytes)}," \
-                f"file_path={self.file_path}," \
-                f"url={self.url}," \
-                f"arr={symbol_or_none(self.arr)})"
+        return f"{self.__class__.__name__}(im_bytes={symbol_or_none(self.im_bytes)}," \
+               f"file_path={self.file_path}," \
+               f"url={self.url}," \
+               f"arr={symbol_or_none(self.arr)})"
 
     class Config:
         # Required for sharing references
@@ -198,6 +204,5 @@ class MaskData(RasterData):
     """
 
 
-
 class ImageData(RasterData, BaseData):
     ...

labelbox/data/annotation_types/geometry/mask.py

@@ -23,10 +23,10 @@ class Mask(Geometry):
     >>> annotations = [
     >>>    ObjectAnnotation(value=Mask(mask=arr, color=1), name="dog"),
     >>>    ObjectAnnotation(value=Mask(mask=arr, color=2), name="cat"),
-    >>> ]
+    >>>]
 
     Args:
-         mask (MaskData): A object containing the actual mask, `MaskData` can
+         mask (MaskData): An object containing the actual mask, `MaskData` can
             be shared across multiple `Masks` to more efficiently store data
             for mutually exclusive segmentations.
          color (Tuple[uint8, uint8, uint8]): RGB color or a single value