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

"""视频封面抽帧工具 — 从视频中抽取帧作为封面，支持标题文字叠加。

统一封面管道（P1 优化后默认本地路径）：
- 默认路径：本地 ffmpeg -ss 单帧 seek 抽取 + cv2 质量评分（清晰度/亮度/色彩），1-2s 完成
- 可选 MediaKit 路径：配置 MEDIAKIT_COVER_ENABLED=true 时启用火山 MediaKit SceneChange 抽帧
- 从已渲染视频抽帧：标题已通过 ASS 字幕烧进视频，帧天然带标题，无需再叠加。
- 从源素材抽帧（API E2 兜底）：源素材无标题，通过 Pillow 在帧上绘制标题文字。
"""

from __future__ import annotations

import logging
import tempfile
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)


def apply_title_overlay(
    image_path: str,
    title_text: str,
    *,
    color: str = "#ffffff",
    position: str = "bottom",
    font_size: int | None = None,
    margin_ratio: float = 0.06,
    stroke_width_ratio: float = 0.04,
) -> str:
    """在图片上绘制标题文字（指定颜色 + 黑色描边/阴影）。"""
    from packages.shared.title_overlay import apply_title_to_image

    if not title_text or not title_text.strip():
        return image_path
    result = apply_title_to_image(
        image_path,
        title_text,
        color=color,
        position=position,
        font_size=font_size,
        margin_ratio=margin_ratio,
        stroke_width_ratio=stroke_width_ratio,
    )
    return result or image_path


def extract_first_frame(
    video_path: str,
    output_path: str | None = None,
    *,
    width: int = -1,
    height: int = -1,
    timeout: int = 30,
    seek_ratio: float = 0.15,
    seek_seconds: float | None = None,
    min_seek_seconds: float = 1.0,
) -> str:
    """抽取视频封面帧（ffmpeg -ss 单帧 seek，<100ms/帧）。

    Args:
        video_path: 视频文件路径
        output_path: 输出图片路径，不传则用临时文件
        width/height: 输出宽高（默认保持原始分辨率）
        timeout: 超时（秒）
        seek_ratio: 抽帧位置占视频时长的比例
        seek_seconds: 指定具体抽帧时间点（秒），优先于 seek_ratio
        min_seek_seconds: 最小抽帧时间
    """
    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:
        if seek_seconds is not None:
            seek_time = max(0.0, float(seek_seconds))
        else:
            try:
                duration = probe_duration(video_path)
                seek_time = max(min_seek_seconds, duration * seek_ratio)
            except Exception:
                seek_time = min_seek_seconds

        seek_str = _format_seek_time(seek_time)

        if width > 0 or height > 0:
            w_str = str(width) if width > 0 else "-1"
            h_str = str(height) if height > 0 else "-1"
            scale_filter = f"scale={w_str}:{h_str}:force_original_aspect_ratio=decrease,format=yuvj420p"
        else:
            scale_filter = "format=yuvj420p"

        # -ss 放在 -i 前面（input seeking，极快），-vframes 1 只取一帧
        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"Cover frame extraction 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:
    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,
    *,
    seek_ratio: float = 0.15,
) -> str:
    """从视频中提取一帧缩略图并上传到 OSS。"""
    from video_processing.oss_helpers import upload_to_oss

    tmp = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
    tmp.close()
    try:
        frame_path = extract_first_frame(video_path, output_path=tmp.name, seek_ratio=seek_ratio)
        url = upload_to_oss(frame_path, storage_key)
        if not url:
            raise RuntimeError(f"上传缩略图到 OSS 失败: {storage_key}")
        return url
    finally:
        Path(tmp.name).unlink(missing_ok=True)


