Merge pull request #540 from Labelbox/jt/al-1921

[AL-1921] Video Bounding Box NDJson Converter

Author: hydak

labelbox/data/serialization/ndjson/label.py

@@ -2,10 +2,11 @@
 from operator import itemgetter
 from typing import Dict, Generator, List, Tuple, Union
 from collections import defaultdict
+import warnings
 
 from pydantic import BaseModel
 
-from ...annotation_types.annotation import ClassificationAnnotation, ObjectAnnotation, VideoClassificationAnnotation
+from ...annotation_types.annotation import ClassificationAnnotation, ObjectAnnotation, VideoClassificationAnnotation, VideoObjectAnnotation
 from ...annotation_types.collection import LabelCollection, LabelGenerator
 from ...annotation_types.data import ImageData, TextData, VideoData
 from ...annotation_types.label import Label
@@ -15,12 +16,13 @@
 
 from .metric import NDScalarMetric, NDMetricAnnotation, NDConfusionMatrixMetric
 from .classification import NDChecklistSubclass, NDClassification, NDClassificationType, NDRadioSubclass
-from .objects import NDObject, NDObjectType
+from .objects import NDObject, NDObjectType, NDSegments
 
 
 class NDLabel(BaseModel):
     annotations: List[Union[NDObjectType, NDClassificationType,
-                            NDConfusionMatrixMetric, NDScalarMetric]]
+                            NDConfusionMatrixMetric, NDScalarMetric,
+                            NDSegments]]
 
     def to_common(self) -> LabelGenerator:
         grouped_annotations = defaultdict(list)
