Merge pull request #2206 from pupil-labs/fix_short_audio_loading

Fix loading short audio streams

Author: papr

pupil_src/shared_modules/audio_playback.py

@@ -81,14 +81,18 @@ def __init__(self, g_pool):
 
         try:
             self.audio_all = load_audio(self.g_pool.rec_dir)
+            logger.debug("Audio_Playback.__init__: Audio loaded successfully")
         except NoAudioLoadedError:
+            logger.debug("Audio_Playback.__init__: No audio loaded")
             return
 
         self.calculate_audio_bounds()
 
         self.filter_graph = None
         self.filter_graph_list = None
+        logger.debug("Audio_Playback.__init__: Initializing PyAudio")
         self.pa = pa.PyAudio()
+        logger.debug("Audio_Playback.__init__: PyAudio initialized")
 
         self._setup_input_audio_part(0)
 

pupil_src/shared_modules/audio_utils.py

@@ -11,6 +11,7 @@
 import collections
 import logging
 import traceback
+import typing as T
 
 import av
 import numpy as np
@@ -25,9 +26,19 @@ class NoAudioLoadedError(Exception):
     pass
 
 
-LoadedAudio = collections.namedtuple(
-    "LoadedAudio", ["container", "stream", "timestamps", "pts"]
-)
+class LoadedAudio(T.NamedTuple):
+    container: T.Any
+    stream: T.Any
+    timestamps: T.List[float]
+    pts: T.List[int]
+
+    def __str__(self):
+        return (
+            f"{type(self).__name__}(container={self.container}, stream={self.stream}, "
+            f"timestamps=(N={len(self.timestamps)}, [{self.timestamps[0]}, "
+            f"{self.timestamps[-1]}]), pts=(N={len(self.pts)}, [{self.pts[0]}, "
+            f"{self.pts[-1]}]))"
+        )
 
 
 def load_audio(rec_dir):
@@ -95,7 +106,8 @@ def _load_audio_single(file_path, return_pts_based_timestamps=False):
 
 
 class Audio_Viz_Transform:
-    def __init__(self, rec_dir, sps_rate=60):
+    def __init__(self, rec_dir, log_scaling=False, sps_rate=60):
+        logger.debug("Audio_Viz_Transform.__init__: Loading audio")
         self.audio_all = iter(load_audio(rec_dir))
         self._setup_next_audio_part()
         self._first_part_start = self.audio.timestamps[0]
@@ -106,13 +118,19 @@ def __init__(self, rec_dir, sps_rate=60):
         self.a_levels = None
         self.a_levels_log = None
         self.final_rescale = True
-        self.log_scaling = False
+        self.log_scaling = log_scaling
 
     def _setup_next_audio_part(self):
         self.audio = next(self.audio_all)
+        logger.debug(
+            f"Audio_Viz_Transform._setup_next_audio_part: Part {self.audio.container} {self.audio.stream}"
+        )
         self.audio_resampler = av.audio.resampler.AudioResampler(
             format=self.audio.stream.format, layout=self.audio.stream.layout, rate=60
         )
+        logger.debug(
+            "Audio_Viz_Transform._setup_next_audio_part: Resampler initialized"
+        )
         self.next_audio_frame = self._next_audio_frame()
         self.start_ts = self.audio.timestamps[0]
 
@@ -190,7 +208,7 @@ def get_data(self, seconds=30.0, height=210, log_scale=False):
                 else:
                     scaled_samples = abs_samples
 
-            else:
+            elif self.a_levels is not None and self.all_abs_samples is not None:
                 new_ts = self.a_levels[::4]  # reconstruct correct ts
 
                 # self.all_abs_samples = np.log10(self.all_abs_samples)
@@ -212,7 +230,17 @@ def get_data(self, seconds=30.0, height=210, log_scale=False):
                     self._setup_next_audio_part()
                 except StopIteration:
                     self.finished = True
-            if not self.finished or self.final_rescale:
+            else:
+                logger.debug(
+                    f"Audio_Viz_Transform.get_data: No audio found in {self.audio}"
+                )
+                new_ts = None
+                try:
+                    self._setup_next_audio_part()
+                except StopIteration:
+                    self.finished = True
+
+            if new_ts is not None and (not self.finished or self.final_rescale):
                 a_levels = self.get_verteces(new_ts, scaled_samples, height)
 
                 if self.a_levels is not None: