Merge pull request #293 from Labelbox/gj/annot-types-polish

Annotation type documentation

Author: Gareth

labelbox/data/annotation_types/annotation.py

@@ -1,47 +1,92 @@
+import abc
 from typing import Any, Dict, List, Union
 
 from .classification import Checklist, Dropdown, Radio, Text
 from .feature import FeatureSchema
-from .geometry import Geometry
+from .geometry import Geometry, Rectangle, Point
 from .ner import TextEntity
 
 
-class BaseAnnotation(FeatureSchema):
+class BaseAnnotation(FeatureSchema, abc.ABC):
     """ Base annotation class. Shouldn't be directly instantiated
     """
     extra: Dict[str, Any] = {}
 
 
 class ClassificationAnnotation(BaseAnnotation):
-    """Class representing classification annotations (annotations that don't have a location) """
+    """Classification annotations (non localized)
+
+    >>> ClassificationAnnotation(
+    >>>     value=Text(answer="my caption message"),
+    >>>     feature_schema_id="my-feature-schema-id"
+    >>> )
+
+    Args:
+        name (Optional[str])
+        feature_schema_id (Optional[Cuid])
+        value (Union[Text, Checklist, Radio, Dropdown])
+        extra (Dict[str, Any])
+     """
 
     value: Union[Text, Checklist, Radio, Dropdown]
 
 
 class ObjectAnnotation(BaseAnnotation):
-    """Class representing objects annotations (non classifications or annotations that have a location)
+    """Generic localized annotation (non classifications)
+
+    >>> ObjectAnnotation(
+    >>>     value=Rectangle(
+    >>>        start=Point(x=0, y=0),
+    >>>        end=Point(x=1, y=1)
+    >>>     ),
+    >>>     feature_schema_id="my-feature-schema-id"
+    >>> )
+
+    Args:
+        name (Optional[str])
+        feature_schema_id (Optional[Cuid])
+        value (Union[TextEntity, Geometry]): Localization of the annotation
+        classifications (Optional[List[ClassificationAnnotation]]): Optional sub classification of the annotation
+        extra (Dict[str, Any])
     """
     value: Union[TextEntity, Geometry]
     classifications: List[ClassificationAnnotation] = []
 
 
 class VideoObjectAnnotation(ObjectAnnotation):
-    """
-    Class for video objects annotations
+    """Video object annotation
+
+    >>> VideoObjectAnnotation(
+    >>>     keyframe=True,
+    >>>     frame=10,
+    >>>     value=Rectangle(
+    >>>        start=Point(x=0, y=0),
+    >>>        end=Point(x=1, y=1)
+    >>>     ),
+    >>>     feature_schema_id="my-feature-schema-id"
+    >>>)
 
     Args:
-        frame: The frame index that this annotation corresponds to
-        keyframe: Whether or not this annotation was a human generated or interpolated annotation
+        name (Optional[str])
+        feature_schema_id (Optional[Cuid])
+        value (Geometry)
+        frame (Int): The frame index that this annotation corresponds to
+        keyframe (bool): Whether or not this annotation was a human generated or interpolated annotation
+        classifications (List[ClassificationAnnotation]) = []
+        extra (Dict[str, Any])
     """
     frame: int
     keyframe: bool
 
 
 class VideoClassificationAnnotation(ClassificationAnnotation):
-    """
-    Class for video classification annotations
+    """Video classification
 
     Args:
-        frame: The frame index that this annotation corresponds to
+        name (Optional[str])
+        feature_schema_id (Optional[Cuid])
+        value (Union[Text, Checklist, Radio, Dropdown])
+        frame (int): The frame index that this annotation corresponds to
+        extra (Dict[str, Any])
     """
     frame: int

labelbox/data/annotation_types/data/raster.py

@@ -1,34 +1,53 @@
-from typing import Callable, Optional
-from io import BytesIO
 from abc import ABC
-
+from io import BytesIO
+from typing import Callable, Optional, Union
+from typing_extensions import Literal
 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 .base_data import BaseData
 from ..types import TypedArray
 
 
 class RasterData(BaseModel, ABC):
-    """
-    Represents an image or segmentation mask.
+    """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):
+    def from_2D_arr(cls, arr: Union[TypedArray[Literal['uint8']],
+                                    TypedArray[Literal['int']]], **kwargs):
+        """Construct from a 2D numpy array
+
+        Args:
+            arr: uint8 compatible numpy array
+
+        Returns:
+            RasterData
+        """
+
         if len(arr.shape) != 2:
             raise ValueError(
-                f"Found array with shape {arr.shape}. Expected two dimensions ([W,H])"
+                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)
 
@@ -153,10 +172,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
@@ -166,7 +185,25 @@ class Config:
 
 
 class MaskData(RasterData):
-    ...
+    """Used to represent a segmentation Mask
+
+    All segments within a mask must be mutually exclusive. At a
+    single cell, only one class can be present. All Mask data is
+    converted to a [H,W,3] image. Classes are
+
+    >>> # 3x3 mask with two classes and back ground
+    >>> MaskData.from_2D_arr([
+    >>>    [0, 0, 0],
+    >>>    [1, 1, 1],
+    >>>    [2, 2, 2],
+    >>>])
+
+    Args:
+        im_bytes: Optional[bytes] = None
+        file_path: Optional[str] = None
+        url: Optional[str] = None
+        arr: Optional[TypedArray[Literal['uint8']]] = None
+    """
 
 
 class ImageData(RasterData, BaseData):

labelbox/data/annotation_types/geometry/geometry.py

@@ -8,8 +8,7 @@
 
 
 class Geometry(BaseModel, ABC):
-    """
-    Base class for geometry objects.
+    """Abstract base class for geometry objects
     """
     extra: Dict[str, Any] = {}
 

