File size: 2,240 Bytes
b1882a9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | #!/usr/bin/env python3
"""Minimal jyoti usage: encode a 16 kHz mono file to a 2.3 kbps stream and decode it back, in process.
python example.py input_16k_mono.wav output.wav [--checkpoint codec.pt]
The stream is bound to the checkpoint's SHA-256; decoding with a different checkpoint raises an error.
"""
import argparse
import hashlib
from pathlib import Path
import sys
import soundfile as sf
import torch
ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))
from tinyasr.codec_factory_v3 import load_codec_checkpoint # noqa: E402
from tinyasr.codec_live_v1 import LiveEncoder # noqa: E402
from tinyasr.codec_phase_streaming_v1 import PhaseLiveDecoder # noqa: E402
def main():
ap = argparse.ArgumentParser()
ap.add_argument('input', type=Path)
ap.add_argument('output', type=Path)
ap.add_argument('--checkpoint', type=Path, default=ROOT / 'codec.pt')
args = ap.parse_args()
if args.output.exists():
raise FileExistsError(args.output)
audio, rate = sf.read(args.input, dtype='float32', always_2d=False)
if rate != 16000 or audio.ndim != 1:
raise ValueError(f'Expected 16 kHz mono input, got {rate} Hz with shape {audio.shape}')
torch.set_num_threads(1)
model_sha256 = hashlib.sha256(args.checkpoint.read_bytes()).hexdigest()
codec = load_codec_checkpoint(args.checkpoint)
# Encode in 20 ms (320-sample) blocks, as a live sender would.
encoder = LiveEncoder(codec, model_sha256)
stream = bytearray(encoder.start())
x = torch.from_numpy(audio)
for start in range(0, x.numel(), 320):
stream += encoder.push(x[start:start + 320])
stream += encoder.finish()
# Decode the byte stream; finish() checks the footer and CRC.
decoder = PhaseLiveDecoder(codec, model_sha256)
parts = [decoder.push(bytes(stream))]
tail, end = decoder.finish()
parts.append(tail)
decoded = torch.cat(parts).numpy()
sf.write(args.output, decoded, 16000, subtype='FLOAT')
seconds = audio.size / 16000
print(f'{seconds:.2f} s audio -> {len(stream)} bytes ({8 * len(stream) / seconds:.0f} bit/s incl. framing); '
f'decoded {decoded.size} samples -> {args.output}')
if __name__ == '__main__':
main()
|