fffiloni's picture
Update app.py
048ef8e
raw
history blame
4.62 kB
import gradio as gr
import io
import numpy as np
from PIL import Image
import pydub
from scipy.io import wavfile
import torch
import torchaudio
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("-i", "--input", help="Input file to process, anything that FFMPEG supports, but wav and mp3 are recommended")
parser.add_argument("-o", "--output", help="Output Image")
parser.add_argument("-m", "--maxvol", default=100, help="Max Volume, 255 for identical results")
parser.add_argument("-p", "--powerforimage", default=0.25, help="Power for Image")
parser.add_argument("-n", "--nmels", default=512, help="n_mels to use for Image, basically width. Higher = more fidelity")
args = parser.parse_args()
def spectrogram_image_from_wav(wav_bytes: io.BytesIO, max_volume: float = 50, power_for_image: float = 0.25, ms_duration: int = 5119) -> Image.Image:
"""
Generate a spectrogram image from a WAV file.
"""
# Read WAV file from bytes
sample_rate, waveform = wavfile.read(wav_bytes)
#sample_rate = 44100 # [Hz]
clip_duration_ms = ms_duration # [ms]
bins_per_image = 512
n_mels = int(args.nmels)
mel_scale = True
# FFT parameters
window_duration_ms = 100 # [ms]
padded_duration_ms = 400 # [ms]
step_size_ms = 10 # [ms]
# Derived parameters
num_samples = int(512 / float(bins_per_image) * clip_duration_ms) * sample_rate
n_fft = int(padded_duration_ms / 1000.0 * sample_rate)
hop_length = int(step_size_ms / 1000.0 * sample_rate)
win_length = int(window_duration_ms / 1000.0 * sample_rate)
# Compute spectrogram from waveform
Sxx = spectrogram_from_waveform(
waveform=waveform,
sample_rate=sample_rate,
n_fft=n_fft,
hop_length=hop_length,
win_length=win_length,
mel_scale=mel_scale,
n_mels=n_mels,
)
# Convert spectrogram to image
image = image_from_spectrogram(Sxx, max_volume=max_volume, power_for_image=power_for_image)
return image
def spectrogram_from_waveform(
waveform: np.ndarray,
sample_rate: int,
n_fft: int,
hop_length: int,
win_length: int,
mel_scale: bool = True,
n_mels: int = 512,
) -> np.ndarray:
"""
Compute a spectrogram from a waveform.
"""
spectrogram_func = torchaudio.transforms.Spectrogram(
n_fft=n_fft,
power=None,
hop_length=hop_length,
win_length=win_length,
)
waveform_tensor = torch.from_numpy(waveform.astype(np.float32)).reshape(1, -1)
Sxx_complex = spectrogram_func(waveform_tensor).numpy()[0]
Sxx_mag = np.abs(Sxx_complex)
if mel_scale:
mel_scaler = torchaudio.transforms.MelScale(
n_mels=n_mels,
sample_rate=sample_rate,
f_min=0,
f_max=10000,
n_stft=n_fft // 2 + 1,
norm=None,
mel_scale="htk",
)
Sxx_mag = mel_scaler(torch.from_numpy(Sxx_mag)).numpy()
return Sxx_mag
def image_from_spectrogram(
data: np.ndarray,
max_volume: float = 50,
power_for_image: float = 0.25
) -> Image.Image:
data = np.power(data, power_for_image)
data = data / (max_volume / 255)
data = 255 - data
data = data[::-1]
image = Image.fromarray(data.astype(np.uint8))
return image
def spectrogram_image_from_file(filename, max_volume: float = 50, power_for_image: float = 0.25) -> Image.Image:
"""
Generate a spectrogram image from an MP3 file.
"""
max_volume = int(max_volume)
power_for_image = float(args.powerforimage)
# Load MP3 file into AudioSegment object
audio = pydub.AudioSegment.from_file(filename)
# Convert to mono and set frame rate
audio = audio.set_channels(1)
audio = audio.set_frame_rate(44100)
length_in_ms = len(audio)
print("ORIGINAL AUDIO LENGTH IN MS:", length_in_ms)
# Extract first 5 seconds of audio data
audio = audio[:5119]
length_in_ms = len(audio)
print("CROPPED AUDIO LENGTH IN MS:", length_in_ms)
# Convert to WAV and save as BytesIO object
wav_bytes = io.BytesIO()
audio.export("clip.wav", format="wav")
audio.export(wav_bytes, format="wav")
wav_bytes.seek(0)
# Generate spectrogram image from WAV file
return spectrogram_image_from_wav(wav_bytes, max_volume=max_volume, power_for_image=power_for_image, ms_duration=length_in_ms)
def convert(audio):
image = spectrogram_image_from_file(audio, 50)
return image
gr.Interface(fn=convert, inputs=[gr.Audio(source="upload", type="filepath")], outputs=[gr.Image()]).launch()