Update app.py

app.py

@@ -2,49 +2,63 @@ import gradio as gr
 import librosa
 import numpy as np
 import torch
-from transformers import AutoFeatureExtractor, AutoModelForAudioClassification, AutoProcessor
+from transformers import AutoFeatureExtractor, AutoModelForAudioClassification
 import tempfile
 import soundfile as sf
 import json
 
+# === 参数设置 ===
 SAMPLE_RATE = 16000
-CHUNK_SIZE = 60
-STEP = 10
+CHUNK_SIZE = 60  # 模型输入60秒
+STEP = 10        # 滑动步长
 MUSIC_THRESHOLD = 0.5
 VOICE_THRESHOLD = 0.3
-MIN_SEG_DURATION = 8
+MIN_SEG_DURATION = 8  # 最小唱段长度(秒)
 
-# === 修正版 ===
+# === 模型加载 ===
+print("Loading models...")
+
+# 🎵 音乐检测模型（AST架构）
 music_model_id = "AI-Music-Detection/ai_music_detection_large_60s"
 music_extractor = AutoFeatureExtractor.from_pretrained("MIT/ast-finetuned-audioset-10-10-0.4593")
 music_model = AutoModelForAudioClassification.from_pretrained(music_model_id)
 
+# 🗣️ 语音活动检测模型（HuBERT）
 voice_model_id = "superb/hubert-large-superb-sid"
-voice_processor = AutoProcessor.from_pretrained(voice_model_id)
+voice_extractor = AutoFeatureExtractor.from_pretrained(voice_model_id)
 voice_model = AutoModelForAudioClassification.from_pretrained(voice_model_id)
 
+print("✅ Models loaded successfully.")
+
+
+# === 模型推理函数 ===
 def predict_music_score(wav):
+    """预测音乐片段概率"""
     wav = librosa.util.fix_length(wav, size=SAMPLE_RATE * CHUNK_SIZE)
-    inputs = music_processor(wav, sampling_rate=SAMPLE_RATE, return_tensors="pt", padding=True)
+    inputs = music_extractor(wav, sampling_rate=SAMPLE_RATE, return_tensors="pt", padding=True)
     with torch.no_grad():
         outputs = music_model(**inputs)
-        scores = torch.softmax(outputs.logits, dim=-1).squeeze()
-        music_score = float(scores[1]) if scores.numel() > 1 else float(scores[0])
-    return music_score
+        probs = torch.softmax(outputs.logits, dim=-1).squeeze()
+        score = float(probs[-1]) if probs.numel() > 1 else float(probs[0])
+    return score
+
 
 def predict_voice_score(wav):
+    """预测语音片段概率"""
     wav = librosa.util.fix_length(wav, size=SAMPLE_RATE * CHUNK_SIZE)
-    inputs = voice_processor(wav, sampling_rate=SAMPLE_RATE, return_tensors="pt", padding=True)
+    inputs = voice_extractor(wav, sampling_rate=SAMPLE_RATE, return_tensors="pt", padding=True)
     with torch.no_grad():
         outputs = voice_model(**inputs)
-        scores = torch.softmax(outputs.logits, dim=-1).squeeze()
-        voice_score = float(scores.mean())  # 简单平均
-    return voice_score
+        probs = torch.softmax(outputs.logits, dim=-1).squeeze()
+        score = float(probs.mean())  # 平均各类别概率
+    return score
+
 
 def detect_singing(audio_path):
-    wav, sr = librosa.load(audio_path, sr=SAMPLE_RATE)
+    """检测唱歌片段"""
+    wav, _ = librosa.load(audio_path, sr=SAMPLE_RATE)
     duration = len(wav) / SAMPLE_RATE
-    results = []
+    raw_segments = []
 
     for start in np.arange(0, max(0, duration - CHUNK_SIZE), STEP):
         end = start + CHUNK_SIZE
@@ -54,11 +68,11 @@ def detect_singing(audio_path):
         voice_score = predict_voice_score(snippet)
 
         if music_score > MUSIC_THRESHOLD and voice_score > VOICE_THRESHOLD:
-            results.append((float(start), float(end)))
+            raw_segments.append((float(start), float(end)))
 
-    # 合并连续窗口
+    # === 合并连续窗口 ===
     merged = []
-    for seg in results:
+    for seg in raw_segments:
         if not merged or seg[0] > merged[-1][1]:
             merged.append(list(seg))
         else:
@@ -67,12 +81,13 @@ def detect_singing(audio_path):
     return merged
 
 
-def analyze_audio(file):
-    if file is None:
-        return "请上传音频文件", None
+# === 主推理函数 ===
+def analyze_audio(file_path):
+    if file_path is None:
+        return "⚠️ 请上传音频文件", None
 
     with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
-        data, sr = librosa.load(file, sr=SAMPLE_RATE)
+        data, sr = librosa.load(file_path, sr=SAMPLE_RATE)
         sf.write(tmp.name, data, sr)
         segments = detect_singing(tmp.name)
 
@@ -86,16 +101,22 @@ def analyze_audio(file):
     return f"检测到 {len(segments)} 段唱歌片段", json_output
 
 
-with gr.Blocks(title="🎵 Singing Segment Detector (Plan A)") as demo:
+# === Gradio UI ===
+with gr.Blocks(title="🎵 Singing Segment Detector (Final)") as demo:
     gr.Markdown(
-        "# 🎤 高精度唱歌片段检测\n"
-        "使用 `AI-Music-Detection/ai_music_detection_large_60s` 模型。\n"
-        "将视频音频分块分析（60s输入），输出唱歌时间戳 JSON。"
+        """
+        # 🎤 唱歌片段自动检测器（AI-Music + HuBERT）
+        - 自动检测视频中的演唱时间段  
+        - 采用 `AI-Music-Detection/ai_music_detection_large_60s` + `HuBERT` 双模型融合  
+        - 输出每段的开始、结束时间与时长  
+        """
     )
-    audio_in = gr.Audio(type="filepath", label="上传音频文件（从视频抽取）")
-    btn = gr.Button("开始分析")
-    status = gr.Textbox(label="分析状态", interactive=False)
-    json_out = gr.Code(label="唱歌片段时间戳（JSON）", language="json")
-    btn.click(fn=analyze_audio, inputs=[audio_in], outputs=[status, json_out])
+
+    audio_input = gr.Audio(type="filepath", label="上传音频（从视频提取）")
+    run_btn = gr.Button("🚀 开始分析")
+    status_box = gr.Textbox(label="分析状态", interactive=False)
+    json_output = gr.Code(label="唱歌片段时间戳（JSON）", language="json")
+
+    run_btn.click(fn=analyze_audio, inputs=[audio_input], outputs=[status_box, json_output])
 
 demo.launch()