app.py

import os
import subprocess
import tempfile
import sample_code_c5_2
import numpy as np
import soundfile as sf
import pyroomacoustics as pra
from auxiva_online import auxiva_online
import matplotlib
matplotlib.use('Agg')  # Hugging Face Spaces対応（ディスプレイなし環境）
import matplotlib.pyplot as plt
from pdb import set_trace as test
from scipy import signal
import uuid
import shutil
import atexit

# ===== 実行環境判定: Hugging Face / Local =====
try:
    import gradio as gr
    USE_GRADIO = True
except ImportError:
    USE_GRADIO = False

# ===== デフォルト値（Hugging Faceもローカルも共通） =====
DEFAULT_WIN_LEN = 1024
DEFAULT_ALPHA = 0.9

# ===== ffmpegで音声抽出（サンプリング周波数は指定しない） =====
def extract_audio(input_video, output_audio):
    cmd = ["ffmpeg", "-i", input_video, "-vn", output_audio]
    subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)

# ===== サンプル混合データ作成 =====
def create_sample_audio(audio_path):
    sample_code_c5_2.main(audio_output=audio_path)

# ===== オフライン IVA音源分離 =====
def separate_sources(audio_path, win_len, alpha=0.95, use_online=False, session_dir=None):
    x, sr = sf.read(audio_path)
    hop = win_len // 2

    _, _, X = signal.stft(x.T, sr, nperseg=win_len, noverlap=hop)
    if use_online:
        Y = auxiva_online(X.transpose((2,1,0)), alpha=alpha)
    else:
        Y = pra.bss.auxiva(X.transpose((2,1,0)), n_iter=50)
    _, y = signal.istft(Y.transpose((2,1,0)), sr, nperseg=win_len, noverlap=hop)

    # セッションディレクトリにファイルを保存
    if session_dir is None:
        session_dir = "."
    
    out1 = os.path.join(session_dir, "separated_1.wav")
    out2 = os.path.join(session_dir, "separated_2.wav")
    sf.write(out1, y[0], sr)
    sf.write(out2, y[1], sr)

    return (out1, out2), y, sr

# ===== 波形プロット（混合音声と分離音声を一緒に表示） =====
def plot_waveform(mixed_data, separated_data):
    fig, ax = plt.subplots(2, 2, figsize=(12, 6))
    
    # 混合音声
    for i in range(2):
        ax[0, i].plot(mixed_data[:, i], alpha=0.7, label=f"Mixed Channel {i+1}")
        ax[0, i].set_title(f"Mixed Audio - Channel {i+1}")
        ax[0, i].set_ylabel("Amplitude")
        ax[0, i].legend()
    
    # 分離音声
    for i in range(2):
        ax[1, i].plot(separated_data[i], alpha=0.7, label=f"Separated {i+1}", color="orange")
        ax[1, i].set_title(f"Separated Audio {i+1}")
        ax[1, i].set_ylabel("Amplitude")
        ax[1, i].set_xlabel("Time")
        ax[1, i].legend()
    
    plt.tight_layout()
    return fig

