XiaoyuMa94/Music_Genre_Classification
1
1"""2Music Genre Classifier — Streamlit demo app3Deploy on Hugging Face Spaces (free CPU tier is enough).4 5Pure PyTorch + timm — NO fastai needed at inference, so there is no6pickle/version coupling with the training environment. The model file is7produced by the "plain export" cell in README.md Step 1a (it saves the8weights + head spec + genre vocab, and numerically verifies the round-trip9inside Kaggle before saving).10 11Required file: genre_model_plain.pt12Optional files: confusion_matrix.png (methodology page)13 per_class_f1.csv (methodology page)14 embeddings.npz (find-similar page: keys 'emb', 'names')15 samples/<Genre>/<clip>.mp3 (one-click demo clips)16 17Spectrogram settings below MUST match how training spectrograms were made.18"""19 20import glob21import io22import os23from pathlib import Path24 25import numpy as np26import streamlit as st27import torch28import torch.nn as nn29import torch.nn.functional as F30 31# ---------------------------------------------------------------- settings32SR = 2205033N_MELS = 12834N_FFT = 204835HOP = 51236CLIP_SECONDS = 30 # length used per training spectrogram37MAX_DECODE_SECONDS = 300 # decode up to this much, then keep the center38IMG_SIZE = 224 # training used Resize(224, method=Squish)39IMAGENET_MEAN = torch.tensor([0.485, 0.456, 0.406])[:, None, None]40IMAGENET_STD = torch.tensor([0.229, 0.224, 0.225])[:, None, None]41 42_pts = sorted(glob.glob("genre_model_plain*.pt"))43MODEL_PATH = _pts[0] if _pts else "genre_model_plain.pt"44SAMPLES_DIR = Path("samples")45 46st.set_page_config(page_title="Music Genre Classifier", page_icon="🎵", layout="wide")47 48 49# ---------------------------------------------------------------- model50class ConcatPool(nn.Module):51 """fastai's AdaptiveConcatPool2d: [max-pool, avg-pool] concatenated."""52 def forward(self, x):53 return torch.cat([F.adaptive_max_pool2d(x, 1),54 F.adaptive_avg_pool2d(x, 1)], dim=1)55 56 57class FeatureBody(nn.Module):58 """Run the timm backbone the way fastai's TimmBody does:59 forward_features() only — skip timm's own head norm/pool."""60 def __init__(self, m):61 super().__init__()62 self.m = m63 64 def forward(self, x):65 return self.m.forward_features(x)66 67 68def build_head(spec):69 """Rebuild the classifier head from the exported layer spec."""70 layers = []71 for s in spec:72 kind = s[0]73 if kind == "concat_pool":74 layers.append(ConcatPool())75 elif kind == "flatten":76 layers.append(nn.Flatten(1))77 elif kind == "bn":78 layers.append(nn.BatchNorm1d(s[1]))79 elif kind == "dropout":80 layers.append(nn.Dropout(s[1]))81 elif kind == "linear":82 layers.append(nn.Linear(s[1], s[2], bias=s[3]))83 elif kind == "relu":84 layers.append(nn.ReLU(inplace=True))85 else:86 raise ValueError(f"unknown head layer: {kind}")87 return nn.Sequential(*layers)88 89 90@st.cache_resource(show_spinner="Loading model…")91def load_model():92 import timm93 d = torch.load(MODEL_PATH, map_location="cpu", weights_only=False)94 body = timm.create_model(d["arch"], pretrained=False, num_classes=0)95 body.load_state_dict(d["body_sd"])96 head = build_head(d["head_spec"])97 head.load_state_dict(d["head_sd"])98 model = nn.Sequential(FeatureBody(body), head).eval()99 return model, list(d["vocab"])100 101 102@st.cache_resource103def load_embeddings():104 if not os.path.exists("embeddings.npz"):105 return None106 d = np.load("embeddings.npz", allow_pickle=True)107 emb = d["emb"].astype(np.float32)108 emb /= np.linalg.norm(emb, axis=1, keepdims=True) + 1e-8109 return emb, list(d["names"])110 111 112# ---------------------------------------------------------------- audio -> image113def _decode_with_av(audio_bytes: bytes):114 """Fallback decoder (m4a/aac and friends) using PyAV's bundled FFmpeg."""115 import av116 container = av.open(io.BytesIO(audio_bytes))117 resampler = av.AudioResampler(format="s16", layout="mono", rate=SR)118 chunks, total, limit = [], 0, SR * MAX_DECODE_SECONDS119 for frame in container.decode(audio=0):120 out = resampler.resample(frame)121 for f in (out if isinstance(out, list) else [out]):122 arr = f.to_ndarray().astype(np.float32).ravel() / 32768.0123 chunks.append(arr)124 total += arr.size125 if total >= limit:126 break127 container.close()128 if not chunks:129 return None130 return np.concatenate(chunks)[:limit]131 132 133def _center_crop(y):134 """Keep the middle CLIP_SECONDS — training clips are mid-song excerpts,135 so classifying a song's intro would be a train/serve mismatch."""136 n = SR * CLIP_SECONDS137 if y is not None and len(y) > n:138 start = (len(y) - n) // 2139 y = y[start:start + n]140 return y141 142 143def load_audio(audio_bytes: bytes):144 """Decode to mono float32 @ SR. librosa/soundfile first, PyAV fallback."""145 import librosa146 try:147 y, _ = librosa.load(io.BytesIO(audio_bytes), sr=SR, mono=True,148 duration=MAX_DECODE_SECONDS)149 return _center_crop(y)150 except Exception:151 try:152 return _center_crop(_decode_with_av(audio_bytes))153 except Exception:154 return None155 156 157def audio_to_melspec(audio_bytes: bytes):158 """Decode audio -> mel-spectrogram (dB), same recipe as training."""159 import librosa160 y = load_audio(audio_bytes)161 if y is None or len(y) < SR: # decode failure, or under 1 s162 return None163 mel = librosa.feature.melspectrogram(y=y, sr=SR, n_mels=N_MELS,164 n_fft=N_FFT, hop_length=HOP)165 return librosa.power_to_db(mel, ref=np.max)166 167 168def melspec_to_image(mel_db):169 """dB mel-spectrogram -> PIL image, EXACTLY as the training notebook:170 per-clip min-max to uint8 grayscale, flipped so low freqs are at the171 bottom. RGB conversion mirrors what the training image loader did."""172 from PIL import Image173 x = ((mel_db - mel_db.min()) / (mel_db.max() - mel_db.min() + 1e-9)174 * 255).astype(np.uint8)175 return Image.fromarray(np.flipud(x)).convert("RGB")176 177 178def preprocess(img):179 """PIL RGB image -> normalized 1x3x224x224 tensor (Squish resize +180 ImageNet stats, matching the training DataBlock)."""181 img = img.resize((IMG_SIZE, IMG_SIZE))182 x = torch.from_numpy(np.array(img)).permute(2, 0, 1).float() / 255.0183 return ((x - IMAGENET_MEAN) / IMAGENET_STD).unsqueeze(0)184 185 186# ---------------------------------------------------------------- inference187def predict(model, vocab, xb):188 with torch.no_grad():189 probs = torch.softmax(model(xb), dim=1)[0].cpu().numpy()190 probs = dict(zip(vocab, probs.tolist()))191 pred = max(probs, key=probs.get)192 return pred, probs193 194 195def find_cam_layer(body):196 body = getattr(body, "m", body) # unwrap FeatureBody197 if hasattr(body, "stages"): # ConvNeXt and many timm models198 return body.stages[-1]199 if hasattr(body, "layer4"): # ResNet family200 return body.layer4201 return None202 203 204def gradcam(model, xb):205 """Grad-CAM heatmap in [0,1] over the input, or None."""206 try:207 target = find_cam_layer(model[0])208 if target is None:209 return None210 acts, grads = [], []211 h1 = target.register_forward_hook(lambda m, i, o: acts.append(o))212 h2 = target.register_full_backward_hook(213 lambda m, gi, go: grads.append(go[0]))214 try:215 out = model(xb)216 out[0, out.argmax(dim=1)].backward()217 finally:218 h1.remove(); h2.remove()219 a, g = acts[0][0], grads[0][0]220 w = g.mean(dim=(1, 2), keepdim=True)221 cam = torch.relu((w * a).sum(0)).detach().cpu().numpy()222 return (cam - cam.min()) / (cam.max() - cam.min() + 1e-8)223 except Exception:224 return None225 226 227def overlay_cam(img, cam):228 import matplotlib.cm as cm229 from PIL import Image230 heat = Image.fromarray(231 (cm.jet(cam)[:, :, :3] * 255).astype(np.uint8)).resize(img.size)232 return Image.blend(img.convert("RGB"), heat, alpha=0.35)233 234 235def embed_clip(model, xb):236 """Concat-pooled backbone features (matches the README embeddings cell)."""237 with torch.no_grad():238 f = model[0](xb)239 v = ConcatPool()(f).flatten().cpu().numpy().astype(np.float32)240 return v / (np.linalg.norm(v) + 1e-8)241 242 243# ---------------------------------------------------------------- ui helpers244def list_samples():245 out = {}246 if SAMPLES_DIR.exists():247 for gdir in sorted(SAMPLES_DIR.iterdir()):248 if gdir.is_dir():249 clips = sorted(gdir.glob("*.mp3")) + sorted(gdir.glob("*.wav"))250 if clips:251 out[gdir.name] = clips252 return out253 254 255def get_audio_bytes():256 st.markdown("Upload a clip (mp3/wav, a few seconds is enough)")257 258 uploaded = st.file_uploader(259 "Choose audio file",260 type=["mp3", "wav", "ogg", "flac", "m4a"],261 accept_multiple_files=False,262 label_visibility="collapsed"263 )264 265 if uploaded is not None:266 bytes_data = bytes(uploaded.getvalue())267 if bytes_data:268 return bytes_data, uploaded.name269 270 samples = list_samples()271 if samples:272 st.caption("…or try a bundled sample:")273 c1, c2 = st.columns(2)274 genre = c1.selectbox("Genre", list(samples))275 clip = c2.selectbox("Clip", samples[genre], format_func=lambda p: p.name)276 if st.button("Use this sample"):277 return clip.read_bytes(), f"{genre}/{clip.name}"278 return None, None279 280 281# ---------------------------------------------------------------- pages282def page_classify(model, vocab):283 st.header("🎵 Classify a clip")284 audio, label = get_audio_bytes()285 if audio is None:286 st.info("Upload a clip or pick a sample to get started.")287 return288 289 st.audio(audio)290 with st.spinner("Analyzing…"):291 mel_db = audio_to_melspec(audio)292 if mel_db is None:293 st.error("Couldn't decode this file (or it's under ~3 seconds). "294 "Try mp3, wav, m4a, ogg, or flac.")295 return296 img = melspec_to_image(mel_db)297 xb = preprocess(img)298 pred, probs = predict(model, vocab, xb)299 cam = gradcam(model, xb)300 301 left, right = st.columns([1.2, 1])302 with left:303 st.subheader("What the model saw")304 st.image(img, caption="Mel-spectrogram", use_container_width=True)305 if cam is not None:306 st.image(overlay_cam(img, cam),307 caption="Grad-CAM — regions that drove the prediction",308 use_container_width=True)309 with right:310 st.subheader("Prediction")311 st.metric("Genre", pred, f"{probs[pred]*100:.1f}% confidence")312 top3 = sorted(probs.items(), key=lambda kv: -kv[1])[:3]313 st.write(" · ".join(f"**{g}** {p*100:.1f}%" for g, p in top3))314 st.subheader("All genres")315 st.bar_chart(probs, horizontal=True)316 317 318def page_similar(model):319 st.header("🔎 Find similar tracks")320 data = load_embeddings()321 if data is None:322 st.info("Add `embeddings.npz` (see README export snippet) to enable "323 "similar-track search.")324 return325 ref_emb, names = data326 327 audio, _ = get_audio_bytes()328 if audio is None:329 return330 st.audio(audio)331 with st.spinner("Embedding…"):332 mel_db = audio_to_melspec(audio)333 if mel_db is None:334 st.error("Clip too short.")335 return336 v = embed_clip(model, preprocess(melspec_to_image(mel_db)))337 sims = ref_emb @ v338 order = np.argsort(-sims)[:10]339 st.subheader("Closest tracks in the reference set")340 for i in order:341 st.write(f"**{names[i]}** — similarity {sims[i]:.3f}")342 343 344def page_methodology():345 st.header("📊 How it works")346 st.markdown(347 f"""348**Pipeline.** Audio → mel-spectrogram ({N_MELS} mels, {N_FFT} FFT,349hop {HOP}, {SR} Hz) → rendered as a grayscale image → **ConvNeXt-Tiny**350(transfer learning from ImageNet-22k, fine-tuned) → genre probabilities.351 352**Dataset.** [FMA — Free Music Archive](https://github.com/mdeff/fma),35316 top-level genres, official train/validation/test split.354 355**Why spectrograms?** They turn audio into images, letting us reuse356powerful pretrained vision models — genre cues (rhythm, timbre,357instrumentation) show up as visual texture.358 """)359 c1, c2 = st.columns(2)360 with c1:361 if os.path.exists("confusion_matrix.png"):362 st.image("confusion_matrix.png", caption="Confusion matrix")363 else:364 st.info("Add `confusion_matrix.png` to show the confusion matrix.")365 with c2:366 if os.path.exists("per_class_f1.csv"):367 import pandas as pd368 st.dataframe(pd.read_csv("per_class_f1.csv"),369 use_container_width=True)370 else:371 st.info("Add `per_class_f1.csv` to show per-genre F1 scores.")372 373 374# ---------------------------------------------------------------- main375def main():376 st.sidebar.title("Music Genre Classifier")377 page = st.sidebar.radio("Page", ["Classify", "Find similar", "How it works"])378 st.sidebar.markdown("---")379 st.sidebar.caption("ConvNeXt on mel-spectrograms · FMA dataset · "380 "PyTorch + Streamlit")381 382 if not os.path.exists(MODEL_PATH):383 st.error(f"Model file `{MODEL_PATH}` not found. Run the plain-export "384 "cell from README.md Step 1a in your Kaggle notebook and "385 "put `genre_model_plain.pt` next to app.py.")386 return387 model, vocab = load_model()388 389 if page == "Classify":390 page_classify(model, vocab)391 elif page == "Find similar":392 page_similar(model)393 else:394 page_methodology()395 396 397if __name__ == "__main__":398 main()399