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,
    seek_ratio: float = 0.15,
    min_seek_seconds: float = 1.0,
) -> str:
    """抽取视频封面图（默认取视频时长 15% 处的帧，避开片头纯色画面）。

    Args:
        video_path: 视频文件路径
        output_path: 输出图片路径，不传则用临时文件
        width: 输出宽度（默认 640，-1 表示按比例缩放）
        height: 输出高度（默认 -1，按比例缩放）
        timeout: 超时时间（秒）
        seek_ratio: 抽帧位置占视频时长的比例（默认 0.15，即 15% 处）
        min_seek_seconds: 最小抽帧时间（秒），避免极短视频 seek 到 0

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

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

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

    try:
        # 计算抽帧时间点：取视频时长 * seek_ratio，最少 min_seek_seconds 秒
        try:
            duration = probe_duration(video_path)
            seek_time = max(min_seek_seconds, duration * seek_ratio)
        except Exception:
            # probe 失败时 fallback 到第1秒
            seek_time = min_seek_seconds

        # 格式化为 HH:MM:SS.xx
        seek_str = _format_seek_time(seek_time)

        # -ss 放在 -i 前面（input seeking，更快但精度稍低，缩略图够用）
        # -vframes 1 只取一帧
        # -q:v 2 jpeg 高质量
        scale_filter = f"scale={width}:{height}:force_original_aspect_ratio=decrease,format=yuvj420p"
        cmd = [
            FFMPEG_BIN,
            "-y",
            "-ss",
            seek_str,
            "-i",
            video_path,
            "-vframes",
            "1",
            "-vf",
            scale_filter,
            "-q:v",
            "2",
            output_path,
        ]

        try:
            run_ffmpeg(cmd, capture_output=True, timeout=timeout)
        except Exception:
            # 失败时退回到第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
    except Exception:
        # 失败时清理自己创建的临时文件
        if _is_temp_output and output_path:
            try:
                Path(output_path).unlink(missing_ok=True)
            except Exception:
                pass
        raise


def _format_seek_time(seconds: float) -> str:
    """将秒数格式化为 HH:MM:SS.xx 格式。"""
    h = int(seconds // 3600)
    m = int((seconds % 3600) // 60)
    s = seconds % 60
    return f"{h:02d}:{m:02d}:{s:05.2f}"


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


def extract_cover_candidates(
    video_path: str,
    num_frames: int = 3,
    *,
    width: int = 640,
    timeout: int = 30,
) -> list[dict]:
    """在视频时长 25%/50%/75% 处各抽一帧，返回候选帧信息列表。

    Args:
        video_path: 视频文件路径
        num_frames: 抽帧数量（默认 3）
        width: 输出宽度
        timeout: 单帧超时（秒）

    Returns:
        [{"local_path": "...", "frame_time": 5.0}, ...]
    """
    from video_processing.ffmpeg_utils import FFMPEG_BIN, probe_duration, run_ffmpeg

    try:
        duration = probe_duration(video_path)
    except Exception:
        duration = 0.0

    if duration <= 0:
        duration = 5.0  # fallback

    # 计算抽帧时间点：25%, 50%, 75%
    ratios = []
    for i in range(1, num_frames + 1):
        ratios.append(i / (num_frames + 1))

    results = []
    for idx, ratio in enumerate(ratios):
        frame_time = max(0.5, duration * ratio)
        tmp = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
        tmp.close()
        output_path = tmp.name

        try:
            seek_str = _format_seek_time(frame_time)
            scale_filter = f"scale={width}:-1:force_original_aspect_ratio=decrease,format=yuvj420p"
            cmd = [
                FFMPEG_BIN, "-y",
                "-ss", seek_str,
                "-i", video_path,
                "-vframes", "1",
                "-vf", scale_filter,
                "-q:v", "2",
                output_path,
            ]
            run_ffmpeg(cmd, capture_output=True, timeout=timeout)

            if Path(output_path).exists() and Path(output_path).stat().st_size > 0:
                results.append({
                    "local_path": output_path,
                    "frame_time": round(frame_time, 2),
                })
            else:
                Path(output_path).unlink(missing_ok=True)
        except Exception as e:
            logger.warning("封面候选帧抽取失败 ratio=%.2f: %s", ratio, e)
            Path(output_path).unlink(missing_ok=True)

    return results


def extract_and_upload_cover_frames(
    video_path: str,
    plan_id: str,
    num_frames: int = 3,
) -> list[dict]:
    """抽取封面候选帧并上传到 OSS。

    Args:
        video_path: 本地视频路径
        plan_id: 剪辑计划 ID（用于 OSS 路径）
        num_frames: 抽帧数量

    Returns:
        [{"image_url": "https://...", "frame_time": 5.0, "storage_key": "covers/xxx/frame_0.jpg"}, ...]
    """
    candidates = extract_cover_candidates(video_path, num_frames=num_frames)
    if not candidates:
        logger.warning("封面候选帧抽取为空: plan_id=%s", plan_id)
        return []

    results = []
    for idx, cand in enumerate(candidates):
        local_path = cand["local_path"]
        frame_time = cand["frame_time"]
        storage_key = f"covers/{plan_id}/frame_{idx}.jpg"

        try:
            from video_processing.oss_helpers import upload_to_oss

            url = upload_to_oss(local_path, storage_key)
            if url:
                results.append({
                    "image_url": url,
                    "frame_time": frame_time,
                    "storage_key": storage_key,
                })
                logger.info(
                    "封面候选帧上传成功: plan_id=%s idx=%d frame_time=%.2f",
                    plan_id, idx, frame_time,
                )
        except Exception as e:
            logger.warning("封面候选帧上传失败: plan_id=%s idx=%d error=%s", plan_id, idx, e)
        finally:
            try:
                Path(local_path).unlink(missing_ok=True)
            except Exception:
                pass

    return results