fix: plotting in colab

pyproject.toml

@@ -63,7 +63,7 @@ nltk = {version = "~=3.9", optional = true}
 sentence-transformers = {version = "~=3.1", optional = true}
 pylangacq = {version = "~=0.19", optional = true}
 mediapipe = {version = "~=0.10", optional = true}
-opencv-python = {version = "~=4.10", optional = true}
+opencv-python-headless = {version = "~=4.10", optional = true}
 ultralytics = {version = "~=8.3", optional = true}
 av = {version = "~=14.2", optional = true}
 
@@ -91,7 +91,7 @@ text = [
 video = [
   "av",
   "mediapipe",
-  "opencv-python",
+  "opencv-python-headless",
   "ultralytics"
 ]
 

src/senselab/audio/tasks/plotting/plotting.py

@@ -5,12 +5,13 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import torch
+from matplotlib.pyplot import Figure
 
 from senselab.audio.data_structures import Audio
 from senselab.utils.data_structures import logger
 
 
-def plot_waveform(audio: Audio, title: str = "Waveform", fast: bool = False) -> None:
+def plot_waveform(audio: Audio, title: str = "Waveform", fast: bool = False) -> Figure:
     """Plots the waveform of an Audio object.
 
     Args:
@@ -43,9 +44,10 @@ def plot_waveform(audio: Audio, title: str = "Waveform", fast: bool = False) ->
     figure.suptitle(title)
     plt.xlabel("Time [s]")
     plt.show(block=False)
+    return figure
 
 
-def plot_specgram(audio: Audio, mel_scale: bool = False, title: str = "Spectrogram", **spect_kwargs: Any) -> None:  # noqa : ANN401
+def plot_specgram(audio: Audio, mel_scale: bool = False, title: str = "Spectrogram", **spect_kwargs: Any) -> Figure:  # noqa : ANN401
     """Plots the spectrogram of an Audio object.
 
     Args:
@@ -107,18 +109,41 @@ def _power_to_db(
 
     # Extract the spectrogram
     if mel_scale:
-        from senselab.audio.tasks.features_extraction.torchaudio import extract_mel_spectrogram_from_audios
+        from senselab.audio.tasks.features_extraction.torchaudio import (
+            extract_mel_spectrogram_from_audios,
+        )
 
-        spectrogram = extract_mel_spectrogram_from_audios([audio], **spect_kwargs)[0]["mel_spectrogram"]
+        result = extract_mel_spectrogram_from_audios([audio], **spect_kwargs)[0]
+        spectrogram = result["mel_spectrogram"]
         y_axis_label = "Mel Frequency"
     else:
-        from senselab.audio.tasks.features_extraction.torchaudio import extract_spectrogram_from_audios
+        from senselab.audio.tasks.features_extraction.torchaudio import (
+            extract_spectrogram_from_audios,
+        )
 
-        spectrogram = extract_spectrogram_from_audios([audio], **spect_kwargs)[0]["spectrogram"]
+        result = extract_spectrogram_from_audios([audio], **spect_kwargs)[0]
+        spectrogram = result["spectrogram"]
         y_axis_label = "Frequency [Hz]"
 
-    if spectrogram.dim() != 2:
-        raise ValueError("Spectrogram must be a 2D tensor.")
+    # Handle edge cases for spectrogram tensor
+    if not isinstance(spectrogram, torch.Tensor):
+        raise ValueError("Spectrogram extraction failed - returned non-tensor result.")
+
+    if torch.isnan(spectrogram).all():
+        error_msg = "Spectrogram extraction failed - all values are NaN. " "Audio may be too short."
+        raise ValueError(error_msg)
+
+    # Handle multi-channel spectrograms (3D tensor: channels x frequency x time)
+    if spectrogram.dim() == 3:
+        # Use the first channel for plotting
+        spectrogram = spectrogram[0]
+    elif spectrogram.dim() == 1:
+        # Handle case where spectrogram is 1D (very short audio)
+        error_msg = "Audio is too short to generate a meaningful spectrogram " "for plotting."
+        raise ValueError(error_msg)
+    elif spectrogram.dim() != 2:
+        error_msg = f"Spectrogram must be a 2D tensor, got {spectrogram.dim()}D tensor."
+        raise ValueError(error_msg)
 
     # Determine time and frequency scale
     num_frames = spectrogram.size(1)
@@ -133,7 +158,7 @@ def _power_to_db(
     else:
         freq_axis = torch.linspace(0, audio.sampling_rate / 2, num_freq_bins)
 
-    plt.figure(figsize=(10, 4))
+    figure = plt.figure(figsize=(12, 4))
     plt.imshow(
         _power_to_db(spectrogram.numpy()),
         aspect="auto",
@@ -146,6 +171,7 @@ def _power_to_db(
     plt.ylabel(y_axis_label)
     plt.xlabel("Time [Sec]")
     plt.show(block=False)
+    return figure
 
 
 def play_audio(audio: Audio) -> None:

src/senselab/utils/tasks/plotting.py

@@ -4,11 +4,12 @@
 
 import matplotlib.cm as cm
 import matplotlib.pyplot as plt
+from matplotlib.pyplot import Figure
 
 from senselab.utils.data_structures import ScriptLine
 
 
-def plot_transcript(transcript: ScriptLine) -> None:
+def plot_transcript(transcript: ScriptLine) -> Figure:
     """Plots the transcript visualization over time.
 
     Args:
@@ -39,7 +40,7 @@ def plot_transcript(transcript: ScriptLine) -> None:
             end_times.append(chunk.end)
 
     # Create a figure and axis
-    _, ax = plt.subplots(figsize=(12, 6))
+    figure, ax = plt.subplots(figsize=(12, 6))
 
     # Plot each text segment and add text label
     for i, text in enumerate(texts):
@@ -55,7 +56,8 @@ def plot_transcript(transcript: ScriptLine) -> None:
     ax.set_title("Transcript Visualization Over Time")
 
     # Show the plot
-    plt.show()
+    plt.show(block=False)
+    return figure
 
 
 def plot_segment(segments: List[ScriptLine]) -> None:
@@ -86,7 +88,7 @@ def plot_segment(segments: List[ScriptLine]) -> None:
             labels.append(segment.speaker)
 
     # Create a figure and axis
-    _, ax = plt.subplots(figsize=(12, 6))
+    figure, ax = plt.subplots(figsize=(12, 6))
 
     # Create a color map based on unique labels
     unique_labels = list(set(labels))
@@ -108,4 +110,5 @@ def plot_segment(segments: List[ScriptLine]) -> None:
     ax.set_title("Segment Visualization Over Time")
 
     # Show the plot
-    plt.show()
+    plt.show(block=False)
+    return figure