redboldcode/spectrogram-to-music

1"""2Audio processing tools to convert between spectrogram images and waveforms.3"""4import io5import typing as T6 7import numpy as np8from PIL import Image9import pydub10from scipy.io import wavfile11import torch12import torchaudio13 14 15def wav_bytes_from_spectrogram_image(image: Image.Image) -> T.Tuple[io.BytesIO, float]:16    """17    Reconstruct a WAV audio clip from a spectrogram image. Also returns the duration in seconds.18    """19 20    max_volume = 5021    power_for_image = 0.2522    Sxx = spectrogram_from_image(image, max_volume=max_volume, power_for_image=power_for_image)23 24    sample_rate = 44100  # [Hz]25    clip_duration_ms = 5000  # [ms]26 27    bins_per_image = 51228    n_mels = 51229 30    # FFT parameters31    window_duration_ms = 100  # [ms]32    padded_duration_ms = 400  # [ms]33    step_size_ms = 10  # [ms]34 35    # Derived parameters36    num_samples = int(image.width / float(bins_per_image) * clip_duration_ms) * sample_rate37    n_fft = int(padded_duration_ms / 1000.0 * sample_rate)38    hop_length = int(step_size_ms / 1000.0 * sample_rate)39    win_length = int(window_duration_ms / 1000.0 * sample_rate)40 41    samples = waveform_from_spectrogram(42        Sxx=Sxx,43        n_fft=n_fft,44        hop_length=hop_length,45        win_length=win_length,46        num_samples=num_samples,47        sample_rate=sample_rate,48        mel_scale=True,49        n_mels=n_mels,50        max_mel_iters=200,51        num_griffin_lim_iters=32,52    )53 54    wav_bytes = io.BytesIO()55    wavfile.write(wav_bytes, sample_rate, samples.astype(np.int16))56    wav_bytes.seek(0)57 58    duration_s = float(len(samples)) / sample_rate59 60    return wav_bytes, duration_s61 62 63def spectrogram_from_image(64    image: Image.Image, max_volume: float = 50, power_for_image: float = 0.2565) -> np.ndarray:66    """67    Compute a spectrogram magnitude array from a spectrogram image.68 69    TODO(hayk): Add image_from_spectrogram and call this out as the reverse.70    """71    # Convert to a numpy array of floats72    data = np.array(image).astype(np.float32)73 74    # Flip Y take a single channel75    data = data[::-1, :, 0]76 77    # Invert78    data = 255 - data79 80    # Rescale to max volume81    data = data * max_volume / 25582 83    # Reverse the power curve84    data = np.power(data, 1 / power_for_image)85 86    return data87 88 89def spectrogram_from_waveform(90    waveform: np.ndarray,91    sample_rate: int,92    n_fft: int,93    hop_length: int,94    win_length: int,95    mel_scale: bool = True,96    n_mels: int = 512,97) -> np.ndarray:98    """99    Compute a spectrogram from a waveform.100    """101 102    spectrogram_func = torchaudio.transforms.Spectrogram(103        n_fft=n_fft,104        power=None,105        hop_length=hop_length,106        win_length=win_length,107    )108 109    waveform_tensor = torch.from_numpy(waveform.astype(np.float32)).reshape(1, -1)110    Sxx_complex = spectrogram_func(waveform_tensor).numpy()[0]111 112    Sxx_mag = np.abs(Sxx_complex)113 114    if mel_scale:115        mel_scaler = torchaudio.transforms.MelScale(116            n_mels=n_mels,117            sample_rate=sample_rate,118            f_min=0,119            f_max=10000,120            n_stft=n_fft // 2 + 1,121            norm=None,122            mel_scale="htk",123        )124 125        Sxx_mag = mel_scaler(torch.from_numpy(Sxx_mag)).numpy()126 127    return Sxx_mag128 129 130def waveform_from_spectrogram(131    Sxx: np.ndarray,132    n_fft: int,133    hop_length: int,134    win_length: int,135    num_samples: int,136    sample_rate: int,137    mel_scale: bool = True,138    n_mels: int = 512,139    max_mel_iters: int = 200,140    num_griffin_lim_iters: int = 32,141    device: str = "cuda:0",142) -> np.ndarray:143    """144    Reconstruct a waveform from a spectrogram.145 146    This is an approximate inverse of spectrogram_from_waveform, using the Griffin-Lim algorithm147    to approximate the phase.148    """149    Sxx_torch = torch.from_numpy(Sxx).to(device)150 151    # TODO(hayk): Make this a class that caches the two things152 153    if mel_scale:154        mel_inv_scaler = torchaudio.transforms.InverseMelScale(155            n_mels=n_mels,156            sample_rate=sample_rate,157            f_min=0,158            f_max=10000,159            n_stft=n_fft // 2 + 1,160            norm=None,161            mel_scale="htk",162            max_iter=max_mel_iters,163        ).to(device)164 165        Sxx_torch = mel_inv_scaler(Sxx_torch)166 167    griffin_lim = torchaudio.transforms.GriffinLim(168        n_fft=n_fft,169        win_length=win_length,170        hop_length=hop_length,171        power=1.0,172        n_iter=num_griffin_lim_iters,173    ).to(device)174 175    waveform = griffin_lim(Sxx_torch).cpu().numpy()176 177    return waveform178 179 180def mp3_bytes_from_wav_bytes(wav_bytes: io.BytesIO) -> io.BytesIO:181    mp3_bytes = io.BytesIO()182    sound = pydub.AudioSegment.from_wav(wav_bytes)183    sound.export(mp3_bytes, format="mp3")184    mp3_bytes.seek(0)185    return mp3_bytes186 187def image_from_spectrogram(spectrogram: np.ndarray, max_volume: float = 50, power_for_image: float = 0.25) -> Image.Image:188    """189    Compute a spectrogram image from a spectrogram magnitude array.190    """191    # Apply the power curve192    data = np.power(spectrogram, power_for_image)193 194    # Rescale to 0-255195    data = data * 255 / max_volume196 197    # Invert198    data = 255 - data199 200    # Convert to a PIL image201    image = Image.fromarray(data.astype(np.uint8))202 203    # Flip Y204    image = image.transpose(Image.FLIP_TOP_BOTTOM)205 206    # Convert to RGB207    image = image.convert("RGB")208 209    return image