Create main.py

main.py

@@ -0,0 +1,118 @@
+import gradio as gr
+import librosa
+import numpy as np
+from spleeter.separator import Separator
+import threading
+import queue
+import matplotlib.pyplot as plt
+import soundfile as sf
+
+# Load the song
+def load_audio(file_path):
+    signal, sr = librosa.load(file_path, sr=None)
+    return signal, sr
+
+# Spleeter separation
+def spleeter_separate(audio, sr):
+    separator = Separator('spleeter:5stems')
+    prediction = separator.separate(audio)
+    return prediction['vocals'], prediction['accompaniment'], prediction['bass'], prediction['drums'], prediction['other']
+
+# Function to control the volume of each stem
+def adjust_volume(stems, volumes):
+    adjusted_stems = []
+    for stem, volume in zip(stems, volumes):
+        adjusted_stem = stem * volume
+        adjusted_stems.append(adjusted_stem)
+    return adjusted_stems
+
+# Function to handle the separation and volume adjustment
+def process_audio(file_path, volumes, result_queue):
+    audio, sr = load_audio(file_path)
+    stems = spleeter_separate(audio, sr)
+    adjusted_stems = adjust_volume(stems, volumes)
+    result_queue.put(adjusted_stems)
+
+# Multithreaded processing
+def process_audio_threaded(file_path, volumes, result_queue):
+    thread = threading.Thread(target=process_audio, args=(file_path, volumes, result_queue))
+    thread.start()
+
+# Gradio interface
+def separate_audio(file, vocals, accompaniment, bass, drums, other):
+    file_path = file.name
+
+    volumes = [vocals, accompaniment, bass, drums, other]
+
+    result_queue = queue.Queue()
+    process_audio_threaded(file_path, volumes, result_queue)
+    adjusted_stems = result_queue.get()
+
+    # Reconstruct the audio with adjusted stems
+    reconstructed_audio = sum(adjusted_stems)
+
+    return reconstructed_audio.astype(np.float32)
+
+# Plot waveform and spectrogram
+def plot_audio(signal, sr, title):
+    plt.figure(figsize=(10, 6))
+    
+    # Waveform plot
+    plt.subplot(2, 1, 1)
+    librosa.display.waveplot(signal, sr=sr)
+    plt.title('Waveform')
+    
+    # Spectrogram plot
+    plt.subplot(2, 1, 2)
+    plt.specgram(signal, NFFT=2048, Fs=2, Fc=0, noverlap=128, cmap='inferno', sides='default', mode='default', scale='dB')
+    plt.title('Spectrogram')
+    
+    plt.tight_layout()
+    plt.suptitle(title, fontsize=14)
+    plt.show()
+
+iface = gr.Interface(
+    fn=separate_audio,
+    inputs=[
+        gr.inputs.Audio(label="Audio file"),
+        gr.inputs.Slider(0.0, 1.0, step=0.1, label="Vocals"),
+        gr.inputs.Slider(0.0, 1.0, step=0.1, label="Accompaniment"),
+        gr.inputs.Slider(0.0, 1.0, step=0.1, label="Bass"),
+        gr.inputs.Slider(0.0, 1.0, step=0.1, label="Drums"),
+        gr.inputs.Slider(0.0, 1.0, step=0.1, label="Other")
+    ],
+    outputs=gr.outputs.Audio(type="numpy", label="Separated Audio"),
+    title="Song Stem Separation",
+    description="Isolate vocals, accompaniment, bass, and drums of any song using the Spleeter model.",
+)
+
+# Add visualizations and support for different audio formats
+def visualize_and_save_audio(inputs, output):
+    audio_file = inputs[0]
+    vocals = inputs[1]
+    accompaniment = inputs[2]
+    bass = inputs[3]
+    drums = inputs[4]
+    other = inputs[5]
+    
+    # Load the original audio file
+    signal, sr = librosa.load(audio_file.name, sr=None)
+    
+    # Plot waveform and spectrogram of the original audio
+    plot_audio(signal, sr, "Original Audio")
+    
+    # Save the separated audio to a file
+    output_file = "separated_audio.wav"
+    sf.write(output_file, output.astype(np.float32), sr)
+    
+    # Load the separated audio file
+    separated_signal, separated_sr = librosa.load(output_file, sr=None)
+    
+    # Plot waveform and spectrogram of the separated audio
+    plot_audio(separated_signal, separated_sr, "Separated Audio")
+    
+    # Provide feedback to the user
+    print("Audio separated successfully! Separated audio saved to 'separated_audio.wav'.")
+
+iface.launch()
+