app.py

import subprocess
import sys
import gradio as gr
from PIL import Image, PngImagePlugin
import tempfile
import os
import zipfile
from datetime import datetime
import shutil
import json # 修正点: jsonライブラリをインポート

"""
Video Frame Extractor (PNG)
--------------------------
* ffmpeg で **フルレンジ RGB24 PNG** を抽出。
* gAMA / cHRM チャンクを削除して sRGB 相当へ正規化（色ずれ防止）。
* 個別ダウンロード (gr.Files) と ZIP (ZIP_STORED) の両方を提供。
* 末尾フレームは `-sseof -0.05` 方式で確実に取得。
* 🆕 2025‑08‑08: ファイル名にフレーム番号 (7 桁ゼロ埋め) を付与。
    * 例) `sample_0000000_start.png`, `sample_0001234_end.png`
* 🛠️ 2025-08-08: ffprobe の情報取得を JSON 形式に変更し、安定性を向上。
"""

###############################
# 1. ユーティリティ
###############################

def install_ffmpeg():
    try:
        subprocess.run(["ffmpeg", "-version"], capture_output=True, check=True)
        return True, "ffmpeg は既にインストール済みです"
    except (subprocess.CalledProcessError, FileNotFoundError):
        if not sys.platform.startswith("linux"):
            return False, f"自動インストール未対応 OS: {sys.platform}"
        try:
            subprocess.run(["apt", "update"], check=True)
            subprocess.run(["apt", "install", "-y", "ffmpeg"], check=True)
            return True, "ffmpeg をインストールしました"
        except subprocess.CalledProcessError as e:
            return False, f"ffmpeg インストール失敗: {e}"

# 修正点: この関数をまるごと入れ替え
def get_video_info(video_path: str):
    """
    ffprobe を使って動画のフレーム数と再生時間を取得します。
    より安定した JSON 出力形式を利用します。
    """
    cmd = [
        "ffprobe",
        "-v", "quiet",
        "-select_streams", "v:0",
        "-show_entries", "stream=nb_frames,duration,r_frame_rate",
        "-of", "json",
        video_path,
    ]
    try:
        # ffprobeの実行と結果の取得
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        data = json.loads(result.stdout)
        stream_data = data.get("streams", [{}])[0]

        # フレーム数の取得 (nb_framesが優先)
        frame_count_str = stream_data.get("nb_frames")
        if frame_count_str and frame_count_str.isdigit():
            frame_count = int(frame_count_str)
        else:
            # nb_framesがなければ、デュレーションとフレームレートから計算
            duration_str = stream_data.get("duration")
            frame_rate_str = stream_data.get("r_frame_rate", "0/1")

            if duration_str and "/" in frame_rate_str:
                duration_val = float(duration_str)
                num, den = map(int, frame_rate_str.split('/'))
                if den > 0:
                    frame_rate = num / den
                    frame_count = int(duration_val * frame_rate)
                else:
                    frame_count = None # 計算不能
            else:
                frame_count = None # 取得不能

        # 再生時間の取得
        duration = float(stream_data.get("duration", 0))

        return frame_count, duration

    except (subprocess.CalledProcessError, json.JSONDecodeError, IndexError, KeyError):
        # エラーが発生した場合はNoneを返す
        return (None, None)

###############################
# 2. PNG 正規化
###############################

def normalize_png(path: str):
    img = Image.open(path).convert("RGB")
    pnginfo = PngImagePlugin.PngInfo()
    img.save(path, "PNG", pnginfo=pnginfo, optimize=True)

###############################
# 3. フレーム抽出
###############################

def extract_frame_with_ffmpeg(video_path: str, output_path: str, *, is_last: bool = False):
    try:
        if is_last:
            cmd = [
                "ffmpeg", "-v", "error", "-sseof", "-0.05", "-i", video_path,
                "-vframes", "1", "-vcodec", "png", "-pix_fmt", "rgb24", "-color_range", "pc",
                "-y", output_path,
            ]
        else:
            cmd = [
                "ffmpeg", "-v", "error", "-i", video_path,
                "-vframes", "1", "-vcodec", "png", "-pix_fmt", "rgb24", "-color_range", "pc",
                "-y", output_path,
            ]
        subprocess.run(cmd, check=True)
        normalize_png(output_path)
        return True, "成功"
    except subprocess.CalledProcessError as e:
        return False, f"抽出失敗: {e}"

###############################
# 4. メイン処理
###############################

def pad(num: int, width: int = 7) -> str:
    return str(num).zfill(width)

