import gradio as gr
import torch
import librosa
import numpy as np

def audio_to_pth(audio):
    # 音声ファイル（ファイルパス）を読み込む
    y, sr = librosa.load(audio, sr=None)

    # メルスペクトログラムに変換
    mel_spectrogram = librosa.feature.melspectrogram(y=y, sr=sr, n_mels=128)

    # メルスペクトログラムを対数スケールに変換
    mel_spectrogram_db = librosa.power_to_db(mel_spectrogram, ref=np.max)

    # メルスペクトログラムをテンソルに変換
    tensor = torch.tensor(mel_spectrogram_db)

    # テンソルを5次元に変換
    tensor = tensor.unsqueeze(0).unsqueeze(0).unsqueeze(0)  # 5次元に拡張

    # テンソルを .pth ファイルに保存
    output_path = "audio_features_5d.pth"
    torch.save(tensor, output_path)

    return output_path

# Gradio インターフェースの設定
iface = gr.Interface(
    fn=audio_to_pth,
    inputs=gr.Audio(type="filepath"),
    outputs="file",
    title="Audio to 5D Tensor .PTH Converter",
    description="Upload an audio file to convert it to a .pth file containing a 5D tensor with audio features in mel spectrogram format."
)

iface.launch()