chore: clean up and track test files

Author: rishisurana-labelbox

libs/labelbox/src/labelbox/data/serialization/ndjson/classification.py

@@ -60,6 +60,22 @@ def serialize_model(self, handler):
         return res
 
 
+class FrameLocation(BaseModel):
+    end: int
+    start: int
+
+
+class VideoSupported(BaseModel):
+    # Note that frames are only allowed as top level inferences for video
+    frames: Optional[List[FrameLocation]] = None
+
+    @model_serializer(mode="wrap")
+    def serialize_model(self, handler):
+        res = handler(self)
+        # This means these are no video frames ..
+        if self.frames is None:
+            res.pop("frames")
+        return res
 
 
 class NDTextSubclass(NDAnswer):
@@ -226,14 +242,13 @@ def from_common(
             name=name,
             schema_id=feature_schema_id,
             uuid=uuid,
-            frames=extra.get("frames"),
             message_id=message_id,
             confidence=text.confidence,
             custom_metrics=text.custom_metrics,
         )
 
 
-class NDChecklist(NDAnnotation, NDChecklistSubclass):
+class NDChecklist(NDAnnotation, NDChecklistSubclass, VideoSupported):
     @model_serializer(mode="wrap")
     def serialize_model(self, handler):
         res = handler(self)
@@ -280,7 +295,7 @@ def from_common(
         )
 
 
-class NDRadio(NDAnnotation, NDRadioSubclass):
+class NDRadio(NDAnnotation, NDRadioSubclass, VideoSupported):
     @classmethod
     def from_common(
         cls,
@@ -410,8 +425,7 @@ def to_common(
     def from_common(
         cls,
         annotation: Union[
-            ClassificationAnnotation,
-            VideoClassificationAnnotation,
+            ClassificationAnnotation, VideoClassificationAnnotation
         ],
         data: GenericDataRowData,
     ) -> Union[NDTextSubclass, NDChecklistSubclass, NDRadioSubclass]:
@@ -434,8 +448,7 @@ def from_common(
     @staticmethod
     def lookup_classification(
         annotation: Union[
-            ClassificationAnnotation,
-            VideoClassificationAnnotation,
+            ClassificationAnnotation, VideoClassificationAnnotation
         ],
     ) -> Union[NDText, NDChecklist, NDRadio]:
         return {Text: NDText, Checklist: NDChecklist, Radio: NDRadio}.get(

libs/labelbox/tests/conftest.py

@@ -688,12 +688,12 @@ def create_label():
             predictions,
         )
         upload_task.wait_until_done(sleep_time_seconds=5)
-        assert upload_task.state == AnnotationImportState.FINISHED, (
-            "Label Import did not finish"
-        )
-        assert len(upload_task.errors) == 0, (
-            f"Label Import {upload_task.name} failed with errors {upload_task.errors}"
-        )
+        assert (
+            upload_task.state == AnnotationImportState.FINISHED
+        ), "Label Import did not finish"
+        assert (
+            len(upload_task.errors) == 0
+        ), f"Label Import {upload_task.name} failed with errors {upload_task.errors}"
 
     project.create_label = create_label
     project.create_label()