labelbox/data/annotation_types/geometry/line.py

@@ -9,6 +9,12 @@
 
 
 class Line(Geometry):
+    """Line annotation
+
+    Args:
+        points (List[Point]): A list of `Point` geometries
+
+    """
     points: List[Point]
 
     @property

labelbox/data/annotation_types/geometry/mask.py

@@ -11,9 +11,30 @@
 
 
 class Mask(Geometry):
-    # Mask data can be shared across multiple masks
+    """Mask used to represent a single class in a larger segmentation mask
+
+    Example of a mutually exclusive class
+
+    >>> arr = MaskData.from_2D_arr([
+    >>>    [0, 0, 0],
+    >>>    [1, 1, 1],
+    >>>    [2, 2, 2],
+    >>>])
+    >>> annotations = [
+    >>>    ObjectAnnotation(value=Mask(mask=arr, color=1), name="dog"),
+    >>>    ObjectAnnotation(value=Mask(mask=arr, color=2), name="cat"),
+    >>>]
+
+    Args:
+         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
+            indicating the values of the class in the `MaskData`
+    """
+
     mask: MaskData
-    color: Tuple[int, int, int]
+    color: Union[Tuple[int, int, int], int]
 
     @property
     def geometry(self):
@@ -31,8 +52,7 @@ def draw(self,
              canvas: Optional[np.ndarray] = None,
              color: Optional[Union[int, Tuple[int, int, int]]] = None,
              thickness=None) -> np.ndarray:
-        """
-        Converts the Mask object into a numpy array
+        """Converts the Mask object into a numpy array
 
         Args:
             height (int): Optionally resize mask height before drawing.
@@ -43,6 +63,7 @@ def draw(self,
                 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
@@ -79,8 +100,9 @@ def create_url(self, signer: Callable[[bytes], str]) -> str:
 
     @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) == 1:
+                color = [color[0]] * 3
             if len(color) != 3:
                 raise ValueError(
                     "Segmentation colors must be either a (r,g,b) tuple or a single grayscale value"

labelbox/data/annotation_types/geometry/point.py

@@ -8,12 +8,21 @@
 
 
 class Point(Geometry):
+    """Point geometry
+
+    >>> Point(x=0, y=0)
+
+    Args:
+        x (float)
+        y (float)
+
+    """
     x: float
     y: float
 
     @property
     def geometry(self) -> geojson.Point:
-        return geojson.Point([self.x, self.y])
+        return geojson.Point((self.x, self.y))
 
     def draw(self,
              height: Optional[int] = None,

labelbox/data/annotation_types/geometry/polygon.py

@@ -1,22 +1,34 @@
 from typing import List, Optional, Union, Tuple
 
-import numpy as np
-import geojson
 import cv2
+import geojson
+import numpy as np
 from pydantic import validator
 
-from .point import Point
 from .geometry import Geometry
+from .point import Point
 
 
 class Polygon(Geometry):
+    """Polygon geometry
+
+    A polygon is created from a collection of points
+
+    >>> Polygon(points=[Point(x=0, y=0), Point(x=1, y=0), Point(x=1, y=1), Point(x=0, y=0)])
+
+    Args:
+        points (List[Point]): List of `Points`, minimum of three points. If you do not
+            close the polygon (the last point and first point are the same) an additional
+            point is added to close it.
+
+    """
     points: List[Point]
 
     @property
-    def geometry(self) -> geojson.MultiPolygon:
+    def geometry(self) -> geojson.Polygon:
         if self.points[0] != self.points[-1]:
             self.points.append(self.points[0])
-        return geojson.Polygon([[[point.x, point.y] for point in self.points]])
+        return geojson.Polygon([[(point.x, point.y) for point in self.points]])
 
     def draw(self,
              height: Optional[int] = None,
@@ -48,4 +60,7 @@ def is_geom_valid(cls, points):
             raise ValueError(
                 f"A polygon must have at least 3 points to be valid. Found {points}"
             )
+        if points[0] != points[-1]:
+            points.append(points[0])
+
         return points

labelbox/data/annotation_types/geometry/rectangle.py

@@ -9,11 +9,13 @@
 
 
 class Rectangle(Geometry):
-    """
-    Represents a 2d rectangle. Also known as a bounding box.
+    """Represents a 2d rectangle. Also known as a bounding box
+
+    >>> Rectangle(start=Point(x=0, y=0), end=Point(x=1, y=1))
 
-    start: Top left coordinate of the rectangle
-    end: Bottom right coordinate of the rectangle
+    Args:
+        start (Point): Top left coordinate of the rectangle
+        end (Point): Bottom right coordinate of the rectangle
     """
     start: Point
     end: Point
@@ -51,3 +53,8 @@ def draw(self,
         if thickness == -1:
             return cv2.fillPoly(canvas, pts, color)
         return cv2.polylines(canvas, pts, True, color, thickness)
+
+    @classmethod
+    def from_xyhw(cls, x: float, y: float, h: float, w: float):
+        """Create Rectangle from x,y, height width format"""
+        return cls(start=Point(x=x, y=y), end=Point(x=x + w, y=y + h))