Spaces:
Paused
Paused
File size: 3,818 Bytes
e319a7b | 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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | """Audio loading and resampling utilities. WAV is handled by the stdlib;
other formats fall back to `soundfile`."""
import wave
from pathlib import Path
from typing import IO
import numpy as np
import torch
from muscriptor.utils.resample import resample_frac
def _read_wav_file(source) -> tuple[torch.Tensor, int]:
"""Load a PCM WAV file using the stdlib `wave` module.
`source` may be a filesystem path or a binary file-like object.
Returns:
(wav, sr) where wav has shape [C, T] and is float32 in [-1, 1].
"""
if hasattr(source, "read"):
opened = wave.open(source, "rb")
else:
opened = wave.open(str(source), "rb")
with opened as wf:
n_channels = wf.getnchannels()
sr = wf.getframerate()
sampwidth = wf.getsampwidth()
n_frames = wf.getnframes()
raw = wf.readframes(n_frames)
if sampwidth == 1:
data = np.frombuffer(raw, dtype=np.uint8).astype(np.float32)
data = (data - 128.0) / 128.0
elif sampwidth == 2:
data = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0
elif sampwidth == 3:
bytes_ = np.frombuffer(raw, dtype=np.uint8).reshape(-1, 3)
as_int32 = (
bytes_[:, 0].astype(np.int32)
| (bytes_[:, 1].astype(np.int32) << 8)
| (bytes_[:, 2].astype(np.int32) << 16)
)
as_int32 = np.where(as_int32 >= (1 << 23), as_int32 - (1 << 24), as_int32)
data = as_int32.astype(np.float32) / float(1 << 23)
elif sampwidth == 4:
data = np.frombuffer(raw, dtype=np.int32).astype(np.float32) / float(1 << 31)
else:
raise ValueError(f"Unsupported WAV sample width: {sampwidth} bytes")
data = data.reshape(-1, n_channels)
return torch.from_numpy(np.ascontiguousarray(data.T)), sr
def _read_non_wav_file(source: str | Path | IO[bytes]) -> tuple[torch.Tensor, int]:
"""Load a non-WAV audio file using `soundfile`.
`source` may be a filesystem path or a binary file-like object (e.g. an
``io.BytesIO`` of an uploaded file), since libsndfile reads either.
Returns:
(wav, sr) where wav has shape [C, T] and is float32 in [-1, 1].
"""
try:
import soundfile as sf
except ImportError as e:
raise ImportError(
"soundfile is required to read non-WAV audio files. "
"Install with: `pip install soundfile` or `uvx --with soundfile`"
) from e
target = str(source) if isinstance(source, (str, Path)) else source
data, sample_rate = sf.read(target, dtype="float32")
if data.ndim == 1:
data = data[:, None]
wav = torch.from_numpy(np.ascontiguousarray(data.T))
return wav, sample_rate
def resample(
waveform: torch.Tensor,
orig_freq: int,
new_freq: int,
) -> torch.Tensor:
"""Sinc resampler via julius `resample_frac`. Operates along the last dim."""
if orig_freq == new_freq:
return waveform
return resample_frac(waveform, int(orig_freq), int(new_freq))
def load_audio(path: str | Path, target_sr: int = 16000) -> torch.Tensor:
"""Load an audio file and return a mono float32 tensor at target_sr.
PCM WAV files are read with the stdlib `wave` module. Other formats (mp3,
flac, ogg, m4a, …) are decoded via `soundfile`. Dispatch is by content, not
file extension, so misnamed files (e.g. an MP3 upload saved as .wav) still
load.
Returns:
Tensor of shape [1, T] at target_sr.
"""
filepath = Path(path)
try:
wav, sr = _read_wav_file(str(filepath))
except (wave.Error, EOFError):
wav, sr = _read_non_wav_file(str(filepath))
if wav.shape[0] > 1:
wav = wav.mean(dim=0, keepdim=True)
if sr != target_sr:
wav = resample(wav, sr, target_sr)
return wav
|