def _compute_clip_boundary_seek_points(
    duration: float,
    clip_boundaries: Optional[list[tuple[float, float]]] = None,
    num_frames: int = 5,
    head_skip_ratio: float = 0.08,
    tail_skip_ratio: float = 0.08,
) -> list[float]:
    """基于clip分段边界计算抽帧时间点（取每段中间帧，效果比均匀抽更好）。

    策略：
    - 如果传入 clip_boundaries（每个元素是 (clip_start_in_timeline, clip_duration)），
      取每个片段的中点作为抽帧候选点
    - 候选点不足 num_frames 时，均匀补充
    - 跳过片头 head_skip_ratio（8%，避免片头黑屏/开场标题）和片尾 tail_skip_ratio（8%）
    - 返回按时间排序的 num_frames 个抽帧点（秒）
    """
    if duration <= 0:
        # 无法probe，均匀分布兜底
        return [max(1.0, duration * (0.1 + 0.8 * i / max(num_frames - 1, 1))) for i in range(num_frames)]

    head_skip = duration * head_skip_ratio
    tail_skip = duration * tail_skip_ratio
    valid_start = head_skip
    valid_end = max(valid_start + 1.0, duration - tail_skip)

    candidates: list[float] = []

    if clip_boundaries:
        # 累加timeline start，取每clip中点
        cur = 0.0
        for _clip_start, clip_dur in clip_boundaries:
            if clip_dur <= 0:
                continue
            mid = cur + clip_dur / 2.0
            if valid_start <= mid <= valid_end:
                candidates.append(mid)
            cur += clip_dur
        # 去重+排序
        candidates = sorted(set(round(c, 3) for c in candidates))

    # 如果候选点不足，均匀补充
    if len(candidates) < num_frames:
        needed = num_frames - len(candidates)
        existing = set(round(c, 1) for c in candidates)
        for i in range(needed * 3):
            ratio = 0.1 + 0.8 * (i + 0.5) / (needed * 3)
            t = valid_start + (valid_end - valid_start) * ratio
            if round(t, 1) not in existing:
                candidates.append(t)
                existing.add(round(t, 1))
            if len(candidates) >= num_frames:
                break

    # 如果还不够，强制均匀
    while len(candidates) < num_frames:
        idx = len(candidates)
        ratio = 0.1 + 0.8 * idx / max(num_frames - 1, 1)
        candidates.append(valid_start + (valid_end - valid_start) * ratio)

    candidates.sort()

    # 如果超过num_frames，均匀选取
    if len(candidates) > num_frames:
        step = len(candidates) / num_frames
        candidates = [candidates[int(i * step)] for i in range(num_frames)]

    return [round(t, 3) for t in candidates[:num_frames]]


def _extract_frames_via_mediakit(
    video_path: str,
    plan_id: str,
    num_frames: int,
) -> list[dict] | None:
    """使用 MediaKit 智能抽帧 API 提取封面帧（fallback 路径，默认不启用）。"""
    import uuid

    from video_processing.oss_helpers import delete_from_oss, get_signed_download_url, upload_to_oss

    from packages.shared.mediakit_client import get_mediakit_client

    client = get_mediakit_client()
    if not client.is_available:
        logger.info("[thumbnail] MediaKit 未配置，跳过智能抽帧")
        return None

    video_storage_key: str = ""
    try:
        video_storage_key = f"temp/{plan_id}/{uuid.uuid4().hex[:8]}_{Path(video_path).name}"
        public_url = upload_to_oss(video_path, video_storage_key)
        if not public_url:
            logger.warning("[thumbnail] 视频上传 OSS 失败，无法使用 MediaKit")
            return None
        video_url = get_signed_download_url(video_storage_key, expires_seconds=3600) or public_url
        logger.info("[thumbnail] 视频已上传 OSS 并生成签名 URL: key=%s", video_storage_key[:80])
    except Exception as e:
        logger.warning("[thumbnail] 视频上传 OSS 异常: %s，降级到本地 ffmpeg", e)
        return None

    try:
        frames = client.extract_frames(
            video_url=video_url,
            strategy="SceneChange",
            max_frames=num_frames * 2,
        )
        if not frames:
            logger.warning("[thumbnail] MediaKit 抽帧返回空")
            return None
        if len(frames) > num_frames:
            step = len(frames) // num_frames
            frames = [frames[i * step] for i in range(num_frames)]
        logger.info("[thumbnail] MediaKit 抽帧成功: %d 帧", len(frames))
        return frames
    except Exception as e:
        logger.warning("[thumbnail] MediaKit 抽帧异常: %s", e)
        return None
    finally:
        try:
            delete_from_oss(video_storage_key)
        except Exception:
            pass


