MeysamSh/SoundClassification
0
1import gradio as gr2import numpy as np3import librosa4import xgboost as xgb5import random6import subprocess7import tempfile8import os9import cv210import difflib11from sklearn.preprocessing import StandardScaler12from sklearn.pipeline import Pipeline13import torch14import torchvision.transforms as T15import torchvision.models as models16 17 18# --- Constants ---19SAMPLE_RATE = 1600020WINDOW_MS = 100 21WINDOW_SAMPLES = int(SAMPLE_RATE * WINDOW_MS / 1000)22N_MFCC = 1323SILENCE_EMOJI = "_"24MIN_SEC = 3.025MAX_SEC = 5.026 27# --- Lightweight pretrained visual backbone ---28device = torch.device("cpu")29 30# mobilenet = models.mobilenet_v2(weights=models.MobileNet_V2_Weights.DEFAULT)31mobilenet = models.mobilenet_v3_small(32 weights=models.MobileNet_V3_Small_Weights.DEFAULT33)34mobilenet = mobilenet.features # remove classifier35mobilenet.eval()36mobilenet.to(device)37 38# ImageNet normalization39video_transform = T.Compose([40 T.ToPILImage(),41 T.Resize((96, 96)), # small input for speed42 T.ToTensor(),43 T.Normalize(44 mean=[0.485, 0.456, 0.406],45 std=[0.229, 0.224, 0.225]46 )47])48 49 50 51def generate_challenge():52 length = random.randint(3, 5)53 seq = []54 for i in range(length):55 seq.append(str(random.choice([0, 1])))56 if i < length - 1:57 seq.append(SILENCE_EMOJI)58 # Return both the mission string and reset visibility to True59 mission = " ".join(seq)60 return mission, gr.update(visible=True, value=mission)61 62def hide_mission(audio_data):63 """Hides the mission textbox once the referee has recorded audio."""64 if audio_data is not None:65 return gr.update(visible=False)66 return gr.update(visible=True)67 68def post_process_video_sequence(69 preds,70 min_segment_frames=10,71 smoothing_window=10,72 background_class=273):74 """75 Post-process frame-level predictions into a clean symbol sequence.76 77 Steps:78 1. Temporal smoothing (majority vote).79 2. Remove very short segments.80 3. Collapse into final sequence.81 82 Args:83 preds: array of class predictions per frame84 min_segment_frames: minimum frames required to accept a symbol85 smoothing_window: neighborhood size for smoothing86 background_class: class index for background87 """88 89 if len(preds) == 0:90 return ""91 92 preds = [int(p) for p in preds]93 94 # -----------------------------------95 # 1. Majority vote smoothing96 # -----------------------------------97 half_w = smoothing_window // 298 smoothed = []99 100 for i in range(len(preds)):101 start = max(0, i - half_w)102 end = min(len(preds), i + half_w + 1)103 neighborhood = preds[start:end]104 smoothed.append(max(set(neighborhood), key=neighborhood.count))105 106 # -----------------------------------107 # 2. Segment compression108 # -----------------------------------109 segments = []110 current = smoothed[0]111 length = 1112 113 for p in smoothed[1:]:114 if p == current:115 length += 1116 else:117 segments.append((current, length))118 current = p119 length = 1120 segments.append((current, length))121 122 # -----------------------------------123 # 3. Filter short segments124 # -----------------------------------125 filtered = []126 for cls, length in segments:127 if cls != background_class and length < min_segment_frames:128 continue129 filtered.append(cls)130 131 # -----------------------------------132 # 4. Collapse duplicates133 # -----------------------------------134 final_seq = []135 for cls in filtered:136 if cls == background_class:137 continue138 if not final_seq or cls != final_seq[-1]:139 final_seq.append(str(cls))140 141 return "_".join(final_seq)142 143 144def post_process_to_emoji(preds, window_ms, min_silence_ms=200):145 """Processes raw AI output, smooths it, enforces silence gaps, and merges duplicates."""146 if len(preds) == 0: return ""147 148 ms_per_step = window_ms / 2149 min_silence_steps = int(min_silence_ms / ms_per_step)150 151 # 1. Majority Vote Smoothing (Temporal Filtering)152 # Reduces "flicker" where a single window might jump to a wrong class153 smoothed = []154 for i in range(len(preds)):155 start = max(0, i - 1)156 end = min(len(preds), i + 2)157 neighborhood = list(preds[start:end])158 smoothed.append(max(set(neighborhood), key=neighborhood.count))159 160 # 2. Silence Enforcement & Transition Logic161 # We only allow a change of class if the silence buffer is respected162 intermediate_sequence = []163 last_val = -1164 silence_count = 0165 166 for p in smoothed:167 p = int(p)168 if p == 2: # Silence Class169 silence_count += 1170 if last_val != 2:171 intermediate_sequence.append(2)172 last_val = 2173 else: # Sound Class (0 or 1)174 if last_val != p:175 # If we were in silence, check if the gap was long enough176 if last_val == -1 or (last_val == 2 and silence_count >= min_silence_steps):177 intermediate_sequence.append(p)178 last_val = p179 silence_count = 0180 # If we are jumping directly from 0 to 1 without silence, 181 # we ignore it or force silence (depending on game strictness)182 183 # 3. Final Merge (The "100110" -> "1010" logic)184 # This removes any accidental back-to-back duplicates185 # print("Intermediate Sequence (post-silence enforcement):", intermediate_sequence)186 final_output = []187 for val in intermediate_sequence:188 # print(f"Processing value: {val}")189 if val != 2:190 # Map back to emoji for silence or string for numbers191 # symbol = SILENCE_EMOJI if val == 2 else str(val)192 193 if not final_output or val != final_output[-1]:194 final_output.append(str(val))195 # print(f"Added {val} to final output {final_output}")196 197 final_output=[char+"_" for char in final_output]198 199 return "".join(final_output[:-1]) # Remove trailing silence if exists200 201def extract_features_sequence(audio_path,validate_duration=True):202 if audio_path is None: return None203 y, sr = librosa.load(audio_path, sr=SAMPLE_RATE, mono=True)204 if len(y) < WINDOW_SAMPLES:205 return None, f"Audio too short ({len(y)/SAMPLE_RATE:.1f}s), needs to be at least {WINDOW_MS/1000:.1f}s."206 elif validate_duration and len(y) > SAMPLE_RATE * 5: # Limit to 30 seconds for performance207 print(f"Audio too long ({len(y)/SAMPLE_RATE:.1f}s), truncating to 5s for feature extraction.")208 y = y[:SAMPLE_RATE * 5]209 210 hop = WINDOW_SAMPLES // 2 # 50% overlap for smoother sequence detection211 feats = []212 for start in range(0, len(y) - WINDOW_SAMPLES, hop):213 w = y[start:start + WINDOW_SAMPLES]214 mfcc = librosa.feature.mfcc(y=w, sr=sr, n_mfcc=N_MFCC, n_fft=512)215 feats.append(mfcc.mean(axis=1))216 return np.array(feats), "OK"217 218def train_player_model(a0, a1, a_silence, player_name):219 X0, msg0 = extract_features_sequence(a0, validate_duration=True)220 X1, msg1 = extract_features_sequence(a1, validate_duration=True)221 X_sil, msg_sil = extract_features_sequence(a_silence, validate_duration=True)222 223 if X0 is None: return None, f"{player_name} Source 0: {msg0}"224 if X1 is None: return None, f"{player_name} Source 1: {msg1}"225 if X_sil is None: return None, f"{player_name} Silence: {msg_sil}"226 227 X = np.vstack([X0, X1, X_sil])228 y = np.concatenate([np.zeros(len(X0)), np.ones(len(X1)), np.full(len(X_sil), 2)])229 230 print(f"{player_name} - Training model with {len(X)} samples: {len(X0)} Source 0, {len(X1)} Source 1, {len(X_sil)} Silence")231 232 model = Pipeline([233 ("scaler", StandardScaler()),234 ("clf", xgb.XGBClassifier(n_estimators=50, max_depth=3, objective='multi:softprob', num_class=3))235 ])236 model.fit(X, y)237 print(f"{player_name} model trained successfully with {len(X)} samples!")238 return model, "OK"239 240def play_game(target_display, ref_audio, p1_0, p1_1, p1_s, p2_0, p2_1, p2_s):241 # Validation and Training logic...242 m1, err1 = train_player_model(p1_0, p1_1, p1_s, "Player 1")243 if m1 is None: return f"### ❌ {err1}"244 245 m2, err2 = train_player_model(p2_0, p2_1, p2_s, "Player 2")246 if m2 is None: return f"### ❌ {err2}"247 248 if not ref_audio: return "### ⚠️ Referee recording missing!"249 250 X_ref, _ = extract_features_sequence(ref_audio, validate_duration=False)251 target_numeric = target_display.replace(" ", "").replace(SILENCE_EMOJI, "2")252 253 res1_emoji = post_process_to_emoji(m1.predict(X_ref), WINDOW_MS)254 res2_emoji = post_process_to_emoji(m2.predict(X_ref), WINDOW_MS)255 256 res1_num = res1_emoji.replace(SILENCE_EMOJI, "2")257 res2_num = res2_emoji.replace(SILENCE_EMOJI, "2")258 259 score1 = round(difflib.SequenceMatcher(None, target_numeric, res1_num).ratio() * 100, 1)260 score2 = round(difflib.SequenceMatcher(None, target_numeric, res2_num).ratio() * 100, 1)261 262 winner = "Player 1" if score1 > score2 else "Player 2"263 if score1 == score2: winner = "It's a Tie!"264 265 # Formatting results with Large Markdown266 return f"""267 # 🏁 BATTLE RESULTS268 269 ## 🎯 Mission Target: {target_display}270 271 ---272 ## 👤 Player 1 `{res1_emoji}` | **Accuracy:** `{score1}%`273 274 ## 👤 Player 2 `{res2_emoji}` | **Accuracy:** `{score2}%`275 276 ---277 # 🏆 WINNER: <span style="color: #ff4b4b; font-size: 40px;">{winner}</span>278 """279 280 281# =========================================================282# VIDEO SECTION283# =========================================================284 285def ensure_readable_video(input_path):286 """Re-encode video to MP4 to avoid WEBM/Opus issues."""287 if input_path is None:288 return None289 290 tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)291 tmp_path = tmp.name292 tmp.close()293 294 cmd = [295 "ffmpeg",296 "-y",297 "-i", input_path,298 "-an", # remove audio299 "-vcodec", "libx264",300 "-preset", "ultrafast",301 tmp_path302 ]303 304 try:305 subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)306 return tmp_path307 except:308 return input_path309 310 311def extract_video_features(video_path, max_frames=300):312 """Extract frame-level features from video."""313 if video_path is None:314 return None, "No video provided"315 316 video_path = ensure_readable_video(video_path)317 318 cap = cv2.VideoCapture(video_path)319 feats = []320 frame_count = 0321 322 while True:323 ret, frame = cap.read()324 if not ret or frame_count >= max_frames:325 break326 327 frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)328 tensor = video_transform(frame_rgb).unsqueeze(0).to(device)329 330 with torch.no_grad():331 feat_map = mobilenet(tensor)332 feat = torch.nn.functional.adaptive_avg_pool2d(feat_map, 1)333 feat = feat.view(-1).cpu().numpy()334 335 feats.append(feat)336 337 # frame = cv2.resize(frame, (64, 64))338 # frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)339 340 # # Basic color statistics341 # mean = frame_rgb.mean(axis=(0, 1))342 # std = frame_rgb.std(axis=(0, 1))343 # brightness = frame_rgb.mean()344 345 # feat = np.concatenate([mean, std, [brightness]])346 # feats.append(feat)347 348 frame_count += 1349 350 cap.release()351 352 if len(feats) == 0:353 return None, "No frames extracted"354 355 return np.array(feats), "OK"356 357 358def train_video_model(v0, v1, v_bg):359 X0, msg0 = extract_video_features(v0)360 X1, msg1 = extract_video_features(v1)361 Xbg, msgbg = extract_video_features(v_bg)362 363 if X0 is None: return None, f"Class 0 error: {msg0}"364 if X1 is None: return None, f"Class 1 error: {msg1}"365 if Xbg is None: return None, f"Background error: {msgbg}"366 367 print(f"Training video model with {len(X0)} frames for Class 0, {len(X1)} frames for Class 1, and {len(Xbg)} frames for Background.")368 369 X = np.vstack([X0, X1, Xbg])370 y = np.concatenate([371 np.zeros(len(X0)),372 np.ones(len(X1)),373 np.full(len(Xbg), 2)374 ])375 376 model = Pipeline([377 ("scaler", StandardScaler()),378 ("clf", xgb.XGBClassifier(379 n_estimators=50,380 max_depth=3,381 objective='multi:softprob',382 num_class=3383 ))384 ])385 386 model.fit(X, y)387 print("Video model trained successfully!")388 return model, "OK"389 390 391def decode_video_sequence(model, video_path):392 X, msg = extract_video_features(video_path)393 if X is None:394 return f"Error: {msg}"395 396 preds = model.predict(X)397 print(f"Raw frame-level predictions: {preds}")398 return post_process_video_sequence(preds)399 400 401def run_video_decoder(v0, v1, v_bg, test_video):402 model, msg = train_video_model(v0, v1, v_bg)403 if model is None:404 return f"❌ {msg}"405 406 result = decode_video_sequence(model, test_video)407 return f"### 🎬 Decoded Sequence: `{result}`"408 409 410# =========================================================411# GRADIO UI WITH DUAL TABS412# =========================================================413 414with gr.Blocks(theme=gr.themes.Soft()) as demo:415 416 with gr.Tabs():417 418 # =====================================419 # TAB 1 — AUDIO GAME (existing)420 # =====================================421 with gr.Tab("🎙️ Audio Sequence Battle"):422 423 hidden_target = gr.State("")424 425 with gr.Row():426 target_seq_ui = gr.Textbox(427 label="📢 Referee's Mission",428 interactive=False429 )430 refresh_btn = gr.Button("🔄 New Mission")431 432 demo.load(generate_challenge, outputs=[hidden_target, target_seq_ui])433 refresh_btn.click(generate_challenge, outputs=[hidden_target, target_seq_ui])434 435 with gr.Accordion("⚖️ Step 1: The Referee", open=True):436 ref_audio = gr.Audio(437 sources=["microphone"],438 type="filepath",439 label="Record the Mission"440 )441 ref_audio.change(hide_mission, inputs=ref_audio, outputs=target_seq_ui)442 443 with gr.Row():444 with gr.Column():445 gr.Markdown("### 👤 Player 1")446 p1_0 = gr.Audio(sources=["microphone"], type="filepath", label="Source 0")447 p1_1 = gr.Audio(sources=["microphone"], type="filepath", label="Source 1")448 p1_s = gr.Audio(sources=["microphone"], type="filepath", label="Silence")449 450 with gr.Column():451 gr.Markdown("### 👤 Player 2")452 p2_0 = gr.Audio(sources=["microphone"], type="filepath", label="Source 0")453 p2_1 = gr.Audio(sources=["microphone"], type="filepath", label="Source 1")454 p2_s = gr.Audio(sources=["microphone"], type="filepath", label="Silence")455 456 btn_fight = gr.Button("🔥 REVEAL WINNER", variant="primary")457 result_display = gr.Markdown("### Results will appear here")458 459 btn_fight.click(460 play_game,461 inputs=[hidden_target, ref_audio, p1_0, p1_1, p1_s, p2_0, p2_1, p2_s],462 outputs=result_display463 )464 465 466 # =====================================467 # TAB 2 — VIDEO DECODER468 # =====================================469 with gr.Tab("🎬 Video Frame Decoder"):470 471 gr.Markdown("## Train video symbols and decode frame-level sequence")472 473 with gr.Row():474 with gr.Column():475 gr.Markdown("### Training Samples")476 v0 = gr.Video(label="Class 0 video",format="mp4")477 v1 = gr.Video(label="Class 1 video",format="mp4")478 v_bg = gr.Video(label="Background video",format="mp4")479 480 with gr.Column():481 gr.Markdown("### Test Video")482 test_video = gr.Video(label="Video to decode",format="mp4")483 484 decode_btn = gr.Button("🎬 Decode Video", variant="primary")485 video_result = gr.Markdown("### Decoded result will appear here")486 487 decode_btn.click(488 run_video_decoder,489 inputs=[v0, v1, v_bg, test_video],490 outputs=video_result491 )492 493demo.launch()494 495 496 497 498# # --- Gradio UI ---499# with gr.Blocks(theme=gr.themes.Soft()) as demo:500# gr.Markdown("# 🎙️ The AI Sequence Battle")501 502# # Store the mission in a hidden state so we can still use it for scoring even when invisible503# hidden_target = gr.State("")504 505# with gr.Row():506# target_seq_ui = gr.Textbox(label="📢 Referee's Mission (Memorize this!)", interactive=False)507# refresh_btn = gr.Button("🔄 New Mission")508 509# # On load and on refresh, update both the UI and the State510# demo.load(generate_challenge, outputs=[hidden_target, target_seq_ui])511# refresh_btn.click(generate_challenge, outputs=[hidden_target, target_seq_ui])512 513# with gr.Accordion("⚖️ Step 1: The Referee", open=True):514# ref_audio = gr.Audio(sources=["microphone"], type="filepath", label="Record the Mission")515# # Trigger hiding when audio is recorded516# ref_audio.change(hide_mission, inputs=ref_audio, outputs=target_seq_ui)517 518# with gr.Row():519# with gr.Column():520# gr.Markdown("### 👤 Player 1 (3-5s samples)")521# p1_0 = gr.Audio(sources=["microphone"], type="filepath", label="Source 0")522# p1_1 = gr.Audio(sources=["microphone"], type="filepath", label="Source 1")523# p1_s = gr.Audio(sources=["microphone"], type="filepath", label="Silence 🤫")524# with gr.Column():525# gr.Markdown("### 👤 Player 2 (3-5s samples)")526# p2_0 = gr.Audio(sources=["microphone"], type="filepath", label="Source 0")527# p2_1 = gr.Audio(sources=["microphone"], type="filepath", label="Source 1")528# p2_s = gr.Audio(sources=["microphone"], type="filepath", label="Silence 🤫")529 530# btn_fight = gr.Button("🔥 REVEAL WINNER", variant="primary", size="lg")531 532# # Using Markdown for large, styled text results533# result_display = gr.Markdown("### Results will appear here after the battle!")534 535# btn_fight.click(536# play_game,537# inputs=[hidden_target, ref_audio, p1_0, p1_1, p1_s, p2_0, p2_1, p2_s],538# outputs=result_display539# )540 541# demo.launch()