ability to run to_common to convert back to common labelbox format

Author: jtsodapop

labelbox/data/serialization/ndjson/label.py

@@ -15,41 +15,48 @@
 
 from .metric import NDScalarMetric, NDMetricAnnotation, NDConfusionMatrixMetric
 from .classification import NDChecklistSubclass, NDClassification, NDClassificationType, NDRadioSubclass
-from .objects import NDObject, NDObjectType
+from .objects import NDObject, NDObjectType, NDFrameObjectType, NDSegments
 
 
 class NDLabel(BaseModel):
     annotations: List[Union[NDObjectType, NDClassificationType,
-                            NDConfusionMatrixMetric, NDScalarMetric]]
+                            NDConfusionMatrixMetric, NDScalarMetric,
+                            NDSegments]]
 
     def to_common(self) -> LabelGenerator:
         grouped_annotations = defaultdict(list)
         for annotation in self.annotations:
             grouped_annotations[annotation.data_row.id].append(annotation)
         return LabelGenerator(
             data=self._generate_annotations(grouped_annotations))
+        # c = LabelGenerator(data=self._generate_annotations(grouped_annotations))
+        # print(next(c), "Wefwefweefw\n")
 
     @classmethod
     def from_common(cls,
                     data: LabelCollection) -> Generator["NDLabel", None, None]:
         for label in data:
-            # print(f"label in data:{label}")
             yield from cls._create_non_video_annotations(label)
             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, NDSegments):
+                    # b = NDSegments.to_common(annotation, annotation.schema_id)
+                    # print(type(b), b)
+                    annots.extend(
+                        NDSegments.to_common(annotation, annotation.schema_id))
 
-                if isinstance(annotation, NDObjectType.__args__):
+                elif isinstance(annotation, NDObjectType.__args__):
                     annots.append(NDObject.to_common(annotation))
-                #TODO: have a check on if the return type needs to extend or not
                 elif isinstance(annotation, NDClassificationType.__args__):
                     annots.extend(NDClassification.to_common(annotation))
                 elif isinstance(annotation,
@@ -58,7 +65,7 @@ def _generate_annotations(
                 else:
                     raise TypeError(
                         f"Unsupported annotation. {type(annotation)}")
-
+            # print(type(annots[0]), annots[0])
             data = self._infer_media_type(annotations)(uid=data_row_id)
             yield Label(annotations=annots, data=data)
 
@@ -115,31 +122,7 @@ def _create_video_annotations(
                         if annotation.keyframe and start_frame <= annotation.frame <= end_frame:
                             segment.append(annotation)
                     segments.append(segment)
-
-                print(segments[0], "\n")
-                print(segments[1], "\n")
-                print(consecutive_frames)
                 yield NDObject.from_common(segments, label.data)
-                # segments = []
-                # seg_frames = []
-                # for cframes in consecutive_frames:
-                #     seg_frames.append(cframes[0])
-                #     seg_frames.append(cframes[1])
-                # print(seg_frames)
-                # for annotation in annotation_group:
-                #     if annotation.frame in seg_frames:
-                #         segments.append(annotation)
-                #     # if annotation.keyframe:
-                # #         segments.append(annotation)
-                # # print(consecutive_frames)
-                # #TODO: current issue is that the way the code is written doesn't account for
-                # #which frames are consecutive. maybe we should just have list of segments
-                # annotations = []
-                # for annotation in segments:
-                #     annotations.append(
-                #         NDObject.from_common(annotation, label.data))
-                # yield annotations[0]
-                # yield {}
 
     @classmethod
     def _create_non_video_annotations(cls, label: Label):

labelbox/data/serialization/ndjson/objects.py

@@ -129,6 +129,15 @@ def from_common(cls, rectangle: Rectangle,
 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,
@@ -143,9 +152,17 @@ class NDSegment(BaseModel):
 
     @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):
+        # print(f"\n to_common in ndsegment {self.keyframes}\n")
+        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)
@@ -160,8 +177,11 @@ def from_common(cls, segment):
 class NDSegments(NDBaseObject):
     segments: List[NDSegment]
 
-    def to_common(self):
-        pass
+    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,
@@ -170,16 +190,10 @@ def from_common(cls, segments: List[VideoObjectAnnotation], data: VideoData,
 
         segments = [NDSegment.from_common(segment) for segment in segments]
 
-        # print("\nbefore class instantiation\n", segments,
-        #       "\nbefore class instantiation\n")
-
-        a = cls(segments=segments,
-                dataRow=DataRow(id=data.uid),
-                schema_id=feature_schema_id,
-                uuid=extra.get('uuid'))
-        # print("\nI am from_common  ndsegments", a,
-        #   "\nI am from_common  ndsegments\n")
-        return a
+        return cls(segments=segments,
+                   dataRow=DataRow(id=data.uid),
+                   schema_id=feature_schema_id,
+                   uuid=extra.get('uuid'))
 
 
 class _URIMask(BaseModel):
@@ -278,11 +292,8 @@ def from_common(
 
         #if it is video segments
         if (obj == NDSegments):
-            print("hello i am ndsegment")
-            #look into segment of segments
-            #look into annotations of segment
-            #check and validate that there are no subclasses
-            #new method for processing segments ?
+            #TODO: need to process segments' classifications
+            #unsure if here or will be elsewhere
 
             return obj.from_common(
                 annotation,
@@ -317,17 +328,8 @@ def lookup_object(annotation: ObjectAnnotation) -> "NDObjectType":
             )
         return result
 
-    # @staticmethod
-    # def process_segments(segments: List[List[VideoObjectAnnotation]]):
-    #     for segment in segments:
-    #         for annotation in segment:
-    #             subclasses = [
-    #                 NDSubclassification.from_common(annot)
-    #                 for annot in annotation.classifications
-    #             ]
-
 
 NDObjectType = Union[NDLine, NDPolygon, NDPoint, NDRectangle, NDMask,
                      NDTextEntity]
 
-NDFrameObjectType = NDFrameRectangle
+NDFrameObjectType = NDFrameRectangle