@@ -37,15 +39,20 @@ def from_common(cls,
             yield from cls._create_video_annotations(label)
 
     def _generate_annotations(
-        self, grouped_annotations: Dict[str, List[Union[NDObjectType,
-                                                        NDClassificationType,
-                                                        NDConfusionMatrixMetric,
-                                                        NDScalarMetric]]]
+        self,
+        grouped_annotations: Dict[str,
+                                  List[Union[NDObjectType, NDClassificationType,
+                                             NDConfusionMatrixMetric,
+                                             NDScalarMetric, NDSegments]]]
     ) -> Generator[Label, None, None]:
         for data_row_id, annotations in grouped_annotations.items():
             annots = []
             for annotation in annotations:
-                if isinstance(annotation, NDObjectType.__args__):
+                if isinstance(annotation, NDSegments):
+                    annots.extend(
+                        NDSegments.to_common(annotation, annotation.schema_id))
+
+                elif isinstance(annotation, NDObjectType.__args__):
                     annots.append(NDObject.to_common(annotation))
                 elif isinstance(annotation, NDClassificationType.__args__):
                     annots.extend(NDClassification.to_common(annotation))
@@ -55,7 +62,6 @@ def _generate_annotations(
                 else:
                     raise TypeError(
                         f"Unsupported annotation. {type(annotation)}")
-
             data = self._infer_media_type(annotations)(uid=data_row_id)
             yield Label(annotations=annots, data=data)
 
@@ -65,7 +71,7 @@ def _infer_media_type(
         types = {type(annotation) for annotation in annotations}
         if TextEntity in types:
             return TextData
-        elif VideoClassificationAnnotation in types:
+        elif VideoClassificationAnnotation in types or VideoObjectAnnotation in types:
             return VideoData
         else:
             return ImageData
@@ -83,26 +89,46 @@ def _get_consecutive_frames(
     def _create_video_annotations(
         cls, label: Label
     ) -> Generator[Union[NDChecklistSubclass, NDRadioSubclass], None, None]:
+
         video_annotations = defaultdict(list)
         for annot in label.annotations:
-            if isinstance(annot, VideoClassificationAnnotation):
+            if isinstance(
+                    annot,
+                (VideoClassificationAnnotation, VideoObjectAnnotation)):
                 video_annotations[annot.feature_schema_id].append(annot)
 
         for annotation_group in video_annotations.values():
             consecutive_frames = cls._get_consecutive_frames(
                 sorted([annotation.frame for annotation in annotation_group]))
-            annotation = annotation_group[0]
-            frames_data = []
-            for frames in consecutive_frames:
-                frames_data.append({'start': frames[0], 'end': frames[-1]})
-            annotation.extra.update({'frames': frames_data})
-            yield NDClassification.from_common(annotation, label.data)
+
+            if isinstance(annotation_group[0], VideoClassificationAnnotation):
+                annotation = annotation_group[0]
+                frames_data = []
+                for frames in consecutive_frames:
+                    frames_data.append({'start': frames[0], 'end': frames[-1]})
+                annotation.extra.update({'frames': frames_data})
+                yield NDClassification.from_common(annotation, label.data)
+
+            elif isinstance(annotation_group[0], VideoObjectAnnotation):
+                warnings.warn(
+                    """Nested classifications are not currently supported
+                    for video object annotations
+                    and will not import alongside the object annotations.""")
+                segments = []
+                for start_frame, end_frame in consecutive_frames:
+                    segment = []
+                    for annotation in annotation_group:
+                        if annotation.keyframe and start_frame <= annotation.frame <= end_frame:
+                            segment.append(annotation)
+                    segments.append(segment)
+                yield NDObject.from_common(segments, label.data)
 
     @classmethod
     def _create_non_video_annotations(cls, label: Label):
         non_video_annotations = [
             annot for annot in label.annotations
-            if not isinstance(annot, VideoClassificationAnnotation)
+            if not isinstance(annot, (VideoClassificationAnnotation,
+                                      VideoObjectAnnotation))
         ]
         for annotation in non_video_annotations:
             if isinstance(annotation, ClassificationAnnotation):

labelbox/data/serialization/ndjson/objects.py

@@ -7,11 +7,13 @@
 from pydantic import BaseModel
 from PIL import Image
 
+from labelbox.data.annotation_types.data.video import VideoData
+
 from ...annotation_types.data import ImageData, TextData, MaskData
 from ...annotation_types.ner import TextEntity
 from ...annotation_types.types import Cuid
 from ...annotation_types.geometry import Rectangle, Polygon, Line, Point, Mask
-from ...annotation_types.annotation import ClassificationAnnotation, ObjectAnnotation
+from ...annotation_types.annotation import ClassificationAnnotation, ObjectAnnotation, VideoObjectAnnotation
 from .classification import NDSubclassification, NDSubclassificationType
 from .base import DataRow, NDAnnotation
 
@@ -20,6 +22,11 @@ class NDBaseObject(NDAnnotation):
     classifications: List[NDSubclassificationType] = []
 
 
+class VideoSupported(BaseModel):
+    #support for video for objects are per-frame basis
+    frame: int
+
+
 class _Point(BaseModel):
     x: float
     y: float
@@ -118,6 +125,75 @@ def from_common(cls, rectangle: Rectangle,
                    classifications=classifications)
 
 
+class NDFrameRectangle(VideoSupported):
+    bbox: Bbox
+
+    def to_common(self, feature_schema_id: Cuid) -> VideoObjectAnnotation:
+        return VideoObjectAnnotation(
+            frame=self.frame,
+            keyframe=True,
+            feature_schema_id=feature_schema_id,
+            value=Rectangle(start=Point(x=self.bbox.left, y=self.bbox.top),
+                            end=Point(x=self.bbox.left + self.bbox.width,
+                                      y=self.bbox.top + self.bbox.height)))
+
+    @classmethod
+    def from_common(cls, frame: int, rectangle: Rectangle):
+        return cls(frame=frame,
+                   bbox=Bbox(top=rectangle.start.y,
+                             left=rectangle.start.x,
+                             height=rectangle.end.y - rectangle.start.y,
+                             width=rectangle.end.x - rectangle.start.x))
+
+
+class NDSegment(BaseModel):
+    keyframes: List[NDFrameRectangle]
+
+    @staticmethod
+    def lookup_segment_object_type(segment: List) -> "NDFrameObjectType":
+        """Used for determining which object type the annotation contains
+        returns the object type"""
+        result = {Rectangle: NDFrameRectangle}.get(type(segment[0].value))
+        return result
+
+    def to_common(self, feature_schema_id: Cuid):
+        return [
+            keyframe.to_common(feature_schema_id) for keyframe in self.keyframes
+        ]
+
+    @classmethod
+    def from_common(cls, segment):
+        nd_frame_object_type = cls.lookup_segment_object_type(segment)
+
+        return cls(keyframes=[
+            nd_frame_object_type.from_common(object_annotation.frame,
+                                             object_annotation.value)
+            for object_annotation in segment
+        ])
+
+
+class NDSegments(NDBaseObject):
+    segments: List[NDSegment]
+
+    def to_common(self, feature_schema_id: Cuid):
+        result = []
+        for segment in self.segments:
+            result.extend(NDSegment.to_common(segment, feature_schema_id))
+        return result
+
+    @classmethod
+    def from_common(cls, segments: List[VideoObjectAnnotation], data: VideoData,
+                    feature_schema_id: Cuid, extra: Dict[str,
+                                                         Any]) -> "NDSegments":
+
+        segments = [NDSegment.from_common(segment) for segment in segments]
+
+        return cls(segments=segments,
+                   dataRow=DataRow(id=data.uid),
+                   schema_id=feature_schema_id,
+                   uuid=extra.get('uuid'))
+
+
 class _URIMask(BaseModel):
     instanceURI: str
     colorRGB: Tuple[int, int, int]
@@ -208,9 +284,20 @@ def to_common(annotation: "NDObjectType") -> ObjectAnnotation:
 
     @classmethod
     def from_common(
-        cls, annotation: ObjectAnnotation, data: Union[ImageData, TextData]
+        cls, annotation: Union[ObjectAnnotation,
+                               List[List[VideoObjectAnnotation]]],
+        data: Union[ImageData, TextData]
     ) -> Union[NDLine, NDPoint, NDPolygon, NDRectangle, NDMask, NDTextEntity]:
         obj = cls.lookup_object(annotation)
+
+        #if it is video segments
+        if (obj == NDSegments):
+            return obj.from_common(
+                annotation,
+                data,
+                feature_schema_id=annotation[0][0].feature_schema_id,
+                extra=annotation[0][0].extra)
+
         subclasses = [
             NDSubclassification.from_common(annot)
             for annot in annotation.classifications
@@ -220,15 +307,19 @@ def from_common(
                                data)
 
     @staticmethod
-    def lookup_object(annotation: ObjectAnnotation) -> "NDObjectType":
-        result = {
-            Line: NDLine,
-            Point: NDPoint,
-            Polygon: NDPolygon,
-            Rectangle: NDRectangle,
-            Mask: NDMask,
-            TextEntity: NDTextEntity
-        }.get(type(annotation.value))
+    def lookup_object(
+            annotation: Union[ObjectAnnotation, List]) -> "NDObjectType":
+        if isinstance(annotation, list):
+            result = NDSegments
+        else:
+            result = {
+                Line: NDLine,
+                Point: NDPoint,
+                Polygon: NDPolygon,
+                Rectangle: NDRectangle,
+                Mask: NDMask,
+                TextEntity: NDTextEntity
+            }.get(type(annotation.value))
         if result is None:
             raise TypeError(
                 f"Unable to convert object to MAL format. `{type(annotation.value)}`"
@@ -238,3 +329,5 @@ def lookup_object(annotation: ObjectAnnotation) -> "NDObjectType":
 
 NDObjectType = Union[NDLine, NDPolygon, NDPoint, NDRectangle, NDMask,
                      NDTextEntity]
+
+NDFrameObjectType = NDFrameRectangle

labelbox/schema/dataset.py

@@ -314,9 +314,7 @@ def convert_item(item):
             items = [future.result() for future in as_completed(futures)]
         # Prepare and upload the desciptor file
         data = json.dumps(items)
-        return self.client.upload_data(data,
-                                       content_type="application/json",
-                                       filename="json_import.json")
+        return self.client.upload_data(data)
 
     def data_rows_for_external_id(self,
                                   external_id,