def extract_and_upload_cover_frames(
    video_path: str,
    plan_id: str,
    *,
    task_id: str = "",
    num_frames: int = 5,
    title_text: str = "",
    title_color: str = "#ffffff",
    title_position: str = "bottom",
    title_font_size: int | None = None,
    clip_boundaries: Optional[list[tuple[float, float]]] = None,
) -> list[dict]:
    """从视频中抽取多帧作为封面候选，通过质量评分选出最佳帧，上传到 OSS。

    默认路径（P1优化）：本地 ffmpeg 单帧 seek 抽帧 + cv2 评分，预期 <2s 完成。
    - 基于 clip 分段边界取各段中间帧（clip_boundaries 参数），效果优于均匀抽帧
    - 无边界信息时均匀分布（10%~90% 之间）
    - 所有帧本地 cv2 清晰度/亮度/色彩三维评分，最高分自动选出

    Fallback（MEDIAKIT_COVER_ENABLED=true）：火山 MediaKit SceneChange 抽帧（~60-90s）。

    Args:
        clip_boundaries: 片段边界列表 [(clip_start, clip_duration), ...]，用于智能取点
    """
    import time

    import httpx
    from video_processing.ffmpeg_utils import probe_duration
    from video_processing.oss_helpers import upload_to_oss

    from packages.shared.config import get_shared_settings

    t0 = time.monotonic()

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

    candidates: list[dict] = []
    _temp_paths: list[str] = []

    try:
        settings = get_shared_settings()
        use_mediakit = getattr(settings, "mediakit_cover_enabled", False)

        if use_mediakit:
            logger.info("[thumbnail] MEDIAKIT_COVER_ENABLED=true，走 MediaKit 路径")
            mediakit_frames = _extract_frames_via_mediakit(video_path, plan_id, num_frames)
            if mediakit_frames:
                for i, frame in enumerate(mediakit_frames):
                    frame_url = frame.get("image_url")
                    if not frame_url:
                        continue
                    tmp = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
                    tmp.close()
                    _temp_paths.append(tmp.name)
                    try:
                        resp = httpx.get(frame_url, timeout=30, follow_redirects=True)
                        resp.raise_for_status()
                        with open(tmp.name, "wb") as f:
                            f.write(resp.content)
                        if title_text and title_text.strip():
                            apply_title_overlay(
                                tmp.name,
                                title_text,
                                color=title_color,
                                position=title_position,
                                font_size=title_font_size,
                            )
                        storage_key = f"covers/{plan_id}/{task_id}/mediakit_frame_{i}.jpg"
                        url = upload_to_oss(tmp.name, storage_key)
                        if url:
                            candidates.append(
                                {
                                    "url": url,
                                    "position": round(frame.get("timestamp", 0.0), 2),
                                    "image_path": tmp.name,
                                }
                            )
                    except Exception as e:
                        logger.warning("[thumbnail] MediaKit 帧 %d 处理失败: %s", i, e)
                if len(candidates) >= num_frames:
                    logger.info("[thumbnail] MediaKit 抽帧完成: %d 帧", len(candidates))

        # ── 默认路径：本地 ffmpeg 单帧 seek ───────────────────────────
        if len(candidates) < num_frames:
            if candidates:
                logger.info("[thumbnail] MediaKit 不足 %d 帧，本地 ffmpeg 补充", num_frames)
            else:
                logger.info("[thumbnail] 使用本地 ffmpeg 抽帧（num=%d, duration=%.1fs）", num_frames, duration)

            seek_points = _compute_clip_boundary_seek_points(duration, clip_boundaries, num_frames)

            for i, seek_t in enumerate(seek_points):
                tmp = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
                tmp.close()
                _temp_paths.append(tmp.name)
                try:
                    frame_path = extract_first_frame(
                        video_path,
                        output_path=tmp.name,
                        seek_seconds=seek_t,
                        min_seek_seconds=0.5,
                    )
                    if title_text and title_text.strip():
                        apply_title_overlay(
                            frame_path,
                            title_text,
                            color=title_color,
                            position=title_position,
                            font_size=title_font_size,
                        )
                    storage_key = f"covers/{plan_id}/{task_id}/frame_{i}.jpg"
                    url = upload_to_oss(frame_path, storage_key)
                    if url:
                        candidates.append(
                            {
                                "url": url,
                                "position": seek_t,
                                "image_path": tmp.name,
                            }
                        )
                except Exception as e:
                    logger.warning("[thumbnail] 封面候选帧 %d 提取失败: %s", i, e)

        # ── 阶段 2：质量评分 ────────────────────────────────────────────
        if len(candidates) > 1:
            try:
                from packages.shared.cover_frame_scorer import score_frames

                candidates = score_frames(candidates)
                elapsed = time.monotonic() - t0
                logger.info(
                    "[thumbnail] 封面帧评分完成: plan_id=%s count=%d best_score=%.1f elapsed=%.2fs path=%s",
                    plan_id,
                    len(candidates),
                    candidates[0].get("score", 0.0) if candidates else 0.0,
                    elapsed,
                    "mediakit" if use_mediakit else "local",
                )
            except Exception:
                logger.warning("[thumbnail] 封面帧质量评分失败，保持原始顺序", exc_info=True)

        for c in candidates:
            c.pop("image_path", None)

        return candidates
    finally:
        for path in _temp_paths:
            try:
                Path(path).unlink(missing_ok=True)
            except Exception:
                pass