feat(openadapt): add progress bar in record.py and visualize.py

openadapt/record.py

@@ -17,6 +17,7 @@
 import threading
 import time
 
+from alive_progress import alive_bar
 from loguru import logger
 from pynput import keyboard, mouse
 import fire
@@ -247,14 +248,18 @@ def write_events(
     utils.set_start_time(recording_timestamp)
     logger.info(f"{event_type=} starting")
     signal.signal(signal.SIGINT, signal.SIG_IGN)
-    while not terminate_event.is_set() or not write_q.empty():
-        try:
-            event = write_q.get_nowait()
-        except queue.Empty:
-            continue
-        assert event.type == event_type, (event_type, event)
-        write_fn(recording_timestamp, event, perf_q)
-        logger.debug(f"{event_type=} written")
+
+    with alive_bar(total=write_q.qsize(), title=f"Writing {event_type} events") as progress:
+        while not terminate_event.is_set() or not write_q.empty():
+            try:
+                event = write_q.get_nowait()
+            except queue.Empty:
+                continue
+            assert event.type == event_type, (event_type, event)
+            write_fn(recording_timestamp, event, perf_q)
+            logger.debug(f"{event_type=} written")
+            progress()
+
     logger.info(f"{event_type=} done")
 
 

openadapt/visualize.py

@@ -4,6 +4,7 @@
 import os
 import string
 
+from alive_progress import alive_bar
 from bokeh.io import output_file, show
 from bokeh.layouts import layout, row
 from bokeh.models.widgets import Div
@@ -188,70 +189,76 @@ def main():
         ),
     ]
     logger.info(f"{len(action_events)=}")
-    for idx, action_event in enumerate(action_events):
-        if idx == MAX_EVENTS:
-            break
-        image = display_event(action_event)
-        diff = display_event(action_event, diff=True)
-        mask = action_event.screenshot.diff_mask
-
-        if SCRUB:
-            image = scrub.scrub_image(image)
-            diff = scrub.scrub_image(diff)
-            mask = scrub.scrub_image(mask)
-
-        image_utf8 = image2utf8(image)
-        diff_utf8 = image2utf8(diff)
-        mask_utf8 = image2utf8(mask)
-        width, height = image.size
-
-        action_event_dict = row2dict(action_event)
-        window_event_dict = row2dict(action_event.window_event)
-
-        if SCRUB:
-            action_event_dict = scrub.scrub_dict(action_event_dict)
-            window_event_dict = scrub.scrub_dict(window_event_dict)
-
-        rows.append(
-            [
-                row(
-                    Div(
-                        text=f"""
-                        <div class="screenshot">
-                            <img
-                                src="{image_utf8}"
-                                style="
-                                    aspect-ratio: {width}/{height};
-                                "
-                            >
-                            <img
-                                src="{diff_utf8}"
-                                style="
-                                    aspect-ratio: {width}/{height};
-                                "
-                            >
-                            <img
-                                src="{mask_utf8}"
-                                style="
-                                    aspect-ratio: {width}/{height};
-                                "
-                            >
-                        </div>
-                        <table>
-                            {dict2html(window_event_dict , None)}
-                        </table>
-                    """,
-                    ),
-                    Div(
-                        text=f"""
-                        <table>
-                            {dict2html(action_event_dict)}
-                        </table>
-                    """
+
+    rows = []
+    num_events = min(MAX_EVENTS, len(action_events)) if MAX_EVENTS is not None else len(action_events)
+    with alive_bar(total=num_events, title="Processing HTML for visualization:") as progress:
+        for idx, action_event in enumerate(action_events):
+            if idx == MAX_EVENTS:
+                break
+            image = display_event(action_event)
+            diff = display_event(action_event, diff=True)
+            mask = action_event.screenshot.diff_mask
+
+            if SCRUB:
+                image = scrub.scrub_image(image)
+                diff = scrub.scrub_image(diff)
+                mask = scrub.scrub_image(mask)
+
+            image_utf8 = image2utf8(image)
+            diff_utf8 = image2utf8(diff)
+            mask_utf8 = image2utf8(mask)
+            width, height = image.size
+
+            action_event_dict = row2dict(action_event)
+            window_event_dict = row2dict(action_event.window_event)
+
+            if SCRUB:
+                action_event_dict = scrub.scrub_dict(action_event_dict)
+                window_event_dict = scrub.scrub_dict(window_event_dict)
+
+            rows.append(
+                [
+                    row(
+                        Div(
+                            text=f"""
+                            <div class="screenshot">
+                                <img
+                                    src="{image_utf8}"
+                                    style="
+                                        aspect-ratio: {width}/{height};
+                                    "
+                                >
+                                <img
+                                    src="{diff_utf8}"
+                                    style="
+                                        aspect-ratio: {width}/{height};
+                                    "
+                                >
+                                <img
+                                    src="{mask_utf8}"
+                                    style="
+                                        aspect-ratio: {width}/{height};
+                                    "
+                                >
+                            </div>
+                            <table>
+                                {dict2html(window_event_dict , None)}
+                            </table>
+                        """,
+                        ),
+                        Div(
+                            text=f"""
+                            <table>
+                                {dict2html(action_event_dict)}
+                            </table>
+                        """
+                        ),
                     ),
-                ),
-            ]
-        )
+                ]
+            )
+
+            progress()
 
     title = f"recording-{recording.id}"
     fname_out = f"recording-{recording.id}.html"

pyproject.toml

@@ -57,6 +57,7 @@ nltk = "3.8.1"
 pywinauto = {version = "^0.6.8", markers = "sys_platform == 'win32'"}
 moviepy = "1.0.3"
 python-levenshtein = "^0.21.1"
+alive-progress = "^3.1.4"
 
 [tool.poetry.dependencies.en_core_web_trf]
 url = "https://github.com/explosion/spacy-models/releases/download/en_core_web_trf-3.5.0/en_core_web_trf-3.5.0.tar.gz"