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
import av
import wave
from pathlib import Path
@@ -17,19 +16,24 @@ class AudioFileWriterProcessor(Processor):
if isinstance(path, str):
path = Path(path)
self.path = path
self.fd = None
self.out_container = None
self.out_stream = None
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.fd = wave.open(self.path.as_posix(), "wb")
self.fd.setnchannels(len(data.layout.channels))
self.fd.setsampwidth(data.format.bytes)
self.fd.setframerate(data.sample_rate)
self.fd.writeframes(data.to_ndarray().tobytes())
self.out_container = av.open(self.path.as_posix(), "w", format="wav")
self.out_stream = self.out_container.add_stream(
"pcm_s16le", rate=data.sample_rate
)
for packet in self.out_stream.encode(data):
self.out_container.mux(packet)
await self.emit(data)
async def _flush(self):
if self.fd:
self.fd.close()
self.fd = None
if self.out_container:
for packet in self.out_stream.encode(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 uuid import uuid4
import io
import wave
import av
@@ -28,12 +27,16 @@ class AudioMergeProcessor(Processor):
# create audio file
uu = uuid4().hex
fd = io.BytesIO()
with wave.open(fd, "wb") as wf:
wf.setnchannels(channels)
wf.setsampwidth(sample_width)
wf.setframerate(sample_rate)
for frame in data:
wf.writeframes(frame.to_ndarray().tobytes())
out_container = av.open(fd, "w", format="wav")
out_stream = out_container.add_stream("pcm_s16le", rate=sample_rate)
for frame in data:
for packet in out_stream.encode(frame):
out_container.mux(packet)
for packet in out_stream.encode(None):
out_container.mux(packet)
out_container.close()
fd.seek(0)
# emit audio file
audiofile = AudioFile(
@@ -44,4 +47,5 @@ class AudioMergeProcessor(Processor):
sample_width=sample_width,
timestamp=data[0].pts * data[0].time_base,
)
await self.emit(audiofile)