Spaces:
Running
Running
| import gradio as gr | |
| import torch | |
| import numpy as np | |
| import librosa | |
| import time | |
| from src.utility import slice_songs | |
| from src.models import MusicCNN, MusicCRNN2D | |
| # Configuration | |
| DESIRED_SR = 22050 | |
| HOP_LENGTH = 512 | |
| DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| NUM_CLASSES = 10 | |
| # Model loading | |
| models = {"cnn": {}, "crnn": {}} | |
| cnn_model_paths = {1: "models/cnn/1s.pth", 3: "models/cnn/3s.pth", 5: "models/cnn/5s.pth", 10: "models/cnn/10s.pth"} | |
| crnn_model_paths = {1: "models/crnn/1s.pth", 3: "models/crnn/3s.pth", 5: "models/crnn/5s.pth", | |
| 10: "models/crnn/10s.pth"} | |
| def get_frames(slice_length): | |
| return int(slice_length * DESIRED_SR / HOP_LENGTH) | |
| # Load cnn models | |
| for slice_len, path in cnn_model_paths.items(): | |
| model = MusicCNN(num_classes=NUM_CLASSES, device=DEVICE) | |
| dummy_input = torch.randn(2, 1, 128, get_frames(slice_len)).to(DEVICE) | |
| _ = model(dummy_input) | |
| model.load_state_dict(torch.load(path, map_location=DEVICE)) | |
| model.to(DEVICE) | |
| model.eval() | |
| models["cnn"][slice_len] = model | |
| # Load crnn models | |
| for slice_len, path in crnn_model_paths.items(): | |
| model = MusicCRNN2D(num_classes=NUM_CLASSES, device=DEVICE) | |
| dummy_input = torch.randn(2, 1, 128, get_frames(slice_len)).to(DEVICE) | |
| _ = model(dummy_input) | |
| model.load_state_dict(torch.load(path, map_location=DEVICE)) | |
| model.to(DEVICE) | |
| model.eval() | |
| models["crnn"][slice_len] = model | |
| GENRE_LABELS = ["Blues", "Classical", "Country", "Disco", "HipHop", "Jazz", "Metal", "Pop", "Reggae", "Rock"] | |
| def predict_genre(audio_file, slice_length, architecture): | |
| slice_length = int(slice_length) | |
| start_time = time.time() | |
| y, sr = librosa.load(audio_file, sr=DESIRED_SR) | |
| target_length = int(np.ceil(len(y) / sr)) * sr | |
| if len(y) < target_length: | |
| y = np.pad(y, (0, target_length - len(y)), mode='constant') | |
| mel_spectrogram = librosa.feature.melspectrogram(y=y, sr=sr, n_fft=2048, hop_length=HOP_LENGTH, n_mels=128) | |
| mel_spectrogram_db = librosa.power_to_db(mel_spectrogram, ref=np.max) | |
| min_val, max_val = np.min(mel_spectrogram_db), np.max(mel_spectrogram_db) | |
| normalized_spectrogram = (mel_spectrogram_db - min_val) / ( | |
| max_val - min_val) if max_val - min_val > 0 else mel_spectrogram_db | |
| X_slices, _, _ = slice_songs([normalized_spectrogram], [0], ["temp"], sr=sr, hop_length=HOP_LENGTH, | |
| length_in_seconds=slice_length) | |
| X_slices = torch.tensor(X_slices, dtype=torch.float32).unsqueeze(1).to(DEVICE) | |
| model_used = models[architecture][slice_length] | |
| with torch.no_grad(): | |
| outputs = model_used(X_slices) | |
| probabilities = torch.nn.functional.softmax(outputs, dim=1).cpu().numpy() | |
| avg_probs = np.mean(probabilities, axis=0) | |
| genre_distribution = {GENRE_LABELS[i]: float(avg_probs[i]) for i in range(NUM_CLASSES)} | |
| inference_time = time.time() - start_time | |
| return genre_distribution, f"Inference Time: {inference_time:.2f} seconds" | |
| slice_length_dropdown = gr.Dropdown(choices=["1", "3", "5", "10"], value="1", label="Slice Length (seconds)") | |
| architecture_dropdown = gr.Dropdown(choices=["cnn", "crnn"], value="cnn", label="Model Architecture") | |
| demo = gr.Interface( | |
| fn=predict_genre, | |
| inputs=[gr.Audio(type="filepath", label="Upload Audio File"), slice_length_dropdown, architecture_dropdown], | |
| outputs=[gr.Label(num_top_classes=10, label="Genre Distribution"), gr.Textbox(label="Inference Time")], | |
| title="Music Genre Classifier", | |
| description="Upload an audio file, select a slice length and model architecture to predict its genre distribution." | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |