xiaoxia-saas/apps/worker/video_processing/thumbnail_generator.py

"""视频缩略图生成工具 — 抽取首帧上传到 OSS。"""

from __future__ import annotations

import logging
import tempfile
from pathlib import Path

logger = logging.getLogger(__name__)


def extract_first_frame(
    video_path: str,
    output_path: str | None = None,
    *,
    width: int = 640,
    height: int = -1,
    timeout: int = 30,
) -> str:
    """抽取视频第一帧作为封面图。

    Args:
        video_path: 视频文件路径
        output_path: 输出图片路径，不传则用临时文件
        width: 输出宽度（默认 640，-1 表示按比例缩放）
        height: 输出高度（默认 -1，按比例缩放）
        timeout: 超时时间（秒）

    Returns:
        生成的缩略图文件路径

    Raises:
        subprocess.CalledProcessError: ffmpeg 执行失败
    """
    from video_processing.ffmpeg_utils import FFMPEG_BIN, run_ffmpeg

    if output_path is None:
        tmp = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
        tmp.close()
        output_path = tmp.name

    # -ss 00:00:01 取第1秒帧（避免首帧黑屏）
    # -vframes 1 只取一帧
    # -q:v 2 jpeg 高质量
    scale_filter = f"scale={width}:{height}:force_original_aspect_ratio=decrease"
    cmd = [
        FFMPEG_BIN,
        "-y",
        "-i",
        video_path,
        "-ss",
        "00:00:01",
        "-vframes",
        "1",
        "-vf",
        scale_filter,
        "-q:v",
        "2",
        output_path,
    ]

    try:
        run_ffmpeg(cmd, capture_output=True, timeout=timeout)
    except Exception:
        # 短视频可能没有第1秒，退回到第0帧
        cmd2 = [
            FFMPEG_BIN,
            "-y",
            "-i",
            video_path,
            "-ss",
            "00:00:00",
            "-vframes",
            "1",
            "-vf",
            scale_filter,
            "-q:v",
            "2",
            output_path,
        ]
        run_ffmpeg(cmd2, capture_output=True, timeout=timeout)

    if not Path(output_path).exists() or Path(output_path).stat().st_size == 0:
        raise RuntimeError(f"Thumbnail generation failed: {output_path}")

    return output_path


def generate_and_upload_thumbnail(
    video_path: str,
    storage_key: str,
) -> str | None:
    """生成缩略图并上传到 OSS，返回 URL。

    Args:
        video_path: 本地视频路径
        storage_key: OSS 存储 key（如 generated/projects/xxx/thumbnails/yyy.jpg）

    Returns:
        上传成功返回 URL，失败返回 None
    """
    thumbnail_path = None
    try:
        thumbnail_path = extract_first_frame(video_path)
    except Exception as e:
        logger.warning("Failed to extract thumbnail from %s: %s", video_path, e)
        return None

    try:
        from video_processing.oss_helpers import upload_to_oss

        url = upload_to_oss(thumbnail_path, storage_key)
        return url
    except Exception as e:
        logger.warning("Failed to upload thumbnail to OSS: %s", e)
        return None
    finally:
        # 清理临时文件
        if thumbnail_path:
            try:
                Path(thumbnail_path).unlink(missing_ok=True)
            except Exception:
                pass