# ===== サンプリング周波数を計算（動画優先、サンプルからフォールバック） =====
def get_sampling_rate(video_file, sample_file):
    """動画ファイルまたはサンプル音声からサンプリング周波数を取得（優先度: video > sample）"""
    # 優先順位: video_file > sample_file
    target = video_file or sample_file
    if not target:
        return "**サンプリング周波数**: 動画またはサンプルが必要です", DEFAULT_WIN_LEN

    try:
        # Gradio Audio/Video コンポーネントの形式を処理
        target_path = None
        
        # Audio コンポーネントの場合、タプル (sample_rate, audio_data) を返すことがある
        if isinstance(target, tuple) and len(target) == 2:
            # (sample_rate, audio_data) の形式
            sr = target[0]
            target_path = None  # パスは不要、既にサンプリング周波数がある
        elif isinstance(target, dict):
            # Video/Audio コンポーネントの辞書形式
            target_path = target.get("name") or target.get("path")
        else:
            # ファイルパス文字列
            target_path = target
        
        # パスから読み込む場合
        if target_path is not None:
            if not os.path.exists(target_path):
                return "**サンプリング周波数**: ファイルが見つかりません", DEFAULT_WIN_LEN
            
            # まず soundfile で読み込みを試みる
            try:
                data, sr = sf.read(target_path)
            except Exception:
                # ffmpegで音声抽出して読み直す
                tmp = "__tmp_extracted_audio.wav"
                try:
                    subprocess.run(["ffmpeg", "-y", "-i", target_path, "-vn", "-acodec", "pcm_s16le", tmp], check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
                    data, sr = sf.read(tmp)
                except Exception as e:
                    print(f"Error extracting audio for sampling rate: {e}")
                    return "**サンプリング周波数**: 読み取りに失敗しました", DEFAULT_WIN_LEN
                finally:
                    if os.path.exists(tmp):
                        try:
                            os.remove(tmp)
                        except Exception:
                            pass

        sr_text = f"**サンプリング周波数**: {sr} Hz"
        candidates = np.array([256, 512, 1024, 2048, 4096, 8192, 16384])
        target_win = int(sr * 0.2)
        closest = int(candidates[np.argmin(np.abs(candidates - target_win))])
        return sr_text, closest
    except Exception as e:
        return f"**サンプリング周波数**: エラー ({str(e)})", DEFAULT_WIN_LEN

# ===== メイン処理 (動画 → 分離) =====
def process(video_file, win_len=DEFAULT_WIN_LEN, alpha=DEFAULT_ALPHA, use_online=False, session_dir=None):
    if session_dir is None:
        session_dir = "."
    
    with tempfile.TemporaryDirectory() as tmpdir:
        audio_path = os.path.join(tmpdir, "audio.wav")
        
        # video_file が Gradio Audio コンポーネントからのタプル (sr, audio_data) の場合に対応
        if isinstance(video_file, tuple) and len(video_file) == 2:
            sr_audio, audio_data = video_file
            # タプル形式の場合は直接ファイルに書き込む
            sf.write(audio_path, audio_data, sr_audio)
        elif video_file is not None and os.path.exists(video_file):
            extract_audio(video_file, audio_path)
        else:
            create_sample_audio(audio_path)

        # 混合音声データを読み込む
        mixed_data, sr = sf.read(audio_path)
        
        (file1, file2), sep_data, sr = separate_sources(audio_path, win_len, alpha, use_online, session_dir=session_dir)
        fig = plot_waveform(mixed_data, sep_data)
        
        # 混合音声をファイルに保存（セッションディレクトリ）
        # - ステレオで保存するファイルと、チャンネルごとのファイルを作成
        mixed_out = os.path.join(session_dir, "mixed_audio.wav")
        mixed_out_ch1 = os.path.join(session_dir, "mixed_audio_ch1.wav")
        mixed_out_ch2 = os.path.join(session_dir, "mixed_audio_ch2.wav")

        try:
            if mixed_data.ndim == 1:
                # モノラル入力はそのまま保存
                sf.write(mixed_out, mixed_data, sr)
                sf.write(mixed_out_ch1, mixed_data, sr)
                sf.write(mixed_out_ch2, mixed_data, sr)
            else:
                # 複数チャンネルがある場合はステレオとして保存し、各チャンネルを個別に書き出す
                chans = mixed_data.shape[1]
                if chans >= 2:
                    sf.write(mixed_out, mixed_data[:, :2], sr)
                    sf.write(mixed_out_ch1, mixed_data[:, 0], sr)
                    sf.write(mixed_out_ch2, mixed_data[:, 1], sr)
                else:
                    # チャンネル数が1の扱い
                    sf.write(mixed_out, mixed_data[:, 0], sr)
                    sf.write(mixed_out_ch1, mixed_data[:, 0], sr)
                    sf.write(mixed_out_ch2, mixed_data[:, 0], sr)
        except Exception as e:
            # 失敗したら最小限の情報を返す（呼び出し側でエラーハンドリング）
            print(f"Warning: failed to write mixed audio files: {e}")
            mixed_out_ch1 = mixed_out_ch2 = None

        return sr, fig, file1, file2, mixed_out_ch1, mixed_out_ch2

# ===== Gradio UI (Hugging Face上のみ) =====
if USE_GRADIO:
    # セッション管理: ユーザーごとに隔離されたディレクトリ
    session_dir = None
    
    def get_session_dir():
        """ユーザーセッション用の一時ディレクトリを取得"""
        global session_dir
        if session_dir is None:
            # セッションディレクトリを作成（アプリ起動時に1回だけ）
            session_id = str(uuid.uuid4())[:8]
            session_dir = os.path.join(tempfile.gettempdir(), f"gradio_session_{session_id}")
            os.makedirs(session_dir, exist_ok=True)
            # パーミッション設定：所有者のみアクセス可能 (mode 0o700)
            os.chmod(session_dir, 0o700)
            print(f"🔹 Session directory created: {session_dir}")
            # プログラム終了時にセッションディレクトリを削除するコールバック登録
            atexit.register(cleanup_session_dir, session_dir)
        return session_dir
    
    def cleanup_session_dir(session_dir):
        """セッション終了時にディレクトリを削除"""
        try:
            if os.path.exists(session_dir):
                shutil.rmtree(session_dir)
                print(f"🔹 Session directory cleaned up: {session_dir}")
        except Exception as e:
            print(f"⚠️ Failed to clean up session directory: {e}")
    
    with gr.Blocks() as demo:
        gr.Markdown("🎧 **IVAステレオ音源分離デモ (pyroomacoustics)**")
        # 小さい動画プレビューにするCSS（アップロードした大きな動画が邪魔なときに有効）
        # ここで高さを制限し、必要ならブラウザで拡大して見るよう促すUIにできます。
        gr.HTML("<style>video{max-height:240px; max-width:100%;}</style>")
        gr.Markdown(
            """ 
            このデモでは、動画ファイルから音声を抽出し、2つの音源に分離します。動画をアップロードするか、サンプル混合音声を生成してから「分離開始」ボタンを押してください。  
            注意:  
             - 個人情報や機密情報はアップロードしないでください
             - アップロードによって生じるトラブルは責任を負いかねます 
            """)
        # 入力ファイル (左:アップローダ, 右:サンプル生成/再生)
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("**📤 入力ファイル**")
                video_input = gr.Video(label="🎥 動画ファイルをアップロード")
                gr.Markdown("**注意:** 動画をアップロードするか、または「🎼 サンプル混合音声を生成・再生」を押してください。どちらかを必ず行ってください。")

            with gr.Column(scale=1):
                gr.Markdown("**🎵 サンプル生成**")
                gen_sample_btn = gr.Button("🎼 サンプル混合音声を生成・再生", variant="secondary", size="lg")
                sample_audio = gr.Audio(label="生成されたサンプル", interactive=False)

        # パラメータ / 処理モード
        with gr.Group():
            gr.Markdown("**⚙️ 処理モードとパラメータ**")
            
            # 窓長説明（最初）
            gr.Markdown("💡 窓長（サンプル数）は表示されたサンプリング周波数の0.2倍程度を推奨します。下のボタンを押すと自動設定されます。")
            
            # サンプリング周波数計算ボタンと表示
            with gr.Row():
                check_sr_btn = gr.Button("▶ サンプリング周波数を計算", variant="secondary", size="sm")
                sr_display = gr.Markdown("**サンプリング周波数**: まだ計算されていません")
            
            # STFT 窓長
            win_len = gr.Dropdown(
                [256, 512, 1024, 2048, 4096, 8192, 16384],
                value=DEFAULT_WIN_LEN,
                label="STFT 窓長 (hopは自動で半分)"
            )
            
            # 処理モード選択
            use_online = gr.Radio(
                choices=["オフライン (AuxIVA)", "オンライン (Online AuxIVA)"],
                value="オフライン (AuxIVA)",
                label="処理モード",
                interactive=True
            )
            
            # オンライン処理専用パラメータ区切り
            gr.Markdown("💡 **以下はオンライン処理の時のみ有効です**")
            
            # Online-specific parameters
            forgetting_factor = gr.Slider(0.90, 0.99, value=DEFAULT_ALPHA, step=0.01, label="忘却係数 (高いほど過去の情報を忘れにくくなります)")

            run_btn = gr.Button("▶ 分離開始", size="lg", variant="primary")

        # ステータス表示（入力チェック等のメッセージを表示）
        status_output = gr.Markdown("", visible=True)
        # 結果表示 (下部)
        with gr.Group():
            gr.Markdown("**📊 結果**")
            fig_out = gr.Plot(label="混合音声 vs 分離音声の波形比較")
            with gr.Row():
                mixed_audio_ch1 = gr.Audio(label="🎵 混合音 - チャンネル1", type="filepath")
                mixed_audio_ch2 = gr.Audio(label="🎵 混合音 - チャンネル2", type="filepath")
                audio1 = gr.Audio(label="🔊 分離音1", type="filepath")
                audio2 = gr.Audio(label="🔊 分離音2", type="filepath")

        # コールバック: サンプル生成
        def generate_and_play_sample():
            try:
                out_path = "sample_mixed_ch1.wav"
                # sample_code_c5_2 (若干改変) を使ってサンプルを生成
                sample_code_c5_2.main(out_path)
                data, sr = sf.read(out_path)
                # 生成されたサンプルを保存（分離処理で2chが必要）
                sf.write(out_path, data, sr)
                return out_path
            except Exception as e:
                print(f"Error generating sample audio: {e}")
                return None

        gen_sample_btn.click(generate_and_play_sample, inputs=[], outputs=[sample_audio])

        # サンプリング周波数計算コールバック（win_lenを更新）
        check_sr_btn.click(get_sampling_rate, inputs=[video_input, sample_audio], outputs=[sr_display, win_len])

        # main processing callback
        def process_and_format(video_file, sample_file, win_len_val, forgetting_val, use_online_mode):
            try:
                use_online_flag = use_online_mode == "オンライン (Online AuxIVA)"

                # 入力チェック: 動画がアップロードされていない && サンプルも生成されていない -> エラー表示
                if not video_file and not sample_file:
                    status_msg = "❗ 動画が選択されていません。動画をアップロードするか、サンプルを生成してください。"
                    return None, None, None, None, status_msg

                print(f"DEBUG: video_file={video_file}, sample_file={sample_file}")
                print(f"DEBUG: video_file type={type(video_file)}, sample_file type={type(sample_file)}")
                
                target = video_file or sample_file
                print(f"DEBUG: target={target}, target type={type(target)}")

                # セッションディレクトリを取得してprocess()に渡す
                sess_dir = get_session_dir()
                sr, fig, file1, file2, mixed_out_ch1, mixed_out_ch2 = process(target, win_len=win_len_val, alpha=forgetting_val, use_online=use_online_flag, session_dir=sess_dir)
                status_msg = "✅ 分離が完了しました。下のオーディオを確認してください。"
                return fig, mixed_out_ch1, mixed_out_ch2, file1, file2, status_msg
            except Exception as e:
                import traceback
                error_msg = f"❌ エラーが発生しました: {str(e)}\n{traceback.format_exc()}"
                print(error_msg)
                return None, None, None, None, None, error_msg

        run_btn.click(
            process_and_format,
            inputs=[video_input, sample_audio, win_len, forgetting_factor, use_online],
            outputs=[fig_out, mixed_audio_ch1, mixed_audio_ch2, audio1, audio2, status_output]
        )

        gr.Markdown(
        """
        ---

        **Resources Used for Demo Site**
        - pyroomacoustics (https://github.com/LCAV/pyroomacoustics)
        - Pythonで学ぶ音源分離（機械学習実践シリーズ）のソースコード (https://github.com/masahitotogami/python_source_separation)
        - CMU ARCTIC Concatenated 15s (https://zenodo.org/records/3066489) 
        - J. Kominek and A. W. Black, “CMU ARCTIC databases for speech synthesis,” CMU-LTI-03-177, 2003.

        **references**
        - N. Ono, *Stable and fast update rules for independent vector analysis based on auxiliary function technique,* Proc. WASPAA, pp. 189-192, Oct. 2011.
        - T. Taniguchi, N. Ono, A. Kawamura, and S. Sagayama, *An auxiliary-function approach to online independent vector analysis for real-time blind source separation,* Proc. HSCMA, pp. 107-111, May. 2014.

        This demo is for educational and research purposes only.
        """
        )

    demo.launch(debug=True)
else:
    # ===== ローカルモード：対話形式でパラメータ入力 =====
    print("🔍 Local mode detected (gradio not available)")
    print("=" * 50)
    
    # test_video 入力
    print("\n【ビデオファイル】")
    video_input = input("ビデオファイルを入力してください (Enterで test_input.mp4, sでサンプル, 他は入力値): ").strip()
    if video_input.lower() in [""]:
        test_video = "test_input.mp4"
    elif video_input.lower() in ["s"]:
        test_video = None
        DEFAULT_WIN_LEN = 2048
    else:
        test_video = video_input
    
    if  test_video is not None and not os.path.exists(test_video):
        raise FileNotFoundError(f"{test_video} が見つかりません。")
    
    # win_len 入力
    print("\n【STFT 窓長】")
    win_len_input = input(f"win_lenを入力してください (デフォルト: {DEFAULT_WIN_LEN}): ").strip()
    if win_len_input == "":
        win_len = DEFAULT_WIN_LEN
    else:
        try:
            win_len = int(win_len_input)
        except ValueError:
            print(f"⚠️  無効な入力。デフォルト値 {DEFAULT_WIN_LEN} を使用します。")
            win_len = DEFAULT_WIN_LEN
    
    # use_online 入力
    print("\n【処理モード】")
    use_online_input = input("オンライン処理を使用しますか？ (yes/no, デフォルト: yes): ").strip().lower()
    use_online = not use_online_input in ["no", "n"]

    # use_online が True の場合のみ alpha 入力
    if use_online:
        print("\n【Alpha パラメータ】")
        alpha_input = input(f"alphaを入力してください (0.90-0.99, デフォルト: {DEFAULT_ALPHA}): ").strip()
        if alpha_input == "":
            alpha = DEFAULT_ALPHA
        else:
            try:
                alpha = float(alpha_input)
                if not (0.90 <= alpha <= 0.99):
                    print(f"⚠️  alpha は 0.90-0.99 の範囲で指定してください。デフォルト値 {DEFAULT_ALPHA} を使用します。")
                    alpha = DEFAULT_ALPHA
            except ValueError:
                print(f"⚠️  無効な入力。デフォルト値 {DEFAULT_ALPHA} を使用します。")
                alpha = DEFAULT_ALPHA
        
    else:
        alpha = DEFAULT_ALPHA
    
    # パラメータ表示
    print("\n" + "=" * 50)
    print("【処理パラメータ】")
    print(f"  ビデオファイル: {test_video if test_video else 'サンプル音声を生成'}")
    print(f"  win_len: {win_len}")
    print(f"  use_online: {use_online}")
    if use_online:
        print(f"  alpha: {alpha}")
    print("=" * 50)
    print("\n処理を開始します...")
    
    # 処理実行
    sr, fig, file1, file2, mixed_ch1, mixed_ch2 = process(test_video, win_len=win_len, alpha=alpha, use_online=use_online)
    
    print(f"✅ 処理完了")
    print(f"Sampling Rate: {sr}")
    print(f"Saved separated files: {file1}, {file2}")
    print(f"Saved mixed channel files: {mixed_ch1}, {mixed_ch2}")
    fig.savefig("separated_plot.png")
    print("Waveform plot saved to separated_plot.png")