215 lines
6.6 KiB
Python
215 lines
6.6 KiB
Python
import subprocess as sp
|
|
from pathlib import Path
|
|
|
|
import json
|
|
import numpy as np
|
|
|
|
import torch
|
|
import torchaudio as ta
|
|
|
|
import typing as tp
|
|
|
|
import lameenc
|
|
import julius
|
|
|
|
from .utils import temp_filenames
|
|
|
|
def _read_info(path):
|
|
stdout_data = sp.check_output([
|
|
'ffprobe', "-loglevel", "panic",
|
|
str(path), '-print_format', 'json', '-show_format', '-show_streams'
|
|
])
|
|
return json.loads(stdout_data.decode('utf-8'))
|
|
|
|
class AudioFile:
|
|
def __init__(self, path: Path):
|
|
self.path = Path(path)
|
|
self._info = None
|
|
|
|
def __repr__(self):
|
|
features = [("path", self.path)]
|
|
features.append(("samplerate", self.samplerate()))
|
|
features.append(("channels", self.channels()))
|
|
features.append(("streams", len(self)))
|
|
features_str = ", ".join(f"{name}={value}" for name, value in features)
|
|
return f"AudioFile({features_str})"
|
|
|
|
@property
|
|
def info(self):
|
|
if self._info is None:
|
|
self._info = _read_info(self.path)
|
|
return self._info
|
|
|
|
@property
|
|
def duration(self):
|
|
return float(self.info['format']['duration'])
|
|
|
|
@property
|
|
def _audio_streams(self):
|
|
return [
|
|
index for index, stream in enumerate(self.info["streams"])
|
|
if stream["codec_type"] == "audio"
|
|
]
|
|
|
|
def __len__(self):
|
|
return len(self._audio_streams)
|
|
|
|
def channels(self, stream=0):
|
|
return int(self.info['streams'][self._audio_streams[stream]]['channels'])
|
|
|
|
def samplerate(self, stream=0):
|
|
return int(self.info['streams'][self._audio_streams[stream]]['sample_rate'])
|
|
|
|
def read(self,
|
|
seek_time=None,
|
|
duration=None,
|
|
streams=slice(None),
|
|
samplerate=None,
|
|
channels=None):
|
|
|
|
streams = np.array(range(len(self)))[streams]
|
|
single = not isinstance(streams, np.ndarray)
|
|
if single:
|
|
streams = [streams]
|
|
|
|
if duration is None:
|
|
target_size = None
|
|
query_duration = None
|
|
else:
|
|
target_size = int((samplerate or self.samplerate()) * duration)
|
|
query_duration = float((target_size + 1) / (samplerate or self.samplerate()))
|
|
|
|
with temp_filenames(len(streams)) as filenames:
|
|
command = ['ffmpeg', '-y']
|
|
command += ['-loglevel', 'panic']
|
|
if seek_time:
|
|
command += ['-ss', str(seek_time)]
|
|
command += ['-i', str(self.path)]
|
|
for stream, filename in zip(streams, filenames):
|
|
command += ['-map', f'0:{self._audio_streams[stream]}']
|
|
if query_duration is not None:
|
|
command += ['-t', str(query_duration)]
|
|
command += ['-threads', '1']
|
|
command += ['-f', 'f32le']
|
|
if samplerate is not None:
|
|
command += ['-ar', str(samplerate)]
|
|
command += [filename]
|
|
|
|
sp.run(command, check=True)
|
|
wavs = []
|
|
for filename in filenames:
|
|
wav = np.fromfile(filename, dtype=np.float32)
|
|
wav = torch.from_numpy(wav)
|
|
wav = wav.view(-1, self.channels()).t()
|
|
if channels is not None:
|
|
wav = convert_audio_channels(wav, channels)
|
|
if target_size is not None:
|
|
wav = wav[..., :target_size]
|
|
wavs.append(wav)
|
|
wav = torch.stack(wavs, dim=0)
|
|
if single:
|
|
wav = wav[0]
|
|
return wav
|
|
|
|
def i16_pcm(wav):
|
|
if wav.dtype.is_floating_point:
|
|
return (wav.clamp_(-1, 1) * (2**15 - 1)).short()
|
|
else:
|
|
return wav
|
|
|
|
|
|
def f32_pcm(wav):
|
|
if wav.dtype.is_floating_point:
|
|
return wav
|
|
else:
|
|
return wav.float() / (2**15 - 1)
|
|
|
|
|
|
def as_dtype_pcm(wav, dtype):
|
|
if wav.dtype.is_floating_point:
|
|
return f32_pcm(wav)
|
|
else:
|
|
return i16_pcm(wav)
|
|
|
|
def encode_mp3(wav, path, samplerate=44100, bitrate=320, quality=2, verbose=False):
|
|
C, T = wav.shape
|
|
wav = i16_pcm(wav)
|
|
encoder = lameenc.Encoder()
|
|
encoder.set_bit_rate(bitrate)
|
|
encoder.set_in_sample_rate(samplerate)
|
|
encoder.set_channels(C)
|
|
encoder.set_quality(quality)
|
|
if not verbose:
|
|
encoder.silence()
|
|
wav = wav.data.cpu()
|
|
wav = wav.transpose(0, 1).numpy()
|
|
mp3_data = encoder.encode(wav.tobytes())
|
|
mp3_data += encoder.flush()
|
|
with open(path, "wb") as f:
|
|
f.write(mp3_data)
|
|
|
|
|
|
def prevent_clip(wav, mode='rescale'):
|
|
if mode is None or mode == 'none':
|
|
return wav
|
|
assert wav.dtype.is_floating_point, "too late for clipping"
|
|
if mode == 'rescale':
|
|
wav = wav / max(1.01 * wav.abs().max(), 1)
|
|
elif mode == 'clamp':
|
|
wav = wav.clamp(-0.99, 0.99)
|
|
elif mode == 'tanh':
|
|
wav = torch.tanh(wav)
|
|
else:
|
|
raise ValueError(f"Invalid mode {mode}")
|
|
return wav
|
|
|
|
def save_audio(wav: torch.Tensor,
|
|
path: tp.Union[str, Path],
|
|
samplerate: int,
|
|
bitrate: int = 320,
|
|
clip: tp.Literal["rescale", "clamp", "tanh", "none"] = 'rescale',
|
|
bits_per_sample: tp.Literal[16, 24, 32] = 16,
|
|
as_float: bool = False,
|
|
preset: tp.Literal[2, 3, 4, 5, 6, 7] = 2):
|
|
|
|
wav = prevent_clip(wav, mode=clip)
|
|
path = Path(path)
|
|
suffix = path.suffix.lower()
|
|
if suffix == ".mp3":
|
|
encode_mp3(wav, path, samplerate, bitrate, preset, verbose=True)
|
|
elif suffix == ".wav":
|
|
if as_float:
|
|
bits_per_sample = 32
|
|
encoding = 'PCM_F'
|
|
else:
|
|
encoding = 'PCM_S'
|
|
ta.save(str(path), wav, sample_rate=samplerate,
|
|
encoding=encoding, bits_per_sample=bits_per_sample)
|
|
elif suffix == ".flac":
|
|
ta.save(str(path), wav, sample_rate=samplerate, bits_per_sample=bits_per_sample)
|
|
else:
|
|
raise ValueError(f"Invalid suffix for path: {suffix}")
|
|
|
|
return {
|
|
"waveform": wav.cpu().unsqueeze(0),
|
|
"sample_rate": samplerate
|
|
}
|
|
|
|
|
|
def convert_audio_channels(wav, channels=2):
|
|
*shape, src_channels, length = wav.shape
|
|
if src_channels == channels:
|
|
pass
|
|
elif channels == 1:
|
|
wav = wav.mean(dim=-2, keepdim=True)
|
|
elif src_channels == 1:
|
|
wav = wav.expand(*shape, channels, length)
|
|
elif src_channels >= channels:
|
|
wav = wav[..., :channels, :]
|
|
else:
|
|
raise ValueError('The audio file has less channels than requested but is not mono.')
|
|
return wav
|
|
|
|
def convert_audio(wav, from_samplerate, to_samplerate, channels) -> torch.Tensor:
|
|
wav = convert_audio_channels(wav, channels)
|
|
return julius.resample_frac(wav, from_samplerate, to_samplerate) |