debug_app.py which choppy audio

debug_001.py

@@ -54,9 +54,9 @@ with open("chunks.pkl", "rb") as f:
             bytes_io.seek(0, io.SEEK_END)
             continue
 
-    with open("frames.pkl", "wb") as f:
-        import pickle
-        pickle.dump(chunks, f)
+    # with open("frames.pkl", "wb") as f:
+    #     import pickle
+    #     pickle.dump(audio_frames, f)
 
 
 if mpv_process.stdin:

debug_app.py

@@ -0,0 +1,167 @@
+import asyncio
+import io
+import logging
+import traceback
+from typing import List
+
+import av
+import numpy as np
+import streamlit as st
+from streamlit_webrtc import WebRtcMode, webrtc_streamer
+import pydub
+
+from dotenv import load_dotenv
+load_dotenv()
+from sample_utils.turn import get_ice_servers
+
+logger = logging.getLogger(__name__)
+
+
+class StreamingMP3ToFrames:
+    def __init__(self):
+        self.append = False
+
+    def process_chunk(self, chunk):
+        audio_frames = []
+        try:
+            if self.append:
+                self.bytes_io.write(chunk)
+                self.append = False
+                self.bytes_io.seek(0)
+            else:
+                self.bytes_io = io.BytesIO(chunk)
+            container = av.open(self.bytes_io, 'r', format='mp3')
+            audio_stream = next(s for s in container.streams if s.type == 'audio')
+            for frame in container.decode(audio_stream):
+                # Convert the audio frame to a NumPy array
+                array = frame.to_ndarray()
+
+                # Now you can use av.AudioFrame.from_ndarray
+                
+                # audio_frame = av.AudioFrame.from_ndarray(array, format='flt', layout='mono')
+                audio_frame = av.AudioFrame.from_ndarray(array, format='fltp', layout='mono')
+                audio_frame.sample_rate = 44100
+
+                audio_frames.append(audio_frame)
+
+            return audio_frames
+
+        except Exception as e:
+            print (e)
+            self.append = True
+            self.bytes_io.seek(0, io.SEEK_END)
+            return audio_frames
+
+
+def video_frame_callback(
+    frame: av.VideoFrame,
+) -> av.VideoFrame:
+    return frame
+
+
+
+streaming_mp3_to_frames = StreamingMP3ToFrames()
+
+with open("chunks.pkl", "rb") as f:
+    import pickle
+    debug_chunks = pickle.load(f)
+    debug_frames = []
+    debug_frame_idx = 0
+    for chunk in debug_chunks:
+        new_frames = streaming_mp3_to_frames.process_chunk(chunk)
+        for frame in new_frames:
+            debug_frames.append(frame)
+            # print (frame)
+
+def dequeue_frame():
+    global debug_frame_idx, debug_frames
+    enqueued_frame = debug_frames[debug_frame_idx]
+    debug_frame_idx += 1
+    if debug_frame_idx >= len(debug_frames):
+        debug_frame_idx = 0
+    return enqueued_frame
+
+# emptry array of type int16
+sample_buffer = np.zeros((0), dtype=np.int16)
+
+def process_frame(old_frame):
+
+    try:
+        output_channels = 2
+        output_sample_rate = 44100
+        required_samples = old_frame.samples
+
+        global sample_buffer
+        while sample_buffer.shape[0] < required_samples:
+            dequeued_frame = dequeue_frame()
+            if dequeued_frame is None:
+                break
+
+            # convert dequeued_frame to same format as old_frame
+            float_samples = dequeued_frame.to_ndarray()
+            max_sample = np.max(np.abs(float_samples))
+            min_sample = np.min(np.abs(float_samples))
+            if max_sample > 1.0 or min_sample > 1.0:
+                print(f"WARNING: max_sample: {max_sample}, min_sample: {min_sample}")
+            int_samples = np.int16(float_samples * 32767)
+            sound = pydub.AudioSegment(
+                data=int_samples.tobytes(),
+                sample_width=2,
+                frame_rate=output_sample_rate,
+                channels=len(dequeued_frame.layout.channels),
+            )
+            sound = sound.set_frame_rate(old_frame.sample_rate)
+
+            samples = np.array(sound.get_array_of_samples(), dtype=np.int16)
+            sample_buffer = np.append(sample_buffer, samples)
+
+        # handle case where we ran out of frames
+        if sample_buffer.shape[0] < required_samples:
+            empty_samples = np.zeros((required_samples - sample_buffer.shape[0]), dtype=np.int16)
+            sample_buffer = np.append(sample_buffer, empty_samples)
+
+        # take the first required_samples samples from the buffer
+        samples = sample_buffer[:required_samples]
+        sample_buffer = sample_buffer[required_samples:]
+
+        # Duplicate mono channel for stereo
+        if output_channels == 2:
+            samples = np.vstack((samples, samples)).reshape((-1,), order='F')
+
+        samples = samples.reshape(1, -1)
+
+        layout = 'stereo' if output_channels == 2 else 'mono'
+        new_frame = av.AudioFrame.from_ndarray(samples, format='s16', layout=layout)
+        new_frame.sample_rate = old_frame.sample_rate
+        new_frame.pts = old_frame.pts
+        return new_frame
+    except Exception as e:
+        print (e)
+        traceback.print_exc() 
+        raise(e)
+
+
+
+def audio_frame_callback(old_frame: av.AudioFrame) -> av.AudioFrame:
+
+    global debug_frame_idx, debug_frames
+
+    new_frame = process_frame(old_frame)
+
+    # print (f"new_frames: {len(new_frames)}, frames: {len(frames)}")
+    print (f"frame:     {old_frame}, pts: {old_frame.pts}")
+    print (f"new_frame: {new_frame}, pts: {new_frame.pts}")
+
+    return new_frame
+    # return old_frame
+
+
+webrtc_streamer(
+    key="delay",
+    mode=WebRtcMode.SENDRECV,
+    rtc_configuration={"iceServers": get_ice_servers()},
+
+    video_frame_callback=video_frame_callback,
+    audio_frame_callback=audio_frame_callback,
+
+)