From 396b5eee97ed7e25088484255fe6411c24a10da1 Mon Sep 17 00:00:00 2001
From: Sepehr Heravi <sepehr.heravi2008@gmail.com>
Date: Wed, 15 Apr 2026 12:05:40 -0700
Subject: [PATCH] Guard plot_stimuli against missing audio / word events

BasePlotBrain.plot_stimuli crashed for silent or speechless clips because
it unconditionally called audio.to_soundarray() and events.type. This made
plot_timesteps(show_stimuli=True) unusable for screen recordings and any
video where WhisperX produces zero words, even though main.py already
handles the missing-modality case by auto-dropping the text extractor.

- Plot the audio waveform only when get_audio(...) returns non-None.
- Skip the word-overlay loop when segment.events is None or has no
  'type' column.
- Behaviour is unchanged when both modalities are present.

Fixes #46
---
 tribev2/plotting/base.py | 43 ++++++++++++++++++++++------------------
 1 file changed, 24 insertions(+), 19 deletions(-)

diff --git a/tribev2/plotting/base.py b/tribev2/plotting/base.py
index fc95ac0b..6aa4d5f0 100644
--- a/tribev2/plotting/base.py
+++ b/tribev2/plotting/base.py
@@ -376,13 +376,16 @@ def plot_stimuli(
 
         TEXT_KEY, SOUND_KEY, VIDEO_KEY = "Text", "Audio", "Video"
 
+        # Audio is optional (e.g., silent screen recordings produce no audio).
         audio = get_audio(
             segments[0], stop_offset=(len(segments) - 1) * segments[0].duration
         )
-        soundarray = audio.to_soundarray().mean(axis=1)
-        axes[SOUND_KEY].plot(soundarray, color="k")
-        axes[SOUND_KEY].set_xlim(0, len(soundarray))
-        axes[SOUND_KEY].axis("off")
+        if audio is not None:
+            soundarray = audio.to_soundarray().mean(axis=1)
+            axes[SOUND_KEY].plot(soundarray, color="k")
+            axes[SOUND_KEY].set_xlim(0, len(soundarray))
+        if SOUND_KEY in axes:
+            axes[SOUND_KEY].axis("off")
         axes[TEXT_KEY].axis("off")
         full_start, full_duration = (
             segments[0].start,
@@ -411,22 +414,24 @@ def plot_stimuli(
                 ax.set_xlim(-margin, img.shape[1] + margin)
                 ax.set_ylim(img.shape[0] + margin, -margin)
                 ax.axis("off")
+            # Text/words are optional (silent or non-speech clips have no word events).
             events = segment.events
-            words = events[events.type == "Word"]
-            for word in words.itertuples():
-                if word.start < full_start:
-                    continue
-                axes[TEXT_KEY].text(
-                    (word.start - full_start) / full_duration,
-                    0.5,
-                    word.text,
-                    color="k",
-                    transform=axes[TEXT_KEY].transAxes,
-                    ha="center",
-                    va="center",
-                    rotation=45,
-                    fontsize=10,
-                )
+            if events is not None and "type" in events.columns:
+                words = events[events.type == "Word"]
+                for word in words.itertuples():
+                    if word.start < full_start:
+                        continue
+                    axes[TEXT_KEY].text(
+                        (word.start - full_start) / full_duration,
+                        0.5,
+                        word.text,
+                        color="k",
+                        transform=axes[TEXT_KEY].transAxes,
+                        ha="center",
+                        va="center",
+                        rotation=45,
+                        fontsize=10,
+                    )
 
     def plot_timesteps_mp4(
         self,