Update app.py

app.py

@@ -43,7 +43,7 @@ def load_model():
 model, threshold = load_model()
 
 # 检测函数
-def detect(dataset, model):
+def detect(dataset):
     """进行音频伪造检测"""
     with torch.no_grad():
         for batch in dataset:
@@ -56,9 +56,19 @@ def detect(dataset, model):
                 'attention_mask': pf['attention_mask'].to(device)
             } for pf in batch['prompt_features']]
 
+            prompt_labels = batch['prompt_labels'].to(device)
             # 模型的前向传播逻辑 (需要补充具体实现)
+            outputs = model({
+                'main_features': main_features,
+                'prompt_features': prompt_features,
+                'prompt_labels': prompt_labels
+            })
+
+            avg_scores = outputs['avg_logits'].softmax(dim=-1) # [batch_size, 2]
+            deepfake_scores = avg_scores[:, 1].cpu() # [batch_size]
+            is_fake = True if deepfake_scores[0] > threshold else False
             # 假设 result 是模型返回的结果
-            result = {"is_fake": True, "confidence": 85.5}  # 示例返回值
+            result = {"is_fake": is_fake, "confidence": deepfake_scores[0]}  # 示例返回值
     return result
 
 # 音频伪造检测主函数
@@ -70,11 +80,12 @@ def audio_deepfake_detection(demonstrations, query_audio_path):
     :return: 检测结果
     """
     demonstration_paths = [audio[0] for audio in demonstrations if audio[0] is not None]
+    demonstration_labels = [audio[1] for audio in demonstrations if audio[1] is not None]
     print(f"Demonstration audio paths: {demonstration_paths}")
     print(f"Query audio path: {query_audio_path}")
 
     # 数据集处理
-    audio_dataset = dataset.DemoDataset(demonstration_paths, query_audio_path)
+    audio_dataset = dataset.DemoDataset(demonstration_paths, demonstration_labels, query_audio_path)
 
     # 调用检测函数
     result = detect(audio_dataset, model)
@@ -93,7 +104,7 @@ def gradio_ui():
             (demonstration_audio2, label2),
             (demonstration_audio3, label3),
         ]
-        return audio_deepfake_detection(demonstrations, query_audio)
+        return audio_deepfake_detection(demonstrations,query_audio)
 
     interface = gr.Interface(
         fn=detection_wrapper,