Merge pull request #204 from Labelbox/ms/annotation-updates

annotation clean up

Author: msokoloff1

labelbox/data/annotation_types/annotation.py

@@ -1,5 +1,7 @@
 from typing import Any, Dict, List, Union
 
+from pydantic.main import BaseModel
+
 from .classification import Checklist, Dropdown, Radio, Text
 from .feature import FeatureSchema
 from .geometry import Geometry
@@ -9,7 +11,6 @@
 class BaseAnnotation(FeatureSchema):
     """ Base annotation class. Shouldn't be directly instantiated
     """
-
     extra: Dict[str, Any] = {}
 
 

labelbox/data/annotation_types/collection.py

@@ -5,8 +5,9 @@
 
 from tqdm import tqdm
 
-from labelbox import OntologyBuilder
+from labelbox.schema import ontology
 from labelbox.orm.model import Entity
+from ..ontology import get_classifications, get_tools
 from ..generator import PrefetchGenerator
 from .label import Label
 
@@ -24,8 +25,8 @@ def __init__(self, data: Iterable[Label]):
         self._data = data
         self._index = 0
 
-    def assign_schema_ids(self,
-                          ontology_builder: OntologyBuilder) -> "LabelList":
+    def assign_schema_ids(
+            self, ontology_builder: "ontology.OntologyBuilder") -> "LabelList":
         """
         Adds schema ids to all FeatureSchema objects in the Labels.
         This is necessary for MAL.
@@ -110,6 +111,16 @@ def add_url_to_data(self, signer, max_concurrency=20) -> "LabelList":
             ...
         return self
 
+    def get_ontology(self) -> ontology.OntologyBuilder:
+        classifications = []
+        tools = []
+        for label in self._data:
+            tools = get_tools(label.object_annotations(), tools)
+            classifications = get_classifications(
+                label.classification_annotations(), classifications)
+        return ontology.OntologyBuilder(tools=tools,
+                                        classifications=classifications)
+
     def _ensure_unique_external_ids(self) -> None:
         external_ids = set()
         for label in self._data:
@@ -122,6 +133,9 @@ def _ensure_unique_external_ids(self) -> None:
                     )
             external_ids.add(label.data.external_id)
 
+    def append(self, label: Label):
+        self._data.append(label)
+
     def __iter__(self) -> "LabelList":
         self._index = 0
         return self
@@ -166,7 +180,8 @@ def as_list(self) -> "LabelList":
         return LabelList(data=list(self))
 
     def assign_schema_ids(
-            self, ontology_builder: OntologyBuilder) -> "LabelGenerator":
+            self,
+            ontology_builder: "ontology.OntologyBuilder") -> "LabelGenerator":
 
         def _assign_ids(label: Label):
             label.assign_schema_ids(ontology_builder)

labelbox/data/annotation_types/data/raster.py

@@ -20,6 +20,15 @@ class RasterData(BaseData):
     url: Optional[str] = None
     arr: Optional[TypedArray[Literal['uint8']]] = None
 
+    @classmethod
+    def from_2D_arr(cls, arr: TypedArray[Literal['uint8']], **kwargs):
+        if len(arr.shape):
+            raise ValueError(
+                f"Found array with shape {arr.shape}. Expected two dimensions ([W,H])"
+            )
+        arr = np.stack((arr,) * 3, axis=-1)
+        return cls(arr=arr, **kwargs)
+
     def bytes_to_np(self, image_bytes: bytes) -> np.ndarray:
         """
         Converts image bytes to a numpy array
