Files
video-gen/video-gen-api/app/services/home_material/watermark_processor.py
T
2026-06-30 17:10:51 +08:00

623 lines
26 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
from __future__ import annotations
import asyncio
import base64
import json
import math
import os
import shutil
import subprocess
import tempfile
from dataclasses import dataclass
from decimal import Decimal
from pathlib import Path
from typing import Any
from fastapi import HTTPException, status
try:
from PIL import Image, ImageDraw, ImageFont
except Exception: # pragma: no cover - 运行时给出明确错误
Image = None # type: ignore[assignment]
ImageDraw = None # type: ignore[assignment]
ImageFont = None # type: ignore[assignment]
from app.config import settings
from app.enums.home_material import (
HomeMaterialMediaType,
HomeMaterialWatermarkPosition,
HomeMaterialWatermarkSizeMode,
HomeMaterialWatermarkType,
)
from app.schemas.home_material import HomeMaterialTextWatermarkConfig, HomeMaterialWatermarkConfig
@dataclass(frozen=True)
class MediaProbeInfo:
width: int | None = None
height: int | None = None
duration_seconds: Decimal | None = None
@dataclass(frozen=True)
class WatermarkProcessResult:
output_path: str
width: int | None
height: int | None
duration_seconds: Decimal | None
file_size_bytes: int
cover_path: str | None = None
class HomeMaterialWatermarkProcessor:
"""首页素材 FFmpeg 水印处理器。纯处理层,不访问数据库。"""
def __init__(self) -> None:
self._image_semaphore = asyncio.Semaphore(int(getattr(settings, "HOME_MATERIAL_IMAGE_WATERMARK_CONCURRENCY", 4)))
self._video_semaphore = asyncio.Semaphore(int(getattr(settings, "HOME_MATERIAL_VIDEO_WATERMARK_CONCURRENCY", 2)))
def _ffmpeg_bin(self) -> str:
configured = getattr(settings, "FFMPEG_BIN", "") or ""
if configured:
return configured
found = shutil.which("ffmpeg") or shutil.which("ffmpeg.exe")
if not found:
raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="未找到 ffmpeg,请配置 FFMPEG_BIN 或安装 ffmpeg")
return found
def _ffprobe_bin(self) -> str:
configured = getattr(settings, "FFMPEG_BIN", "") or ""
if configured:
p = Path(configured)
candidate = p.with_name("ffprobe.exe" if p.name.endswith(".exe") else "ffprobe")
if candidate.exists():
return str(candidate)
found = shutil.which("ffprobe") or shutil.which("ffprobe.exe")
if not found:
raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="未找到 ffprobe,请确认 ffmpeg 环境完整")
return found
@staticmethod
def _run_blocking(args: list[str], timeout: int) -> tuple[str, str]:
"""在线程里执行 FFmpeg/ffprobe。
Windows 下 uvicorn/watchfiles 有概率使用不支持子进程的 SelectorEventLoop
asyncio.create_subprocess_exec 会直接抛 NotImplementedError。这里统一改为
subprocess.run + asyncio.to_thread,仍然不会阻塞 Web Server 事件循环。
"""
creationflags = 0
if os.name == "nt":
creationflags = getattr(subprocess, "CREATE_NO_WINDOW", 0)
try:
completed = subprocess.run(
args,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
timeout=timeout,
check=False,
creationflags=creationflags,
)
except subprocess.TimeoutExpired as exc:
raise RuntimeError("FFmpeg 处理超时") from exc
out = completed.stdout.decode("utf-8", errors="ignore")
err = completed.stderr.decode("utf-8", errors="ignore")
if completed.returncode != 0:
raise RuntimeError(err[-2000:] or f"FFmpeg 退出码异常:{completed.returncode}")
return out, err
async def _run(self, args: list[str], timeout: int | None = None) -> tuple[str, str]:
effective_timeout = timeout or int(getattr(settings, "HOME_MATERIAL_FFMPEG_TIMEOUT_SECONDS", 600))
return await asyncio.to_thread(self._run_blocking, args, effective_timeout)
async def probe(self, path: str) -> MediaProbeInfo:
args = [
self._ffprobe_bin(),
"-v",
"error",
"-select_streams",
"v:0",
"-show_entries",
"stream=width,height,duration",
"-of",
"json",
path,
]
stdout, _ = await self._run(args, timeout=30)
try:
data = json.loads(stdout or "{}")
stream = (data.get("streams") or [{}])[0]
duration_raw = stream.get("duration")
return MediaProbeInfo(
width=int(stream["width"]) if stream.get("width") is not None else None,
height=int(stream["height"]) if stream.get("height") is not None else None,
duration_seconds=Decimal(str(duration_raw)).quantize(Decimal("0.001")) if duration_raw not in (None, "N/A") else None,
)
except Exception:
return MediaProbeInfo()
@staticmethod
def _value(config: HomeMaterialWatermarkConfig | HomeMaterialTextWatermarkConfig | dict[str, Any], key: str, default: Any = None) -> Any:
if isinstance(config, dict):
return config.get(key, default)
return getattr(config, key, default)
def _watermark_type(self, config: HomeMaterialWatermarkConfig | dict[str, Any]) -> str:
raw = self._value(config, "watermark_type", None) or HomeMaterialWatermarkType.IMAGE.value
return HomeMaterialWatermarkType(raw).value
def _watermark_width(self, source_width: int | None, config: HomeMaterialWatermarkConfig | dict[str, Any]) -> int:
size_mode = self._value(config, "size_mode", HomeMaterialWatermarkSizeMode.RATIO.value)
if hasattr(size_mode, "value"):
size_mode = size_mode.value
width_px = self._value(config, "width_px", None)
width_ratio = self._value(config, "width_ratio", 0.18) or 0.18
if size_mode == HomeMaterialWatermarkSizeMode.PX.value and width_px:
return max(1, int(width_px))
base_width = source_width or 1080
return max(1, int(base_width * float(width_ratio)))
def _overlay_xy(self, config: HomeMaterialWatermarkConfig | dict[str, Any]) -> tuple[str, str]:
position = self._value(config, "position", HomeMaterialWatermarkPosition.BOTTOM_RIGHT.value)
if hasattr(position, "value"):
position = position.value
margin_x = int(self._value(config, "margin_x", 24) or 24)
margin_y = int(self._value(config, "margin_y", 24) or 24)
custom_x_ratio = self._value(config, "custom_x_ratio", None)
custom_y_ratio = self._value(config, "custom_y_ratio", None)
if position == HomeMaterialWatermarkPosition.CUSTOM.value:
x_ratio = float(custom_x_ratio if custom_x_ratio is not None else 0.5)
y_ratio = float(custom_y_ratio if custom_y_ratio is not None else 0.5)
return f"(main_w-overlay_w)*{x_ratio:.6f}", f"(main_h-overlay_h)*{y_ratio:.6f}"
positions: dict[str, tuple[str, str]] = {
HomeMaterialWatermarkPosition.TOP_LEFT.value: (str(margin_x), str(margin_y)),
HomeMaterialWatermarkPosition.TOP_CENTER.value: ("(main_w-overlay_w)/2", str(margin_y)),
HomeMaterialWatermarkPosition.TOP_RIGHT.value: (f"main_w-overlay_w-{margin_x}", str(margin_y)),
HomeMaterialWatermarkPosition.MIDDLE_LEFT.value: (str(margin_x), "(main_h-overlay_h)/2"),
HomeMaterialWatermarkPosition.CENTER.value: ("(main_w-overlay_w)/2", "(main_h-overlay_h)/2"),
HomeMaterialWatermarkPosition.MIDDLE_RIGHT.value: (f"main_w-overlay_w-{margin_x}", "(main_h-overlay_h)/2"),
HomeMaterialWatermarkPosition.BOTTOM_LEFT.value: (str(margin_x), f"main_h-overlay_h-{margin_y}"),
HomeMaterialWatermarkPosition.BOTTOM_CENTER.value: ("(main_w-overlay_w)/2", f"main_h-overlay_h-{margin_y}"),
HomeMaterialWatermarkPosition.BOTTOM_RIGHT.value: (f"main_w-overlay_w-{margin_x}", f"main_h-overlay_h-{margin_y}"),
}
return positions.get(str(position), positions[HomeMaterialWatermarkPosition.BOTTOM_RIGHT.value])
def _opacity(self, config: HomeMaterialWatermarkConfig | HomeMaterialTextWatermarkConfig | dict[str, Any]) -> float:
level = int(self._value(config, "opacity_level", 6) or 6)
return min(max(level, 1), 10) / 10
def _filter_complex_image_watermark(self, source_width: int | None, config: HomeMaterialWatermarkConfig | dict[str, Any]) -> str:
wm_width = self._watermark_width(source_width, config)
opacity = self._opacity(config)
x, y = self._overlay_xy(config)
return f"[1:v]format=rgba,colorchannelmixer=aa={opacity:.2f},scale={wm_width}:-1[wm];[0:v][wm]overlay={x}:{y}[v]"
@staticmethod
def _hex_to_rgb(value: str) -> tuple[int, int, int]:
raw = (value or "#ffffff").strip()
if not raw.startswith("#") or len(raw) != 7:
raw = "#ffffff"
try:
return int(raw[1:3], 16), int(raw[3:5], 16), int(raw[5:7], 16)
except Exception:
return 255, 255, 255
def _font_path(self) -> str:
configured = str(getattr(settings, "HOME_MATERIAL_TEXT_WATERMARK_FONT", "") or "").strip()
if not configured:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="重复文字水印字体未配置,请配置 HOME_MATERIAL_TEXT_WATERMARK_FONT 为开源可商用字体文件路径,例如 Noto Sans CJK SC / Source Han Sans SC。",
)
path = Path(configured)
if not path.exists() or not path.is_file():
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail=f"重复文字水印字体文件不存在或不可读:{configured}",
)
return str(path)
def _assert_pillow_available(self) -> None:
if Image is None or ImageDraw is None or ImageFont is None:
raise HTTPException(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
detail="重复文字水印依赖 Pillow,请先安装 pillow,并配置 HOME_MATERIAL_TEXT_WATERMARK_FONT。",
)
def _text_config(self, config: HomeMaterialWatermarkConfig | HomeMaterialTextWatermarkConfig | dict[str, Any]) -> HomeMaterialTextWatermarkConfig | dict[str, Any]:
if isinstance(config, HomeMaterialTextWatermarkConfig):
return config
if isinstance(config, HomeMaterialWatermarkConfig):
if config.text_watermark is None:
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="重复文字水印配置不能为空")
return config.text_watermark
text_config = config.get("text_watermark")
if not isinstance(text_config, dict):
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="重复文字水印配置不能为空")
return text_config
def generate_repeated_text_layer_file(
self,
*,
width: int,
height: int,
text_config: HomeMaterialTextWatermarkConfig | dict[str, Any],
output_path: str,
) -> str:
"""用固定开源字体生成透明重复文字水印层。
该方法同时用于前端精准预览和最终 FFmpeg overlay,保证前端看到的透明层和实际叠加层来自同一套渲染逻辑。
"""
self._assert_pillow_available()
font_path = self._font_path()
safe_width = max(1, int(width))
safe_height = max(1, int(height))
max_text_length = int(getattr(settings, "HOME_MATERIAL_TEXT_WATERMARK_MAX_TEXT_LENGTH", 64))
max_font_size = int(getattr(settings, "HOME_MATERIAL_TEXT_WATERMARK_MAX_FONT_SIZE", 160))
max_gap = int(getattr(settings, "HOME_MATERIAL_TEXT_WATERMARK_MAX_GAP", 2000))
text = str(self._value(text_config, "text", "") or "").strip()[:max_text_length]
if not text:
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="重复文字水印内容不能为空")
font_size = max(8, min(max_font_size, int(self._value(text_config, "font_size_px", 28) or 28)))
rotate_deg = max(-90, min(90, int(self._value(text_config, "rotate_deg", -30) or 0)))
gap_x = max(20, min(max_gap, int(self._value(text_config, "gap_x", 220) or 220)))
gap_y = max(20, min(max_gap, int(self._value(text_config, "gap_y", 140) or 140)))
staggered = bool(self._value(text_config, "staggered", True))
opacity = self._opacity(text_config)
r, g, b = self._hex_to_rgb(str(self._value(text_config, "color", "#ffffff") or "#ffffff"))
alpha = int(round(opacity * 255))
font = ImageFont.truetype(font_path, font_size)
measure = Image.new("RGBA", (1, 1), (0, 0, 0, 0))
measure_draw = ImageDraw.Draw(measure)
bbox = measure_draw.textbbox((0, 0), text, font=font)
text_w = max(1, bbox[2] - bbox[0])
text_h = max(1, bbox[3] - bbox[1])
padding = max(8, int(font_size * 0.8))
text_img = Image.new("RGBA", (text_w + padding * 2, text_h + padding * 2), (0, 0, 0, 0))
text_draw = ImageDraw.Draw(text_img)
text_draw.text((padding - bbox[0], padding - bbox[1]), text, font=font, fill=(r, g, b, alpha))
if rotate_deg:
text_img = text_img.rotate(rotate_deg, resample=Image.Resampling.BICUBIC, expand=True)
layer = Image.new("RGBA", (safe_width, safe_height), (0, 0, 0, 0))
tile_w, tile_h = text_img.size
step_x = max(gap_x, int(tile_w * 0.8))
step_y = max(gap_y, int(tile_h * 0.8))
start_x = -tile_w
start_y = -tile_h
end_x = safe_width + tile_w
end_y = safe_height + tile_h
row = 0
y = start_y
while y <= end_y:
offset = step_x // 2 if staggered and row % 2 == 1 else 0
x = start_x + offset
while x <= end_x:
layer.alpha_composite(text_img, (int(x), int(y)))
x += step_x
y += step_y
row += 1
target = Path(output_path)
target.parent.mkdir(parents=True, exist_ok=True)
layer.save(str(target), format="PNG")
return str(target)
def generate_repeated_text_layer_data_url(
self,
*,
width: int,
height: int,
text_config: HomeMaterialTextWatermarkConfig | dict[str, Any],
) -> str:
max_width = int(getattr(settings, "HOME_MATERIAL_TEXT_WATERMARK_PREVIEW_MAX_WIDTH", 8192))
max_height = int(getattr(settings, "HOME_MATERIAL_TEXT_WATERMARK_PREVIEW_MAX_HEIGHT", 8192))
safe_width = max(1, min(max_width, int(width)))
safe_height = max(1, min(max_height, int(height)))
with tempfile.NamedTemporaryFile(prefix="home_material_text_wm_", suffix=".png", delete=False) as tmp:
tmp_path = tmp.name
try:
self.generate_repeated_text_layer_file(width=safe_width, height=safe_height, text_config=text_config, output_path=tmp_path)
raw = Path(tmp_path).read_bytes()
return "data:image/png;base64," + base64.b64encode(raw).decode("ascii")
finally:
try:
Path(tmp_path).unlink(missing_ok=True)
except Exception:
pass
async def apply_watermark(
self,
*,
media_type: HomeMaterialMediaType | str,
source_path: str,
watermark_path: str | None,
output_path: str,
config: HomeMaterialWatermarkConfig | dict[str, Any],
cover_path: str | None = None,
) -> WatermarkProcessResult:
media = HomeMaterialMediaType(media_type)
watermark_type = self._watermark_type(config)
if media == HomeMaterialMediaType.IMAGE:
async with self._image_semaphore:
if watermark_type == HomeMaterialWatermarkType.REPEATED_TEXT.value:
return await self._apply_image_repeated_text_watermark(source_path, output_path, config)
if not watermark_path:
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="图片水印文件不能为空")
return await self._apply_image_watermark(source_path, watermark_path, output_path, config)
async with self._video_semaphore:
if watermark_type == HomeMaterialWatermarkType.REPEATED_TEXT.value:
return await self._apply_video_repeated_text_watermark(source_path, output_path, config, cover_path=cover_path)
if not watermark_path:
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="图片水印文件不能为空")
return await self._apply_video_watermark(source_path, watermark_path, output_path, config, cover_path=cover_path)
async def _apply_image_watermark(
self,
source_path: str,
watermark_path: str,
output_path: str,
config: HomeMaterialWatermarkConfig | dict[str, Any],
) -> WatermarkProcessResult:
source_info = await self.probe(source_path)
tmp_path = output_path + ".part"
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
filter_complex = self._filter_complex_image_watermark(source_info.width, config)
args = [
self._ffmpeg_bin(),
"-y",
"-i",
source_path,
"-i",
watermark_path,
"-filter_complex",
filter_complex,
"-map",
"[v]",
"-frames:v",
"1",
"-f",
"image2",
"-vcodec",
"png",
tmp_path,
]
try:
await self._run(args)
os.replace(tmp_path, output_path)
finally:
if os.path.exists(tmp_path):
os.remove(tmp_path)
output_info = await self.probe(output_path)
return WatermarkProcessResult(
output_path=output_path,
width=output_info.width or source_info.width,
height=output_info.height or source_info.height,
duration_seconds=None,
file_size_bytes=os.path.getsize(output_path),
)
async def _apply_video_watermark(
self,
source_path: str,
watermark_path: str,
output_path: str,
config: HomeMaterialWatermarkConfig | dict[str, Any],
cover_path: str | None = None,
) -> WatermarkProcessResult:
source_info = await self.probe(source_path)
max_duration = int(getattr(settings, "HOME_MATERIAL_MAX_VIDEO_DURATION_SECONDS", 300))
if source_info.duration_seconds is not None and source_info.duration_seconds > Decimal(max_duration):
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=f"视频时长超限,最大 {max_duration} 秒")
tmp_path = output_path + ".part"
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
filter_complex = self._filter_complex_image_watermark(source_info.width, config)
args = [
self._ffmpeg_bin(),
"-y",
"-i",
source_path,
"-i",
watermark_path,
"-filter_complex",
filter_complex,
"-map",
"[v]",
"-map",
"0:a?",
"-c:v",
"libx264",
"-preset",
"veryfast",
"-crf",
"23",
"-c:a",
"copy",
"-movflags",
"+faststart",
"-f",
"mp4",
tmp_path,
]
try:
await self._run(args)
os.replace(tmp_path, output_path)
finally:
if os.path.exists(tmp_path):
os.remove(tmp_path)
output_info = await self.probe(output_path)
generated_cover = None
if cover_path:
generated_cover = await self.generate_cover(output_path, cover_path)
return WatermarkProcessResult(
output_path=output_path,
width=output_info.width or source_info.width,
height=output_info.height or source_info.height,
duration_seconds=output_info.duration_seconds or source_info.duration_seconds,
file_size_bytes=os.path.getsize(output_path),
cover_path=generated_cover,
)
async def _apply_image_repeated_text_watermark(
self,
source_path: str,
output_path: str,
config: HomeMaterialWatermarkConfig | dict[str, Any],
) -> WatermarkProcessResult:
source_info = await self.probe(source_path)
width = int(source_info.width or 1080)
height = int(source_info.height or 1920)
tmp_path = output_path + ".part"
layer_path = output_path + ".text-layer.png"
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
self.generate_repeated_text_layer_file(width=width, height=height, text_config=self._text_config(config), output_path=layer_path)
args = [
self._ffmpeg_bin(),
"-y",
"-i",
source_path,
"-i",
layer_path,
"-filter_complex",
"[1:v]format=rgba[wm];[0:v][wm]overlay=0:0[v]",
"-map",
"[v]",
"-frames:v",
"1",
"-f",
"image2",
"-vcodec",
"png",
tmp_path,
]
try:
await self._run(args)
os.replace(tmp_path, output_path)
finally:
for p in (tmp_path, layer_path):
if os.path.exists(p):
os.remove(p)
output_info = await self.probe(output_path)
return WatermarkProcessResult(
output_path=output_path,
width=output_info.width or source_info.width,
height=output_info.height or source_info.height,
duration_seconds=None,
file_size_bytes=os.path.getsize(output_path),
)
async def _apply_video_repeated_text_watermark(
self,
source_path: str,
output_path: str,
config: HomeMaterialWatermarkConfig | dict[str, Any],
cover_path: str | None = None,
) -> WatermarkProcessResult:
source_info = await self.probe(source_path)
max_duration = int(getattr(settings, "HOME_MATERIAL_MAX_VIDEO_DURATION_SECONDS", 300))
if source_info.duration_seconds is not None and source_info.duration_seconds > Decimal(max_duration):
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=f"视频时长超限,最大 {max_duration} 秒")
width = int(source_info.width or 1080)
height = int(source_info.height or 1920)
tmp_path = output_path + ".part"
layer_path = output_path + ".text-layer.png"
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
self.generate_repeated_text_layer_file(width=width, height=height, text_config=self._text_config(config), output_path=layer_path)
args = [
self._ffmpeg_bin(),
"-y",
"-i",
source_path,
"-loop",
"1",
"-i",
layer_path,
"-filter_complex",
"[1:v]format=rgba[wm];[0:v][wm]overlay=0:0[v]",
"-map",
"[v]",
"-map",
"0:a?",
"-c:v",
"libx264",
"-preset",
"veryfast",
"-crf",
"23",
"-c:a",
"copy",
"-movflags",
"+faststart",
"-shortest",
"-f",
"mp4",
tmp_path,
]
try:
await self._run(args)
os.replace(tmp_path, output_path)
finally:
for p in (tmp_path, layer_path):
if os.path.exists(p):
os.remove(p)
output_info = await self.probe(output_path)
generated_cover = None
if cover_path:
generated_cover = await self.generate_cover(output_path, cover_path)
return WatermarkProcessResult(
output_path=output_path,
width=output_info.width or source_info.width,
height=output_info.height or source_info.height,
duration_seconds=output_info.duration_seconds or source_info.duration_seconds,
file_size_bytes=os.path.getsize(output_path),
cover_path=generated_cover,
)
async def generate_cover(self, video_path: str, cover_path: str) -> str:
tmp_path = cover_path + ".part"
Path(cover_path).parent.mkdir(parents=True, exist_ok=True)
seek = getattr(settings, "VIDEO_COVER_SEEK_TIME", "00:00:01") or "00:00:01"
width = int(getattr(settings, "VIDEO_COVER_WIDTH", 720))
args = [
self._ffmpeg_bin(),
"-y",
"-ss",
seek,
"-i",
video_path,
"-frames:v",
"1",
"-vf",
f"scale={width}:-2",
"-f",
"image2",
"-vcodec",
"mjpeg",
tmp_path,
]
try:
await self._run(args, timeout=int(getattr(settings, "VIDEO_COVER_TIMEOUT_SECONDS", 15)))
os.replace(tmp_path, cover_path)
return cover_path
except Exception:
if os.path.exists(tmp_path):
os.remove(tmp_path)
fallback_seek = getattr(settings, "VIDEO_COVER_FALLBACK_SEEK_TIME", "00:00:00") or "00:00:00"
fallback_args = args.copy()
fallback_args[fallback_args.index(seek)] = fallback_seek
await self._run(fallback_args, timeout=int(getattr(settings, "VIDEO_COVER_TIMEOUT_SECONDS", 15)))
os.replace(tmp_path, cover_path)
return cover_path
watermark_processor = HomeMaterialWatermarkProcessor()