feat: updated progress bar (#60)

Co-authored-by: anna-grim <anna.grim@alleninstitute.org>

Author: anna-grim

src/segmentation_skeleton_metrics/swc_utils.py

@@ -110,29 +110,28 @@ def parse_cloud_paths(cloud_dict, min_size, anisotropy):
     zip_paths = utils.list_gcs_filenames(bucket, cloud_dict["path"], ".zip")
     print("Downloading predicted swc files from cloud...")
     print("# zip files:", len(zip_paths))
+
+    # Assign processes
+    chunk_size = int(len(zip_paths) * 0.02)
+    cnt = 1
+    t0, t1 = utils.init_timers()
     with ProcessPoolExecutor() as executor:
         processes = []
-        for path in zip_paths:
+        for i, path in enumerate(zip_paths):
             zip_content = bucket.blob(path).download_as_bytes()
             processes.append(
                 executor.submit(download, zip_content, anisotropy, min_size)
             )
-    print("Processes Assigned!\n")
+            if i > cnt * chunk_size:
+                cnt, t1 = utils.report_progress(
+                    i, len(zip_paths), chunk_size, cnt, t0, t1
+                )
 
     # Store results
-    chunk_size = int(len(zip_paths) * 0.02)
-    cnt = 1
     valid_labels = dict()
-    t0, t1 = utils.init_timers()
     for i, process in enumerate(as_completed(processes)):
         valid_labels.update(process.result())
-        if i > cnt * chunk_size:
-            utils.progress_bar(i + 1, len(zip_paths))
-            cnt += 1
-
-    # Report Results
     print("\n#Valid Labels:", len(valid_labels))
-    print("")
     return valid_labels
 
 

src/segmentation_skeleton_metrics/utils.py

@@ -439,6 +439,29 @@ def init_timers():
     return time(), time()
 
 
+# -- Utils --
+def report_progress(current, total, chunk_size, cnt, t0, t1):
+    eta = get_eta(current, total, chunk_size, t1)
+    runtime = get_runtime(current, total, chunk_size, t0, t1)
+    progress_bar2(current, total, eta=eta, runtime=runtime)
+    return cnt + 1, time()
+
+
+def get_eta(current, total, chunk_size, t0, return_str=True):
+    chunk_runtime = time() - t0
+    remaining = total - current
+    eta = remaining * (chunk_runtime / chunk_size)
+    t, unit = time_writer(eta)
+    return f"{round(t, 4)} {unit}" if return_str else eta
+
+
+def get_runtime(current, total, chunk_size, t0, t1):
+    eta = get_eta(current, total, chunk_size, t1, return_str=False)
+    total_runtime = time() - t0 + eta
+    t, unit = time_writer(total_runtime)
+    return f"{round(t, 4)} {unit}"
+
+
 def progress_bar(current, total, bar_length=50):
     """
     Reports the progress of completing some process.
@@ -462,3 +485,11 @@ def progress_bar(current, total, bar_length=50):
         f"[{'=' * progress}{' ' * (bar_length - progress)}] {current}/{total}"
     )
     print(f"\r{bar}", end="", flush=True)
+
+def progress_bar2(current, total, bar_length=50, eta=None, runtime=None):
+    progress = int(current / total * bar_length)
+    n_completed = f"Completed: {current}/{total}"
+    bar = f"[{'=' * progress}{' ' * (bar_length - progress)}]"
+    eta = f"Time Remaining: {eta}" if eta else ""
+    runtime = f"Estimated Total Runtime: {runtime}" if runtime else ""
+    print(f"\r{bar} {n_completed} | {eta} | {runtime}    ", end="", flush=True)