File size: 2,794 Bytes
31e78c9 0a766cb 31e78c9 67c4f45 31e78c9 67c4f45 31e78c9 67c4f45 31e78c9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | import gradio as gr
import torch
import torchaudio
import numpy as np
from transformers import ASTFeatureExtractor, ASTForAudioClassification
import os
MODEL_ID = "jananiramaseshan/ast-music-genre-classifier"
def load_assets():
try:
# Try loading from Hub
feature_extractor = ASTFeatureExtractor.from_pretrained(MODEL_ID)
model = ASTForAudioClassification.from_pretrained(MODEL_ID)
except Exception as e:
# Fallback to local directory if on a Space with model files
print(f"Loading from Hub failed: {e}. Trying local...")
feature_extractor = ASTFeatureExtractor.from_pretrained("./")
model = ASTForAudioClassification.from_pretrained("./")
return feature_extractor, model
extractor, model = load_assets()
id2label = model.config.id2label
def predict_genre(audio):
if audio is None:
return None
sr, data = audio
if data.dtype == np.int16:
data = data.astype(np.float32) / 32768.0
elif data.dtype == np.int32:
data = data.astype(np.float32) / 2147483648.0
waveform = torch.from_numpy(data.astype(np.float32))
if waveform.ndim > 1:
waveform = waveform.mean(dim=1, keepdim=True).T
else:
waveform = waveform.unsqueeze(0)
if sr != 16000:
resampler = torchaudio.transforms.Resample(sr, 16000)
waveform = resampler(waveform)
full_audio = waveform.squeeze(0).numpy()
target_samples = 16000 * 10 # 10 seconds segment
audio_len = len(full_audio)
if audio_len > target_samples:
offsets = [0, (audio_len - target_samples)//2, audio_len - target_samples]
else:
offsets = [0]
all_logits = []
for offset in offsets:
window = full_audio[offset:offset+target_samples]
if len(window) < target_samples:
window = np.pad(window, (0, target_samples - len(window)))
inputs = extractor(window, sampling_rate=16000, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
all_logits.append(outputs.logits)
avg_logits = torch.stack(all_logits).mean(dim=0)
probs = torch.softmax(avg_logits, dim=-1)[0]
results = {id2label[i]: float(probs[i]) for i in range(len(probs))}
return results
demo = gr.Interface(
fn=predict_genre,
inputs=gr.Audio(type="numpy", label="Upload Song"),
outputs=gr.Label(num_top_classes=5, label="Predicted Genre"),
title="🎵 Music Genre Classifier",
description="Upload a song file (WAV, MP3, etc.) to classify its genre using the Audio Spectrogram Transformer (AST).",
examples=[], # You can add example audio files here
flagging_mode="never"
)
if __name__ == "__main__":
demo.launch()
|