Merge branch 'tpeharda/DIAG-904-import-prediction-progress' of https://github.com/Labelbox/labelbox-python into develop

Author: Matt Sokoloff

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,28 @@ 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) 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 pbar:
+                pbar.update(self.progress)
+
+        if pbar:
+            pbar.update(100)
+            pbar.close()
 
     @backoff.on_exception(
         backoff.expo,

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=[