server: replace wave module with pyav directly

Closes #87
This commit is contained in:
2023-08-16 10:40:40 +02:00
committed by Mathieu Virbel
parent d470696a49
commit 33ab54a626
2 changed files with 26 additions and 18 deletions

View File

@@ -1,6 +1,5 @@
from reflector.processors.base import Processor from reflector.processors.base import Processor
import av import av
import wave
from pathlib import Path from pathlib import Path
@@ -17,19 +16,24 @@ class AudioFileWriterProcessor(Processor):
if isinstance(path, str): if isinstance(path, str):
path = Path(path) path = Path(path)
self.path = path self.path = path
self.fd = None self.out_container = None
self.out_stream = None
async def _push(self, data: av.AudioFrame): async def _push(self, data: av.AudioFrame):
if not self.fd: if not self.out_container:
self.path.parent.mkdir(parents=True, exist_ok=True) self.path.parent.mkdir(parents=True, exist_ok=True)
self.fd = wave.open(self.path.as_posix(), "wb") self.out_container = av.open(self.path.as_posix(), "w", format="wav")
self.fd.setnchannels(len(data.layout.channels)) self.out_stream = self.out_container.add_stream(
self.fd.setsampwidth(data.format.bytes) "pcm_s16le", rate=data.sample_rate
self.fd.setframerate(data.sample_rate) )
self.fd.writeframes(data.to_ndarray().tobytes()) for packet in self.out_stream.encode(data):
self.out_container.mux(packet)
await self.emit(data) await self.emit(data)
async def _flush(self): async def _flush(self):
if self.fd: if self.out_container:
self.fd.close() for packet in self.out_stream.encode(None):
self.fd = None self.out_container.mux(packet)
self.out_container.close()
self.out_container = None
self.out_stream = None

View File

@@ -3,7 +3,6 @@ from reflector.processors.types import AudioFile
from time import monotonic_ns from time import monotonic_ns
from uuid import uuid4 from uuid import uuid4
import io import io
import wave
import av import av
@@ -28,12 +27,16 @@ class AudioMergeProcessor(Processor):
# create audio file # create audio file
uu = uuid4().hex uu = uuid4().hex
fd = io.BytesIO() fd = io.BytesIO()
with wave.open(fd, "wb") as wf:
wf.setnchannels(channels) out_container = av.open(fd, "w", format="wav")
wf.setsampwidth(sample_width) out_stream = out_container.add_stream("pcm_s16le", rate=sample_rate)
wf.setframerate(sample_rate) for frame in data:
for frame in data: for packet in out_stream.encode(frame):
wf.writeframes(frame.to_ndarray().tobytes()) out_container.mux(packet)
for packet in out_stream.encode(None):
out_container.mux(packet)
out_container.close()
fd.seek(0)
# emit audio file # emit audio file
audiofile = AudioFile( audiofile = AudioFile(
@@ -44,4 +47,5 @@ class AudioMergeProcessor(Processor):
sample_width=sample_width, sample_width=sample_width,
timestamp=data[0].pts * data[0].time_base, timestamp=data[0].pts * data[0].time_base,
) )
await self.emit(audiofile) await self.emit(audiofile)