Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
+
 import io
-import typing as T
 
 import numpy as np
 from PIL import Image
@@ -8,36 +8,56 @@ import pydub
 from scipy.io import wavfile
 import torch
 import torchaudio
+import argparse
 
-def convert(audio):
-    # read uploaded file to wav
-    rate, data = wavfile.read(audio)
-    
-    # resample from 48000 to 44100
-    from scipy.signal import resample
-    data = resample(data, int(data.shape[0] * 44100 / 48000))
-    
-    # convert to mono
-    data = np.mean(data, axis=0)
-    
-    # convert to float32
-    data = data.astype(np.float32)
-    
-    # take a random 7 second slice of the audio
-    data = data[rate*7:rate*14]
-    
-    spectrogram = spectrogram_from_waveform(
-        waveform=data,
-        sample_rate=rate,
-        # width=768,
-        n_fft=8192,
-        hop_length=512,
-        win_length=8192,
+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,
     )
 
-    spec = image_from_spectrogram(spectrogram)
+    # Convert spectrogram to image
+    image = image_from_spectrogram(Sxx, max_volume=max_volume, power_for_image=power_for_image)
 
-    return spec
+    return image
 
 def spectrogram_from_waveform(
     waveform: np.ndarray,
@@ -80,28 +100,51 @@ def spectrogram_from_waveform(
     return Sxx_mag
 
 def image_from_spectrogram(
-    spectrogram: np.ndarray, max_volume: float = 50, power_for_image: float = 0.25
+        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:
     """
-    Compute a spectrogram image from a spectrogram magnitude array.
+    Generate a spectrogram image from an MP3 file.
     """
-    # Apply the power curve
-    data = np.power(spectrogram, power_for_image)
 
-    # Rescale to 0-255
-    data = data * 255 / max_volume
+    max_volume = int(args.maxvol)
+    power_for_image = float(args.powerforimage)
 
-    # Invert
-    data = 255 - data
+    # Load MP3 file into AudioSegment object
+    audio = pydub.AudioSegment.from_file(filename)
 
-    # Convert to a PIL image
-    image = Image.fromarray(data.astype(np.uint8))
+    # 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)
 
-    # Flip Y
-    image = image.transpose(Image.FLIP_TOP_BOTTOM)
+    # 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)
 
-    # Convert to RGB
-    image = image.convert("RGB")
+    # 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(filename)
 
     return image