Update app.py

app.py

@@ -2,33 +2,39 @@ import gradio as gr
 import librosa
 import numpy as np
 import torch
-from transformers import pipeline, AutoModelForAudioClassification, AutoFeatureExtractor
-from tqdm import tqdm
+from transformers import AutoModelForAudioClassification, AutoFeatureExtractor, pipeline
 import tempfile
-import json
 import soundfile as sf
+import json
 
-# ==== 参数 ====
 SAMPLE_RATE = 16000
-WINDOW = 5
-STEP = 2
+WINDOW = 10
+STEP = 5
 MUSIC_THRESHOLD = 0.4
 VOICE_THRESHOLD = 0.3
 MIN_SING_DURATION = 8
 
-# ==== 模型加载 ====
+# === 模型加载 ===
 music_model_id = "AI-Music-Detection/ai_music_detection_large_60s"
 music_feature_extractor = AutoFeatureExtractor.from_pretrained("MIT/ast-finetuned-audioset-10-10-0.4593")
 music_model = AutoModelForAudioClassification.from_pretrained(music_model_id)
-music_pipe = pipeline(
-    task="audio-classification",
-    model=music_model,
-    feature_extractor=music_feature_extractor
-)
-voice_pipe = pipeline(
-    "audio-classification",
-    model="superb/hubert-large-superb-sid"
-)
+voice_pipe = pipeline("audio-classification", model="superb/hubert-large-superb-sid")
+
+def predict_music_score(snippet):
+    """
+    直接手动跑 feature_extractor + model
+    避免 pipeline 自动切片问题
+    """
+    inputs = music_feature_extractor(snippet, sampling_rate=SAMPLE_RATE, return_tensors="pt", truncation=True, padding="max_length")
+    with torch.no_grad():
+        outputs = music_model(**inputs)
+        scores = torch.nn.functional.softmax(outputs.logits, dim=-1)[0]
+        labels = music_model.config.id2label
+        label_scores = {labels[i].lower(): float(scores[i]) for i in range(len(scores))}
+    # 找 music 或 singing 相关标签
+    music_score = max([v for k, v in label_scores.items() if "music" in k or "sing" in k] or [0])
+    return music_score
+
 
 def detect_singing(audio_path):
     wav, sr = librosa.load(audio_path, sr=SAMPLE_RATE)
@@ -39,27 +45,22 @@ def detect_singing(audio_path):
         end = start + WINDOW
         snippet = wav[int(start * SAMPLE_RATE):int(end * SAMPLE_RATE)]
 
-        # === 修复：AST模型要求固定60秒输入 ===
+        # 模型输入安全长度
         max_len = SAMPLE_RATE * 60
-        if len(snippet) < max_len:
-            pad = np.zeros(max_len)
-            pad[:len(snippet)] = snippet
-            snippet = pad
-        elif len(snippet) > max_len:
+        if len(snippet) < SAMPLE_RATE * 3:  # 过短片段跳过
+            continue
+        if len(snippet) > max_len:
             snippet = snippet[:max_len]
 
-        # === 音乐检测 ===
-        music_pred = music_pipe(snippet, sampling_rate=SAMPLE_RATE)
-        music_score = max([p['score'] for p in music_pred if 'music' in p['label'].lower()] or [0])
+        music_score = predict_music_score(snippet)
 
-        # === 人声检测 ===
         voice_pred = voice_pipe(snippet, sampling_rate=SAMPLE_RATE)
         voice_score = max([p['score'] for p in voice_pred if 'speech' in p['label'].lower()] or [0])
 
         if music_score > MUSIC_THRESHOLD and voice_score > VOICE_THRESHOLD:
             results.append((float(start), float(end)))
 
-    # === 合并连续片段 ===
+    # 合并连续窗口
     merged = []
     for seg in results:
         if not merged or seg[0] > merged[-1][1]:
@@ -74,8 +75,7 @@ def analyze_audio(file):
     if file is None:
         return "请上传音频文件", None
 
-    audio_path = file  # type="filepath" 返回的是路径字符串
-
+    audio_path = file
     with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
         data, sr = librosa.load(audio_path, sr=SAMPLE_RATE)
         sf.write(tmp.name, data, sr)
@@ -91,7 +91,6 @@ def analyze_audio(file):
     return f"检测到 {len(segments)} 段唱歌片段", json_output
 
 
-# ==== Gradio UI ====
 with gr.Blocks(title="🎵 Singing Segment Detector") as demo:
     gr.Markdown("# 🎤 自动识别唱歌片段\n上传音频文件（从视频提取后），返回检测到的唱歌时间段 JSON。")
     audio_in = gr.Audio(type="filepath", label="上传音频文件（WAV）")