import gradio as gr
from transformers import pipeline
import librosa
import numpy as np

# --- 設定模型 ---
# 這裡目前使用一個通用的音訊分類模型 (Audio Spectrogram Transformer) 作為 Demo
# 它可以辨識 "Vibration", "Rumble", "Explosion" 等類別
# 未來您可以將此處替換為您訓練好的地震波專用模型 (例如: "your-username/seismic-wave-classifier")
MODEL_ID = "MIT/ast-finetuned-audioset-10-10-0.4593"

print(f"正在載入模型: {MODEL_ID} ...")
# 建立 Audio Pipeline，這行就是書中第 5 章的概念
pipe = pipeline("audio-classification", model=MODEL_ID)

def classify_seismic_audio(audio_filepath):
    """
    接收音訊檔案路徑，回傳前 5 名可能的分類結果
    """
    if audio_filepath is None:
        return "請上傳或錄製一段音訊"
    
    # 使用 Pipeline 進行推論
    # Pipeline 會自動處理讀取音檔、重取樣 (Resample) 等繁瑣工作
    preds = pipe(audio_filepath)
    
    # 整理輸出格式給 Gradio (將 list 轉為 dictionary)
    # 輸出範例: {'P-wave': 0.95, 'Background Noise': 0.05}
    output = {}
    for p in preds:
        output[p["label"]] = p["score"]
        
    return output

# --- 建立 Gradio 介面 ---
description_text = """
### 🌍 地震波音訊分類器 (Seismic Event Classifier)
此範例展示如何使用 **Hugging Face Audio Pipeline** 來判讀訊號。
您可以上傳一段地震波轉成的音訊檔（.wav/.mp3），模型將判斷其特徵。

* **目前模型**：通用環境音模型 (MIT/AST) - 用於展示架構
* **應用場景**：分辨 P波 (P-wave)、S波 (S-wave) 或 背景雜訊 (Noise)
"""

demo = gr.Interface(
    fn=classify_seismic_audio,
    inputs=gr.Audio(type="filepath", label="上傳地震波音訊"),
    outputs=gr.Label(num_top_classes=5, label="分類結果與信心分數"),
    title="🔊 AI 地震訊號辨識",
    description=description_text,
    examples=[
        ["sample_audio.wav"] # 如果您有範例檔案可以放這裡
    ]
)

if __name__ == "__main__":
    demo.launch()