Add LineAnnotation and LinePrediction support (#224)

* Add LineAnnotation and LinePrediction support

* Fix top level imports

* Update test case to check right thing

* Add Line types to utils.py

Author: flubstep

nucleus/__init__.py

@@ -15,7 +15,8 @@
     "DatasetItemRetrievalError",
     "Frame",
     "LidarScene",
-    "VideoScene",
+    "LineAnnotation",
+    "LinePrediction",
     "Model",
     "ModelCreationError",
     # "MultiCategoryAnnotation", # coming soon!
@@ -31,6 +32,7 @@
     "SegmentationAnnotation",
     "SegmentationPrediction",
     "Slice",
+    "VideoScene",
 ]
 
 import os
@@ -49,6 +51,7 @@
     BoxAnnotation,
     CategoryAnnotation,
     CuboidAnnotation,
+    LineAnnotation,
     MultiCategoryAnnotation,
     Point,
     Point3D,
@@ -119,6 +122,7 @@
     BoxPrediction,
     CategoryPrediction,
     CuboidPrediction,
+    LinePrediction,
     PolygonPrediction,
     SegmentationPrediction,
 )

nucleus/annotation.py

@@ -1,7 +1,7 @@
 import json
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Dict, List, Optional, Sequence, Union
+from typing import Dict, List, Optional, Sequence, Type, Union
 from urllib.parse import urlparse
 
 from .constants import (
@@ -16,6 +16,7 @@
     INDEX_KEY,
     LABEL_KEY,
     LABELS_KEY,
+    LINE_TYPE,
     MASK_TYPE,
     MASK_URL_KEY,
     METADATA_KEY,
@@ -46,18 +47,17 @@ class Annotation:
     @classmethod
     def from_json(cls, payload: dict):
         """Instantiates annotation object from schematized JSON dict payload."""
-        if payload.get(TYPE_KEY, None) == BOX_TYPE:
-            return BoxAnnotation.from_json(payload)
-        elif payload.get(TYPE_KEY, None) == POLYGON_TYPE:
-            return PolygonAnnotation.from_json(payload)
-        elif payload.get(TYPE_KEY, None) == CUBOID_TYPE:
-            return CuboidAnnotation.from_json(payload)
-        elif payload.get(TYPE_KEY, None) == CATEGORY_TYPE:
-            return CategoryAnnotation.from_json(payload)
-        elif payload.get(TYPE_KEY, None) == MULTICATEGORY_TYPE:
-            return MultiCategoryAnnotation.from_json(payload)
-        else:
-            return SegmentationAnnotation.from_json(payload)
+        type_key_to_type: Dict[str, Type[Annotation]] = {
+            BOX_TYPE: BoxAnnotation,
+            LINE_TYPE: LineAnnotation,
+            POLYGON_TYPE: PolygonAnnotation,
+            CUBOID_TYPE: CuboidAnnotation,
+            CATEGORY_TYPE: CategoryAnnotation,
+            MULTICATEGORY_TYPE: MultiCategoryAnnotation,
+        }
+        type_key = payload.get(TYPE_KEY, None)
+        AnnotationCls = type_key_to_type.get(type_key, SegmentationAnnotation)
+        return AnnotationCls.from_json(payload)
 
     def to_payload(self) -> dict:
         """Serializes annotation object to schematized JSON dict."""
@@ -177,6 +177,88 @@ def to_payload(self) -> dict:
         return {X_KEY: self.x, Y_KEY: self.y}
 
 
+@dataclass
+class LineAnnotation(Annotation):
+    """A polyline annotation consisting of an ordered list of 2D points.
+    A LineAnnotation differs from a PolygonAnnotation by not forming a closed
+    loop, and by having zero area.
+
+    ::
+
+        from nucleus import LineAnnotation
+
+        line = LineAnnotation(
+            label="face",
+            vertices=[Point(100, 100), Point(200, 300), Point(300, 200)],
+            reference_id="person_image_1",
+            annotation_id="person_image_1_line_1",
+            metadata={"camera_mode": "portrait"},
+        )
+
+    Parameters:
+        label (str): The label for this annotation.
+        vertices (List[:class:`Point`]): The list of points making up the line.
+        reference_id (str): User-defined ID of the image to which to apply this
+            annotation.
+        annotation_id (Optional[str]): The annotation ID that uniquely identifies
+            this annotation within its target dataset item. Upon ingest, a matching
+            annotation id will be ignored by default, and updated if update=True
+            for dataset.annotate.
+        metadata (Optional[Dict]): Arbitrary key/value dictionary of info to
+            attach to this annotation.  Strings, floats and ints are supported best
+            by querying and insights features within Nucleus. For more details see
+            our `metadata guide <https://nucleus.scale.com/docs/upload-metadata>`_.
+    """
+
+    label: str
+    vertices: List[Point]
+    reference_id: str
+    annotation_id: Optional[str] = None
+    metadata: Optional[Dict] = None
+
+    def __post_init__(self):
+        self.metadata = self.metadata if self.metadata else {}
+        if len(self.vertices) > 0:
+            if not hasattr(self.vertices[0], X_KEY) or not hasattr(
+                self.vertices[0], "to_payload"
+            ):
+                try:
+                    self.vertices = [
+                        Point(x=vertex[X_KEY], y=vertex[Y_KEY])
+                        for vertex in self.vertices
+                    ]
+                except KeyError as ke:
+                    raise ValueError(
+                        "Use a point object to pass in vertices. For example, vertices=[nucleus.Point(x=1, y=2)]"
+                    ) from ke
+
+    @classmethod
+    def from_json(cls, payload: dict):
+        geometry = payload.get(GEOMETRY_KEY, {})
+        return cls(
+            label=payload.get(LABEL_KEY, 0),
+            vertices=[
+                Point.from_json(_) for _ in geometry.get(VERTICES_KEY, [])
+            ],
+            reference_id=payload[REFERENCE_ID_KEY],
+            annotation_id=payload.get(ANNOTATION_ID_KEY, None),
+            metadata=payload.get(METADATA_KEY, {}),
+        )
+
+    def to_payload(self) -> dict:
+        payload = {
+            LABEL_KEY: self.label,
+            TYPE_KEY: LINE_TYPE,
+            GEOMETRY_KEY: {
+                VERTICES_KEY: [_.to_payload() for _ in self.vertices]
+            },
+            REFERENCE_ID_KEY: self.reference_id,
+            ANNOTATION_ID_KEY: self.annotation_id,
+            METADATA_KEY: self.metadata,
+        }
+        return payload
+
+
 @dataclass
 class PolygonAnnotation(Annotation):
     """A polygon annotation consisting of an ordered list of 2D points.
@@ -499,6 +581,7 @@ def to_payload(self) -> dict:
 
 class AnnotationTypes(Enum):
     BOX = BOX_TYPE
+    LINE = LINE_TYPE
     POLYGON = POLYGON_TYPE
     CUBOID = CUBOID_TYPE
     CATEGORY = CATEGORY_TYPE
@@ -600,6 +683,7 @@ class AnnotationList:
     """Wrapper class separating a list of annotations by type."""
 
     box_annotations: List[BoxAnnotation] = field(default_factory=list)
+    line_annotations: List[LineAnnotation] = field(default_factory=list)
     polygon_annotations: List[PolygonAnnotation] = field(default_factory=list)
     cuboid_annotations: List[CuboidAnnotation] = field(default_factory=list)
     category_annotations: List[CategoryAnnotation] = field(
@@ -620,6 +704,8 @@ def add_annotations(self, annotations: List[Annotation]):
 
             if isinstance(annotation, BoxAnnotation):
                 self.box_annotations.append(annotation)
+            elif isinstance(annotation, LineAnnotation):
+                self.line_annotations.append(annotation)
             elif isinstance(annotation, PolygonAnnotation):
                 self.polygon_annotations.append(annotation)
             elif isinstance(annotation, CuboidAnnotation):
@@ -637,6 +723,7 @@ def add_annotations(self, annotations: List[Annotation]):
     def __len__(self):
         return (
             len(self.box_annotations)
+            + len(self.line_annotations)
             + len(self.polygon_annotations)
             + len(self.cuboid_annotations)
             + len(self.category_annotations)

nucleus/constants.py

@@ -5,6 +5,7 @@
 ANNOTATION_METADATA_SCHEMA_KEY = "annotation_metadata_schema"
 BACKFILL_JOB_KEY = "backfill_job"
 BOX_TYPE = "box"
+LINE_TYPE = "line"
 POLYGON_TYPE = "polygon"
 MASK_TYPE = "mask"
 SEGMENTATION_TYPE = "segmentation"
@@ -13,6 +14,7 @@
 MULTICATEGORY_TYPE = "multicategory"
 ANNOTATION_TYPES = (
     BOX_TYPE,
+    LINE_TYPE,
     POLYGON_TYPE,
     SEGMENTATION_TYPE,
     CUBOID_TYPE,

nucleus/prediction.py

@@ -4,12 +4,13 @@
 such as confidence or probability distributions.
 """
 from dataclasses import dataclass, field
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional, Type, Union
 
 from .annotation import (
     BoxAnnotation,
     CategoryAnnotation,
     CuboidAnnotation,
+    LineAnnotation,
     Point,
     Point3D,
     PolygonAnnotation,
@@ -28,6 +29,7 @@
     GEOMETRY_KEY,
     HEIGHT_KEY,
     LABEL_KEY,
+    LINE_TYPE,
     MASK_URL_KEY,
     METADATA_KEY,
     POLYGON_TYPE,
@@ -45,16 +47,16 @@
 
 def from_json(payload: dict):
     """Instantiates prediction object from schematized JSON dict payload."""
-    if payload.get(TYPE_KEY, None) == BOX_TYPE:
-        return BoxPrediction.from_json(payload)
-    elif payload.get(TYPE_KEY, None) == POLYGON_TYPE:
-        return PolygonPrediction.from_json(payload)
-    elif payload.get(TYPE_KEY, None) == CUBOID_TYPE:
-        return CuboidPrediction.from_json(payload)
-    elif payload.get(TYPE_KEY, None) == CATEGORY_TYPE:
-        return CategoryPrediction.from_json(payload)
-    else:
-        return SegmentationPrediction.from_json(payload)
+    type_key_to_type: Dict[str, Type[Prediction]] = {
+        BOX_TYPE: BoxPrediction,
+        LINE_TYPE: LinePrediction,
+        POLYGON_TYPE: PolygonPrediction,
+        CUBOID_TYPE: CuboidPrediction,
+        CATEGORY_TYPE: CategoryPrediction,
+    }
+    type_key = payload.get(TYPE_KEY, None)
+    PredictionCls = type_key_to_type.get(type_key, SegmentationPrediction)
+    return PredictionCls.from_json(payload)
 
 
 class SegmentationPrediction(SegmentationAnnotation):
@@ -203,6 +205,74 @@ def from_json(cls, payload: dict):
         )
 
 
+class LinePrediction(LineAnnotation):
+    """Prediction of a line.
+
+    Parameters:
+        label (str): The label for this prediction (e.g. car, pedestrian, bicycle).
+        vertices List[:class:`Point`]: The list of points making up the line.
+        reference_id (str): User-defined ID of the image to which to apply this
+            annotation.
+        confidence: 0-1 indicating the confidence of the prediction.
+        annotation_id (Optional[str]): The annotation ID that uniquely identifies
+            this annotation within its target dataset item. Upon ingest, a matching
+            annotation id will be ignored by default, and updated if update=True
+            for dataset.annotate.
+        metadata (Optional[Dict]): Arbitrary key/value dictionary of info to
+            attach to this prediction.  Strings, floats and ints are supported best
+            by querying and insights features within Nucleus. For more details see
+            our `metadata guide <https://nucleus.scale.com/docs/upload-metadata>`_.
+        class_pdf: An optional complete class probability distribution on this
+            annotation. Each value should be between 0 and 1 (inclusive), and sum up to
+            1 as a complete distribution. This can be useful for computing entropy to
+            surface places where the model is most uncertain.
+    """
+
+    def __init__(
+        self,
+        label: str,
+        vertices: List[Point],
+        reference_id: str,
+        confidence: Optional[float] = None,
+        annotation_id: Optional[str] = None,
+        metadata: Optional[Dict] = None,
+        class_pdf: Optional[Dict] = None,
+    ):
+        super().__init__(
+            label=label,
+            vertices=vertices,
+            reference_id=reference_id,
+            annotation_id=annotation_id,
+            metadata=metadata,
+        )
+        self.confidence = confidence
+        self.class_pdf = class_pdf
+
+    def to_payload(self) -> dict:
+        payload = super().to_payload()
+        if self.confidence is not None:
+            payload[CONFIDENCE_KEY] = self.confidence
+        if self.class_pdf is not None:
+            payload[CLASS_PDF_KEY] = self.class_pdf
+
+        return payload
+
+    @classmethod
+    def from_json(cls, payload: dict):
+        geometry = payload.get(GEOMETRY_KEY, {})
+        return cls(
+            label=payload.get(LABEL_KEY, 0),
+            vertices=[
+                Point.from_json(_) for _ in geometry.get(VERTICES_KEY, [])
+            ],
+            reference_id=payload[REFERENCE_ID_KEY],
+            confidence=payload.get(CONFIDENCE_KEY, None),
+            annotation_id=payload.get(ANNOTATION_ID_KEY, None),
+            metadata=payload.get(METADATA_KEY, {}),
+            class_pdf=payload.get(CLASS_PDF_KEY, None),
+        )
+
+
 class PolygonPrediction(PolygonAnnotation):
     """Prediction of a polygon.
 
@@ -404,6 +474,7 @@ def from_json(cls, payload: dict):
 
 Prediction = Union[
     BoxPrediction,
+    LinePrediction,
     PolygonPrediction,
     CuboidPrediction,
     CategoryPrediction,
@@ -416,6 +487,7 @@ class PredictionList:
     """Wrapper class separating a list of predictions by type."""
 
     box_predictions: List[BoxPrediction] = field(default_factory=list)
+    line_predictions: List[LinePrediction] = field(default_factory=list)
     polygon_predictions: List[PolygonPrediction] = field(default_factory=list)
     cuboid_predictions: List[CuboidPrediction] = field(default_factory=list)
     category_predictions: List[CategoryPrediction] = field(
@@ -429,6 +501,8 @@ def add_predictions(self, predictions: List[Prediction]):
         for prediction in predictions:
             if isinstance(prediction, BoxPrediction):
                 self.box_predictions.append(prediction)
+            elif isinstance(prediction, LinePrediction):
+                self.line_predictions.append(prediction)
             elif isinstance(prediction, PolygonPrediction):
                 self.polygon_predictions.append(prediction)
             elif isinstance(prediction, CuboidPrediction):
@@ -444,6 +518,7 @@ def add_predictions(self, predictions: List[Prediction]):
     def __len__(self):
         return (
             len(self.box_predictions)
+            + len(self.line_predictions)
             + len(self.polygon_predictions)
             + len(self.cuboid_predictions)
             + len(self.category_predictions)