fix(visualize): Handle no screenshots (#763)

* handle no screenshot in visualize.py / plotting.py * black/flake8
OpenAdaptAI · Jun 18, 2024 · 92b74f0 · 92b74f0
1 parent 220c607
commit 92b74f0
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 24 deletions.
diff --git a/openadapt/plotting.py b/openadapt/plotting.py
@@ -249,6 +249,12 @@ def display_event(
     recording = action_event.recording
     window_event = action_event.window_event
     screenshot = action_event.screenshot
+
+    if not screenshot:
+        logger.warning(
+            f"{screenshot=} for {action_event=} {window_event=} {recording=}"
+        )
+        return None
     if diff and screenshot.diff:
         image = screenshot.diff.convert("RGBA")
     else:

diff --git a/openadapt/record.py b/openadapt/record.py
@@ -328,7 +328,6 @@ def write_events(
     started_counter: multiprocessing.Value,
     pre_callback: Callable[[float], dict] | None = None,
     post_callback: Callable[[dict], None] | None = None,
-    event_type_modifier: str = "",
 ) -> None:
     """Write events of a specific type to the db using the provided write function.
 
@@ -345,7 +344,6 @@ def write_events(
             timestamp as only argument, returns a state dict.
         post_callback: Optional function to call after main loop. Takes state dict as
             only argument, returns None.
-        event_type_modifier: Optional string to differentiate identical event_types
     """
     utils.set_start_time(recording.timestamp)
 
@@ -368,7 +366,7 @@ def write_events(
                 total_events = num_events.value
                 progress = tqdm(
                     total=total_events,
-                    desc=f"Writing {event_type}{event_type_modifier} events...",
+                    desc=f"Writing {event_type} events...",
                     unit="event",
                     colour="green",
                     dynamic_ncols=True,

diff --git a/openadapt/visualize.py b/openadapt/visualize.py
@@ -313,28 +313,36 @@ def main(
                     logger.warning(exc)
                     continue
 
-                if diff_video:
-                    frame_image = frames[idx]
-                    diff_image = compute_diff(
-                        frame_image, action_event.screenshot.image
-                    )
-
-                    # TODO: rename
-                    diff = frame_image
-                    mask = diff_image
+                if image:
+                    if diff_video:
+                        frame_image = frames[idx]
+                        diff_image = compute_diff(
+                            frame_image, action_event.screenshot.image
+                        )
+
+                        # TODO: rename
+                        diff = frame_image
+                        mask = diff_image
+                    else:
+                        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
                 else:
-                    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
+                    # TODO: display a placeholder image
+                    image_utf8 = ""
+                    diff_utf8 = ""
+                    mask_utf8 = ""
+                    width = 0
+                    height = 1
 
                 action_event_dict = row2dict(action_event)
                 window_event_dict = row2dict(action_event.window_event)