@@ -38,9 +47,9 @@ def np_to_bytes(self, arr: np.ndarray) -> bytes:
         Returns:
             png encoded bytes
         """
-        if len(arr.shape) not in [2, 3]:
-            raise ValueError("unsupported image format")
-
+        if len(arr.shape) != 3:
+            raise ValueError("unsupported image format. Must be 3D ([H,W,C])."
+                             "Use RasterData.from_2D_arr to construct from 2D")
         if arr.dtype != np.uint8:
             raise TypeError(f"image data type must be uint8. Found {arr.dtype}")
 
@@ -72,6 +81,9 @@ def data(self) -> np.ndarray:
         else:
             raise ValueError("Must set either url, file_path or im_bytes")
 
+    def set_fetch_fn(self, fn):
+        object.__setattr__(self, 'fetch_remote', lambda: fn(self))
+
     def fetch_remote(self) -> bytes:
         """
         Method for accessing url.
@@ -122,12 +134,19 @@ def validate_args(cls, values):
                 raise TypeError(
                     "Numpy array representing segmentation mask must be np.uint8"
                 )
-            elif len(arr.shape) not in [2, 3]:
-                raise TypeError(
-                    f"Numpy array must have 2 or 3 dims. Found shape {arr.shape}"
-                )
+            elif len(arr.shape) != 3:
+                raise ValueError(
+                    "unsupported image format. Must be 3D ([H,W,C])."
+                    "Use RasterData.from_2D_arr to construct from 2D")
         return values
 
+    def __repr__(self) -> str:
+        symbol_or_none = lambda data: '...' if data is not None else None
+        return  f"RasterData(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
         copy_on_model_validation = False

labelbox/data/annotation_types/data/text.py

@@ -36,6 +36,9 @@ def data(self) -> str:
         else:
             raise ValueError("Must set either url, file_path or im_bytes")
 
+    def set_fetch_fn(self, fn):
+        object.__setattr__(self, 'fetch_remote', lambda: fn(self))
+
     def fetch_remote(self) -> str:
         """
         Method for accessing url.
@@ -79,6 +82,11 @@ def validate_date(cls, values):
                 "One of `file_path`, `text`, `uid`, or `url` required.")
         return values
 
+    def __repr__(self) -> str:
+        return  f"TextData(file_path={self.file_path}," \
+                f"text={self.text[:30] + '...' if self.text is not None else None}," \
+                f"url={self.url})"
+
     class config:
         # Required for discriminating between data types
         extra = 'forbid'

labelbox/data/annotation_types/data/video.py

@@ -86,6 +86,9 @@ def __getitem__(self, idx: int) -> np.ndarray:
             )
         return self.frames[idx]
 
+    def set_fetch_fn(self, fn):
+        object.__setattr__(self, 'fetch_remote', lambda: fn(self))
+
     def fetch_remote(self, local_path) -> None:
         """
         Method for downloading data from self.url
@@ -153,6 +156,11 @@ def validate_data(cls, values):
                 "One of `file_path`, `frames`, `uid`, or `url` required.")
         return values
 
+    def __repr__(self) -> str:
+        return  f"TextData(file_path={self.file_path}," \
+                f"frames={'...' if self.frames is not None else None}," \
+                f"url={self.url})"
+
     class Config:
         # Required for discriminating between data types
         extra = 'forbid'

labelbox/data/annotation_types/geometry/line.py

@@ -20,9 +20,9 @@ def raster(self,
                height: int,
                width: int,
                thickness=1,
-               color=255) -> np.ndarray:
+               color=(255, 255, 255)) -> np.ndarray:
         """
-        Draw the line onto a 2d mask
+        Draw the line onto a 3d mask
 
         Args:
             height (int): height of the mask
@@ -32,8 +32,7 @@ def raster(self,
         Returns:
             numpy array representing the mask with the line drawn on it.
         """
