Merge pull request #339 from Labelbox/tpeharda/DIAG-904-import-prediction-progress

[DIAG-904] Show Annotation Import progress

Author: msokoloff1

labelbox/data/annotation_types/classification/classification.py

@@ -25,10 +25,10 @@ class ClassificationAnswer(FeatureSchema):
     - Because it inherits from FeatureSchema
         the option can be represented with either the name or feature_schema_id
 
-    - The key frame arg only applies to video classifications.
-      Each answer can have a key frame indepdent of the others.
+    - The keyframe arg only applies to video classifications.
+      Each answer can have a keyframe independent of the others.
         So unlike object annotations, classification annotations
-          track key frames at a classification answer level.
+          track keyframes at a classification answer level.
     """
     extra: Dict[str, Any] = {}
     keyframe: Optional[bool] = None

labelbox/schema/annotation_import.py

@@ -4,6 +4,7 @@
 import os
 import time
 from typing import Any, Dict, List, BinaryIO
+from tqdm import tqdm  # type: ignore
 
 import backoff
 import ndjson
@@ -25,6 +26,7 @@ class AnnotationImport(DbObject):
     input_file_url = Field.String("input_file_url")
     error_file_url = Field.String("error_file_url")
     status_file_url = Field.String("status_file_url")
+    progress = Field.String("progress")
 
     created_by = Relationship.ToOne("User", False, "created_by")
 
@@ -76,18 +78,30 @@ def statuses(self) -> List[Dict[str, Any]]:
         self.wait_until_done()
         return self._fetch_remote_ndjson(self.status_file_url)
 
-    def wait_until_done(self, sleep_time_seconds: int = 10) -> None:
+    def wait_until_done(self,
+                        sleep_time_seconds: int = 10,
+                        show_progress: bool = False) -> None:
         """Blocks import job until certain conditions are met.
         Blocks until the AnnotationImport.state changes either to
         `AnnotationImportState.FINISHED` or `AnnotationImportState.FAILED`,
         periodically refreshing object's state.
         Args:
-            sleep_time_seconds (str): a time to block between subsequent API calls
+            sleep_time_seconds (int): a time to block between subsequent API calls
+            show_progress (bool): should show progress bar
         """
+        pbar = tqdm(total=100,
+                    bar_format="{n}% |{bar}| [{elapsed}, {rate_fmt}{postfix}]"
+                   ) if show_progress else None
         while self.state.value == AnnotationImportState.RUNNING.value:
             logger.info(f"Sleeping for {sleep_time_seconds} seconds...")
             time.sleep(sleep_time_seconds)
             self.__backoff_refresh()
+            if self.progress and self.progress and pbar:
+                pbar.update(int(self.progress.replace("%", "")) - pbar.n)
+
+        if pbar:
+            pbar.update(100 - pbar.n)
+            pbar.close()
 
     @backoff.on_exception(
         backoff.expo,
@@ -499,7 +513,7 @@ def create_from_objects(cls, client: "labelbox.Client", project_id: str,
             client: Labelbox Client for executing queries
             project_id: Project to import labels into
             name: Name of the import job. Can be used to reference the task later
-            labels: List of labels 
+            labels: List of labels
         Returns:
             LabelImport
         """

setup.py

@@ -27,11 +27,12 @@
         "requests>=2.22.0",
         "google-api-core>=1.22.1",
         "pydantic>=1.8,<2.0",
+        "tqdm",
     ],
     extras_require={
         'data': [
             "shapely", "geojson", "numpy", "rasterio", "PILLOW",
-            "opencv-python", "typeguard", "tqdm", "imagesize"
+            "opencv-python", "typeguard", "imagesize"
         ],
     },
     classifiers=[