yaoayanui/src-tauri/resources/cover-python/cover_preview_generator.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
实时封面预览生成器（移植版：调用 cover_generator.py，不依赖 subtitle_cover_generator_simple）
"""

import os
import sys
import json
import hashlib
import tempfile
import subprocess
from typing import Dict, Any

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
if SCRIPT_DIR not in sys.path:
    sys.path.insert(0, SCRIPT_DIR)

COVER_GENERATOR = os.path.join(SCRIPT_DIR, "cover_generator.py")

STYLE_CONFIGS = {
    "default": {
        "titlePosition": "top",
        "titleColor": "#FFFFFF",
        "titleFontSize": 72,
        "titleStrokeWidth": 0,
    },
    "blur-bg": {
        "titlePosition": "top",
        "titleColor": "#FFFFFF",
        "titleFontSize": 72,
        "backgroundBlurEnabled": True,
    },
    "outline": {
        "titlePosition": "top",
        "titleColor": "#FFFFFF",
        "titleFontSize": 72,
        "titleStrokeWidth": 4,
        "titleStrokeColor": "#000000",
    },
    "gradient": {
        "titlePosition": "top",
        "titleColor": "#FFFFFF",
        "titleFontSize": 72,
    },
    "split": {
        "titlePosition": "center",
        "titleColor": "#FFFFFF",
        "titleFontSize": 64,
    },
}


class CoverPreviewGenerator:
    def __init__(self):
        self.temp_dir = os.path.join(tempfile.gettempdir(), "aiclient-cover-previews")
        os.makedirs(self.temp_dir, exist_ok=True)
        self.preview_cache = {}
        self.COVER_STYLES = list(STYLE_CONFIGS.keys())

    def get_video_hash(self, video_path: str) -> str:
        h = hashlib.md5()
        with open(video_path, "rb") as f:
            for chunk in iter(lambda: f.read(4096), b""):
                h.update(chunk)
        return h.hexdigest()[:16]

    def _run_cover_generator(self, video_path: str, title: str, output_path: str, style: str) -> bool:
        if not os.path.isfile(COVER_GENERATOR):
            print(f"Missing {COVER_GENERATOR}", file=sys.stderr)
            return False
        config = dict(STYLE_CONFIGS.get(style, STYLE_CONFIGS["default"]))
        python_exe = os.environ.get("AICLIENT_PYTHON_PATH") or sys.executable
        cmd = [
            python_exe,
            COVER_GENERATOR,
            "--video",
            video_path,
            "--title",
            title,
            "--output",
            output_path,
            "--config",
            json.dumps(config, ensure_ascii=False),
            "--timestamp",
            "1",
        ]
        env = os.environ.copy()
        env["PYTHONIOENCODING"] = "utf-8"
        try:
            r = subprocess.run(
                cmd,
                cwd=SCRIPT_DIR,
                capture_output=True,
                text=True,
                timeout=120,
                env=env,
            )
            if r.returncode != 0:
                print(r.stderr or r.stdout, file=sys.stderr)
            return r.returncode == 0 and os.path.isfile(output_path)
        except Exception as e:
            print(f"cover_generator failed: {e}", file=sys.stderr)
            return False

    def generate_preview(self, video_path: str, style: str, output_path: str) -> bool:
        return self._run_cover_generator(
            video_path,
            "封面预览",
            output_path,
            style,
        )

    def generate_all_previews(self, video_path: str) -> Dict[str, Any]:
        if not os.path.exists(video_path):
            return {"error": f"Video file not found: {video_path}"}

        video_hash = self.get_video_hash(video_path)
        if video_hash in self.preview_cache:
            return self.preview_cache[video_hash]

        previews = []
        errors = []
        for style in self.COVER_STYLES:
            output_path = os.path.join(self.temp_dir, f"preview_{video_hash}_{style}.jpg")
            ok = self.generate_preview(video_path, style, output_path)
            if ok:
                previews.append({"style": style, "path": output_path})
            else:
                errors.append(style)

        response = {
            "previews": previews,
            "videoHash": video_hash,
            "tempDir": self.temp_dir,
        }
        if errors:
            response["errors"] = errors
        self.preview_cache[video_hash] = response
        return response


def main():
    import argparse

    parser = argparse.ArgumentParser(description="Generate cover style previews")
    parser.add_argument("video_path", help="Path to the video file")
    parser.add_argument("--output-dir", help="Output directory for previews")
    args = parser.parse_args()

    if not os.path.exists(args.video_path):
        print(json.dumps({"error": f"not found: {args.video_path}"}, ensure_ascii=False))
        sys.exit(1)

    gen = CoverPreviewGenerator()
    if args.output_dir:
        gen.temp_dir = args.output_dir
        os.makedirs(gen.temp_dir, exist_ok=True)

    result = gen.generate_all_previews(args.video_path)
    print(json.dumps(result, ensure_ascii=False))


if __name__ == "__main__":
    main()