def extract_frames_from_multiple_videos(video_files):
    if not video_files:
        return None, None, None, None, "⚠️ ビデオを選択してください"

    ok, msg_ffm = install_ffmpeg()
    if not ok:
        return None, None, None, None, f"❌ {msg_ffm}"

    results = [f"ℹ️ {msg_ffm}"]
    temp_dir = tempfile.mkdtemp()
    saved_files, first_imgs, last_imgs = [], [], []

    for fobj in video_files:
        base = os.path.splitext(os.path.basename(fobj.name))[0]
        frame_count, duration = get_video_info(fobj.name)

        # ファイル名生成
        first_png = os.path.join(temp_dir, f"{base}_{pad(0)}_start.png")
        if frame_count is not None and frame_count > 0:
            last_png = os.path.join(temp_dir, f"{base}_{pad(frame_count - 1)}_end.png")
        else:
            last_png = os.path.join(temp_dir, f"{base}_unknown_end.png")

        # 抽出
        ok_first, _ = extract_frame_with_ffmpeg(fobj.name, first_png, is_last=False)
        if not ok_first:
            results.append(f"❌ {base}: start 抽出失敗")
            continue
        first_imgs.append(Image.open(first_png)); saved_files.append(first_png)

        ok_last, _ = extract_frame_with_ffmpeg(fobj.name, last_png, is_last=True)
        if ok_last:
            last_imgs.append(Image.open(last_png)); saved_files.append(last_png)
        else:
            shutil.copy2(first_png, last_png)
            last_imgs.append(Image.open(last_png)); saved_files.append(last_png)
            results.append(f"⚠️ {base}: end 抽出失敗 → start 流用")

        results.append(
            f"✅ {base}: frames {frame_count or '?'} / {duration or '?'}s")

    # ZIP 生成
    zip_path = os.path.join(temp_dir, f"frames_{datetime.now():%Y%m%d_%H%M%S}.zip")
    with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_STORED) as z:
        for p in saved_files:
            z.write(p, os.path.basename(p))

    status = "\n".join(results)
    return (
        first_imgs[0] if first_imgs else None,
        last_imgs[0]  if last_imgs  else None,
        saved_files,
        zip_path,
        status,
    )

###############################
# 5. ギャラリー
###############################

def create_gallery(video_files, is_last=False):
    imgs, tmp = [], tempfile.mkdtemp()
    for f in video_files:
        base = os.path.splitext(os.path.basename(f.name))[0]
        # シンプルなギャラリー名 (番号不要)
        p = os.path.join(tmp, f"{base}_{'end' if is_last else 'start'}_gal.png")
        if extract_frame_with_ffmpeg(f.name, p, is_last=is_last)[0]:
            imgs.append(Image.open(p))
    return imgs

###############################
# 6. Gradio UI
###############################

# オシャレ配色のテーマ作成
def create_custom_theme():
    return gr.Theme(
        primary_hue="slate",
        secondary_hue="stone",
        neutral_hue="zinc",
        text_size="md",
        spacing_size="lg",
        radius_size="sm",
        font=[
            "Hiragino Sans",
            "Noto Sans JP",
            "Yu Gothic",
            "system-ui",
            "sans-serif"
        ],
        font_mono=[
            "SF Mono",
            "Monaco",
            "monospace"
        ]
    ).set(
        # 背景・文字色
        body_background_fill="#ffffff",
        body_text_color="#2A4359",  # ダークブルー
        # プライマリボタン
        button_primary_background_fill="#F2163E",  # 鮮やかなレッド
        button_primary_background_fill_hover="#C177F2",  # ホバー時パープル
        button_primary_text_color="#ffffff",
        # セカンダリボタン
        button_secondary_background_fill="#D9BFD4",  # 落ち着いたピンクグレー
        button_secondary_text_color="#2A4359",
        # 入力フォーム
        input_background_fill="#ffffff",
        input_border_color="#77D9D9",  # アクアブルーの枠線
        input_border_color_focus="#2A4359",  # フォーカス時濃紺
        # ブロック・パネル
        block_background_fill="#ffffff",
        block_border_color="#D9BFD4",
        panel_background_fill="#ffffff",
        panel_border_color="#D9BFD4",
        # スライダー
        slider_color="#F2163E"
    )

# Gradio UIにテーマ適用
custom_theme = create_custom_theme()

with gr.Blocks(theme=custom_theme, title="Video Frame Extractor (PNG)") as demo:

    # ヘッダー
    gr.HTML(f"""
    <div style='
        text-align: center;
        margin-bottom: 2rem;
        padding: 2rem;
        background: linear-gradient(135deg, #F2163E 0%, #77D9D9 100%);
        color: white;
        border-radius: 12px;
    '>
        <h1 style='
            font-size: 2.5rem;
            margin-bottom: 0.5rem;
            text-shadow: 2px 2px 4px rgba(0,0,0,0.3);
        '>
            🎞️ Video Frame Extractor
        </h1>
        <p style='
            font-size: 1.1rem;
            margin: 0;
            color: #ffffffcc;
        '>
            抽出した最初と最後のフレームを美しく保存
        </p>
    </div>
    """)

    with gr.Row():
        video_input = gr.File(
            label="動画をアップロード (複数可)",
            file_types=["video"],
            file_count="multiple"
        )

    extract_btn = gr.Button("フレームを抽出", variant="primary")
    status_box = gr.Textbox(label="ステータス", interactive=False, lines=8)

    with gr.Row():
        with gr.Column():
            first_preview = gr.Image(label="最初のフレーム", type="pil")
            last_preview  = gr.Image(label="最後のフレーム", type="pil")
        with gr.Column():
            download_files = gr.Files(label="個別フレーム (PNG)")
            download_zip   = gr.File(label="まとめて ZIP")

    with gr.Row():
        gallery_first = gr.Gallery(label="全ての最初", columns=3, rows=2, height="auto")
        gallery_last  = gr.Gallery(label="全ての最後", columns=3, rows=2, height="auto")

    def process(videos):
        first, last, files, zipf, status = extract_frames_from_multiple_videos(videos)
        g_first = create_gallery(videos, is_last=False)
        g_last  = create_gallery(videos, is_last=True)
        return first, last, files, zipf, status, g_first, g_last

    extract_btn.click(
        fn=process,
        inputs=video_input,
        outputs=[
            first_preview,
            last_preview,
            download_files,
            download_zip,
            status_box,
            gallery_first,
            gallery_last,
        ],
    )

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