Update app.py

app.py

@@ -1,32 +1,37 @@
 import gradio as gr
 import os
-import dataset
 import torch
-from model import Wav2Vec2BERT_Llama
+from model import Wav2Vec2BERT_Llama  # 自定义模型模块
+import dataset  # 自定义数据集模块
 
-# init
+# 初始化设备
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# # init model
-# model = Wav2Vec2BERT_Llama().to(device)
-# checkpoint_path = "ckpt/model_checkpoint.pth"
-# if os.path.exists(checkpoint_path):
-#     checkpoint = torch.load(checkpoint_path)
-#     model_state_dict = checkpoint['model_state_dict']
-    
-#     # 处理模型状态字典
-#     if hasattr(model, 'module') and not any(key.startswith('module.') for key in model_state_dict.keys()):
-#         model_state_dict = {'module.' + key: value for key, value in model_state_dict.items()}
-#     elif not hasattr(model, 'module') and any(key.startswith('module.') for key in model_state_dict.keys()):
-#         model_state_dict = {key.replace('module.', ''): value for key, value in model_state_dict.items()}
-        
-#     model.load_state_dict(model_state_dict)
-#     model.eval()
-# else:
-#     raise FileNotFoundError(f"Not found checkpoint: {checkpoint_path}")
+# 初始化模型
+# def load_model():
+#     model = Wav2Vec2BERT_Llama().to(device)
+#     checkpoint_path = "ckpt/model_checkpoint.pth"
+#     if os.path.exists(checkpoint_path):
+#         checkpoint = torch.load(checkpoint_path)
+#         model_state_dict = checkpoint['model_state_dict']
 
+#         # 处理模型状态字典的 key
+#         if hasattr(model, 'module') and not any(key.startswith('module.') for key in model_state_dict.keys()):
+#             model_state_dict = {'module.' + key: value for key, value in model_state_dict.items()}
+#         elif not hasattr(model, 'module') and any(key.startswith('module.') for key in model_state_dict.keys()):
+#             model_state_dict = {key.replace('module.', ''): value for key, value in model_state_dict.items()}
 
+#         model.load_state_dict(model_state_dict)
+#         model.eval()
+#     else:
+#         raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
+#     return model
+
+# model = load_model()
+
+# 检测函数
 def detect(dataset, model):
+    """进行音频伪造检测"""
     with torch.no_grad():
         for batch in dataset:
             main_features = {
@@ -37,81 +42,65 @@ def detect(dataset, model):
                 'input_features': pf['input_features'].to(device),
                 'attention_mask': pf['attention_mask'].to(device)
             } for pf in batch['prompt_features']]
-        
 
+            # 模型的前向传播逻辑 (需要补充具体实现)
+            # 假设 result 是模型返回的结果
+            result = {"is_fake": True, "confidence": 85.5}  # 示例返回值
+    return result
 
+# 音频伪造检测主函数
+def audio_deepfake_detection(demonstrations, query_audio_path):
+    """
+    音频伪造检测函数
+    :param demonstrations: 演示音频路径和标签的列表
+    :param query_audio_path: 查询音频路径
+    :return: 检测结果
+    """
+    demonstration_paths = [audio[0] for audio in demonstrations if audio[0] is not None]
+    print(f"Demonstration audio paths: {demonstration_paths}")
+    print(f"Query audio path: {query_audio_path}")
 
-def audio_deepfake_detection(demonstrations_container, audio_path):
-    """Audio deepfake detection function"""
-    # 收集所有demonstration的路径和标签
-    demonstration_paths = []
-    for child in demonstrations_container.children:
-        if isinstance(child, gr.Row):
-            audio = child.children[0].children[0].value
-            if audio is not None:
-                demonstration_paths.append(audio)
+    # 数据集处理
+    audio_dataset = dataset.DemoDataset(demonstration_paths, query_audio_path)
 
-    print("Demonstration audio paths: {}".format(demonstration_paths))
-    print("Query audio path: {}".format(audio_path))
+    # 调用检测函数
+    result = detect(audio_dataset, model)
 
-    # dataset
-    dataset = dataset.DemoDataset(demonstration_paths, audio_path)
-    # Example return value, modify according to your model
-    result = detect(dataset, model)
-    
-    # Return detection results and confidence scores
+    # 返回结果
     return {
         "Is AI Generated": result["is_fake"],
         "Confidence": f"{result['confidence']:.2f}%"
     }
 
-with gr.Blocks() as demo:
-    gr.Markdown(
-        """
-        # Audio Deepfake Detection System
-        
-        This demo helps you detect whether an audio clip is AI-generated or authentic.
-        """
-    )
-    
-    with gr.Column() as demonstrations_container:
-        gr.Markdown("## Demonstration Audios (Optional)")
-        
-        # 创建3个固定的demonstration组件
-        for i in range(3):
-            with gr.Row():
-                with gr.Column(scale=8):
-                    audio = gr.Audio(
-                        sources=["upload"],
-                        type="filepath",
-                        label=f"Demonstration Audio {i+1}"
-                    )
-                with gr.Column(scale=3):
-                    label = gr.Dropdown(
-                        choices=["bonafide", "spoof"],
-                        value="bonafide",
-                        label="Label"
-                    )
-
-    # Query audio input component
-    query_audio_input = gr.Audio(
-        sources=["upload"],
-        label="Query Audio (Audio for Detection)",
-        type="filepath",
-    )
-
-    # Submit button
-    submit_btn = gr.Button(value="Start Detection", variant="primary")
-
-    # Output results
-    output_labels = gr.Json(label="Detection Results")
+# Gradio 界面
+def gradio_ui():
+    def detection_wrapper(demonstration_audio1, label1, demonstration_audio2, label2, demonstration_audio3, label3, query_audio):
+        # 将输入音频和标签封装成列表
+        demonstrations = [
+            (demonstration_audio1, label1),
+            (demonstration_audio2, label2),
+            (demonstration_audio3, label3),
+        ]
+        return audio_deepfake_detection(demonstrations, query_audio)
 
-    # Set click event
-    submit_btn.click(
-        fn=audio_deepfake_detection,
-        inputs=[demonstrations_container, query_audio_input],
-        outputs=[output_labels]
+    # 构建 Gradio 界面
+    interface = gr.Interface(
+        fn=detection_wrapper,  # 主函数
+        inputs=[
+            gr.Audio(source="upload", type="filepath", label="Demonstration Audio 1"),
+            gr.Dropdown(choices=["bonafide", "spoof"], value="bonafide", label="Label 1"),
+            gr.Audio(source="upload", type="filepath", label="Demonstration Audio 2"),
+            gr.Dropdown(choices=["bonafide", "spoof"], value="bonafide", label="Label 2"),
+            gr.Audio(source="upload", type="filepath", label="Demonstration Audio 3"),
+            gr.Dropdown(choices=["bonafide", "spoof"], value="bonafide", label="Label 3"),
+            gr.Audio(source="upload", type="filepath", label="Query Audio (Audio for Detection)")
+        ],
+        outputs=gr.JSON(label="Detection Results"),
+        title="Audio Deepfake Detection System",
+        description="Upload demonstration audios and a query audio to detect whether the query is AI-generated.",
     )
+    return interface
 
 if __name__ == "__main__":
-    demo.launch(share=False)
+    demo = gradio_ui()
+    demo.launch()