-
-        canvas = np.zeros((height, width), dtype=np.uint8)
+        canvas = np.zeros((height, width, 3), dtype=np.uint8)
         pts = np.array(self.geometry['coordinates']).astype(np.int32)
         return cv2.polylines(canvas,
                              pts,

labelbox/data/annotation_types/geometry/mask.py

@@ -1,9 +1,10 @@
-from typing import Callable, Tuple, Union
+from typing import Callable, Optional, Tuple, Union
 
 import numpy as np
 from pydantic.class_validators import validator
 from rasterio.features import shapes
 from shapely.geometry import MultiPolygon, shape
+import cv2
 
 from ..data.raster import RasterData
 from .geometry import Geometry
@@ -13,7 +14,7 @@ class Mask(Geometry):
     # Raster data can be shared across multiple masks... or not
     mask: RasterData
     # RGB or Grayscale
-    color: Union[int, Tuple[int, int, int]]
+    color: Tuple[int, int, int]
 
     @property
     def geometry(self):
@@ -25,21 +26,26 @@ def geometry(self):
             if val >= 1)
         return MultiPolygon(polygons).__geo_interface__
 
-    def raster(self, binary=False) -> np.ndarray:
+    def raster(self,
+               height: Optional[int] = None,
+               width: Optional[int] = None,
+               binary=False) -> np.ndarray:
         """
         Removes all pixels from the segmentation mask that do not equal self.color
 
+        Args:
+            height:
+
         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 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
-        elif isinstance(self.color, int):
-            return mask * self.color
         else:
             color_image = np.zeros((mask.shape[0], mask.shape[1], 3),
                                    dtype=np.uint8)

labelbox/data/annotation_types/geometry/point.py

@@ -17,9 +17,9 @@ def raster(self,
                height: int,
                width: int,
                thickness: int = 1,
-               color=255) -> np.ndarray:
+               color=(255, 255, 255)) -> np.ndarray:
         """
-        Draw the point onto a 2d mask
+        Draw the point onto a 3d mask
 
         Args:
             height (int): height of the mask
@@ -29,7 +29,7 @@ def raster(self,
         Returns:
             numpy array representing the mask with the point drawn on it.
         """
-        canvas = np.zeros((height, width), dtype=np.uint8)
+        canvas = np.zeros((height, width, 3), dtype=np.uint8)
         return cv2.circle(canvas, (int(self.x), int(self.y)),
                           radius=thickness,
                           color=color,

labelbox/data/annotation_types/geometry/polygon.py

@@ -18,9 +18,10 @@ def geometry(self) -> geojson.MultiPolygon:
             self.points.append(self.points[0])
         return geojson.Polygon([[[point.x, point.y] for point in self.points]])
 
-    def raster(self, height: int, width: int, color=255) -> np.ndarray:
+    def raster(self, height: int, width: int,
+               color=(255, 255, 255)) -> np.ndarray:
         """
-        Draw the polygon onto a 2d mask
+        Draw the polygon onto a 3d mask
 
         Args:
             height (int): height of the mask
@@ -29,7 +30,7 @@ def raster(self, height: int, width: int, color=255) -> np.ndarray:
         Returns:
             numpy array representing the mask with the polygon drawn on it.
         """
-        canvas = np.zeros((height, width), dtype=np.uint8)
+        canvas = np.zeros((height, width, 3), dtype=np.uint8)
         pts = np.array(self.geometry['coordinates']).astype(np.int32)
         return cv2.fillPoly(canvas, pts=pts, color=color)
 

labelbox/data/annotation_types/geometry/rectangle.py

@@ -26,9 +26,10 @@ def geometry(self) -> geojson.geometry.Geometry:
             [self.start.x, self.start.y],
         ]])
 
-    def raster(self, height: int, width: int, color: int = 255) -> np.ndarray:
+    def raster(self, height: int, width: int,
+               color=(255, 255, 255)) -> np.ndarray:
         """
-        Draw the rectangle onto a 2d mask
+        Draw the rectangle onto a 3d mask
 
         Args:
             height (int): height of the mask
@@ -37,6 +38,6 @@ def raster(self, height: int, width: int, color: int = 255) -> np.ndarray:
         Returns:
             numpy array representing the mask with the rectangle drawn on it.
         """
-        canvas = np.zeros((height, width), dtype=np.uint8)
+        canvas = np.zeros((height, width, 3), dtype=np.uint8)
         pts = np.array(self.geometry['coordinates']).astype(np.int32)
         return cv2.fillPoly(canvas, pts=pts, color=color)