cat-shark
5688f38d62
- 默认 crop 由 [0,0.82,1,0.18] 调整为 [0,0.75,1,0.25]：字幕很少出现在 画面上半部分，扩大裁切高度提升 OCR 召回 - 同步更新 ocr-subtitle 工作流 DAG 与参数覆盖测试
236 lines
9.0 KiB
Python
236 lines
9.0 KiB
Python
"""视频抽帧节点。
|
||
|
||
按**帧间隔**从视频抽取帧：解析视频帧率后，帧间隔 step = round(间隔秒 × fps)，
|
||
用 ffmpeg 的 select 过滤器按帧号（n mod step == 0）精确取帧——每帧都是真实
|
||
视频帧，帧时间 = 帧号 / fps，避免按时间 seek（-ss 秒）造成的取整漂移。
|
||
再按 crop 相对比例 [x,y,w,h]（0~1）裁切出字幕区域，供 vlm-ocr 等下游节点
|
||
使用。产物为 frames.json 清单：[{"time": 帧时间秒, "image_uri": 帧图片路径}, ...]。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import re
|
||
import subprocess
|
||
import time
|
||
from pathlib import Path
|
||
|
||
from nodes.ffmpeg import _ffmpeg_bin
|
||
from wov_app.logging import get_logger
|
||
from wov_sdk.models import InvokeRequest, InvokeResponse
|
||
|
||
logger = get_logger("frame-extract")
|
||
|
||
# 默认裁切：画面底部 1/4 区域（字幕区，字幕很少出现在画面上半部分）。
|
||
DEFAULT_CROP = [0.0, 0.75, 1.0, 0.25]
|
||
|
||
|
||
def _parse_crop(raw) -> list[float] | None:
|
||
"""解析并校验 crop 相对比例 [x,y,w,h]（0~1 且不越出画面）。"""
|
||
try:
|
||
crop = [float(value) for value in raw]
|
||
except (TypeError, ValueError):
|
||
return None
|
||
if len(crop) != 4:
|
||
return None
|
||
x, y, w, h = crop
|
||
if not (0 <= x <= 1 and 0 <= y <= 1 and 0 <= w <= 1 and 0 <= h <= 1):
|
||
return None
|
||
if x + w > 1.001 or y + h > 1.001:
|
||
return None
|
||
return crop
|
||
|
||
|
||
def _video_size(video: Path, ffmpeg_bin: str) -> tuple[int, int] | None:
|
||
"""从 ffmpeg -i 输出解析视频分辨率，避免依赖 ffprobe。"""
|
||
result = subprocess.run(
|
||
[ffmpeg_bin, "-i", str(video)],
|
||
capture_output=True,
|
||
text=True,
|
||
)
|
||
for line in (result.stderr or "").splitlines():
|
||
if "Video:" not in line:
|
||
continue
|
||
match = re.search(r"(\d{2,5})x(\d{2,5})", line)
|
||
if match:
|
||
return int(match.group(1)), int(match.group(2))
|
||
return None
|
||
|
||
|
||
def _video_duration(video: Path, ffmpeg_bin: str) -> float | None:
|
||
"""从 ffmpeg -i 输出解析总时长（秒）。"""
|
||
result = subprocess.run(
|
||
[ffmpeg_bin, "-i", str(video)],
|
||
capture_output=True,
|
||
text=True,
|
||
)
|
||
match = re.search(r"Duration:\s*(\d+):(\d+):(\d+(?:\.\d+)?)", result.stderr or "")
|
||
if match:
|
||
hours = int(match.group(1))
|
||
minutes = int(match.group(2))
|
||
seconds = float(match.group(3))
|
||
return hours * 3600 + minutes * 60 + seconds
|
||
return None
|
||
|
||
|
||
def _video_fps(video: Path, ffmpeg_bin: str) -> float | None:
|
||
"""从 ffmpeg -i 输出解析帧率，支持小数（29.97 fps）与有理数（30000/1001 fps）。"""
|
||
result = subprocess.run(
|
||
[ffmpeg_bin, "-i", str(video)],
|
||
capture_output=True,
|
||
text=True,
|
||
)
|
||
for line in (result.stderr or "").splitlines():
|
||
if "Video:" not in line:
|
||
continue
|
||
match = re.search(r"(\d+)/(\d+)\s*fps", line)
|
||
if match:
|
||
return int(match.group(1)) / int(match.group(2))
|
||
match = re.search(r"(\d+(?:\.\d+)?)\s*fps", line)
|
||
if match:
|
||
return float(match.group(1))
|
||
return None
|
||
|
||
|
||
def _frame_step(fps: float, interval: float) -> int:
|
||
"""帧间隔换算：每 step 帧取一帧（step=round(间隔秒×fps)，至少为 1）。"""
|
||
return max(1, int(round(interval * fps)))
|
||
|
||
|
||
def _parse_progress_line(line: str) -> int | None:
|
||
"""解析 ffmpeg -progress 输出行，返回 frame=N 的数值；其他行返回 None。
|
||
|
||
-progress 每 ~0.5s 输出一组 key=value（frame/fps/out_time_ms/progress 等），
|
||
这里只关心 frame=（已写出的选中帧数），用于实时进度与速度统计。
|
||
"""
|
||
stripped = line.strip()
|
||
if not stripped.startswith("frame="):
|
||
return None
|
||
try:
|
||
return int(stripped.split("=", 1)[1])
|
||
except ValueError:
|
||
return None
|
||
|
||
|
||
def _sorted_frame_files(frames_dir: Path) -> list[Path]:
|
||
"""按文件名中的帧号数值排序返回帧文件列表（自然排序，非字典序）。
|
||
|
||
关键点：ffmpeg 的 %04d 编号在超过 9999 帧后会自动扩为 5 位
|
||
（frame_10000.png 等），此时 sorted() 默认的字典序会把 5 位编号排在
|
||
4 位编号之前（如 "frame_10009" < "frame_1009"），导致帧号回退、
|
||
manifest 时间与图像错位（曾真实发生于 run_339ec7ee437f 的 14236 帧
|
||
任务，全片后半段时间轴全部错乱）。必须解析出帧号按数值排序，
|
||
才能保证"第 k 个文件 = 第 k 个选中帧 = 时间 index*step/fps"成立。
|
||
"""
|
||
def frame_number(path: Path) -> int:
|
||
# 文件名形如 frame_0001.png，取下划线后的数字部分。
|
||
return int(path.stem.split("_", 1)[1])
|
||
|
||
return sorted(frames_dir.glob("frame_*.png"), key=frame_number)
|
||
|
||
|
||
def invoke(request: InvokeRequest) -> InvokeResponse:
|
||
"""按帧间隔抽取并裁切视频帧，输出 frames.json 清单。"""
|
||
video_uri = request.inputs.get("video_uri")
|
||
if not video_uri:
|
||
return InvokeResponse(status="failed", error="video_uri is required")
|
||
video = Path(video_uri)
|
||
if not video.is_file():
|
||
return InvokeResponse(status="failed", error="video file not found")
|
||
|
||
interval = float(request.params.get("interval_seconds", 0.5))
|
||
crop = _parse_crop(request.params.get("crop", DEFAULT_CROP))
|
||
if crop is None:
|
||
return InvokeResponse(status="failed", error="invalid crop")
|
||
if interval <= 0:
|
||
return InvokeResponse(status="failed", error="invalid interval_seconds")
|
||
|
||
ffmpeg_bin = _ffmpeg_bin()
|
||
size = _video_size(video, ffmpeg_bin)
|
||
if size is None:
|
||
return InvokeResponse(status="failed", error="cannot read video size")
|
||
duration = _video_duration(video, ffmpeg_bin)
|
||
if duration is None or duration <= 0:
|
||
return InvokeResponse(status="failed", error="cannot read video duration")
|
||
fps = _video_fps(video, ffmpeg_bin)
|
||
if fps is None or fps <= 0:
|
||
return InvokeResponse(status="failed", error="cannot read video fps")
|
||
|
||
# 把"每多少秒一帧"换算为"每多少帧取一帧"，按帧号取帧是帧精确的。
|
||
step = _frame_step(fps, interval)
|
||
width, height = size
|
||
x_px = int(round(crop[0] * width))
|
||
y_px = int(round(crop[1] * height))
|
||
w_px = max(1, int(round(crop[2] * width)))
|
||
h_px = max(1, int(round(crop[3] * height)))
|
||
|
||
output_dir = Path(request.output_dir)
|
||
frames_dir = output_dir / "frames"
|
||
frames_dir.mkdir(parents=True, exist_ok=True)
|
||
# 单次解码全片：select 按帧号（n mod step == 0）精确取帧，随后对选中帧
|
||
# 裁切字幕区域并压缩到 720p 内（仅缩小）——过大的输入图会触发 glm-ocr
|
||
# 的重复循环 bug（M-RoPE delta），源头规避。
|
||
# -progress 把 frame=N 进度写入管道，实时打印抽帧进度与速度。
|
||
expected_frames = (int(duration * fps) + step - 1) // step
|
||
process = subprocess.Popen(
|
||
[
|
||
ffmpeg_bin,
|
||
"-y",
|
||
"-hide_banner",
|
||
"-loglevel",
|
||
"error",
|
||
"-nostats",
|
||
"-progress",
|
||
"pipe:1",
|
||
"-i",
|
||
str(video),
|
||
"-vf",
|
||
(
|
||
f"select='not(mod(n\\,{step}))',"
|
||
f"crop={w_px}:{h_px}:{x_px}:{y_px},"
|
||
"scale=1280:720:force_original_aspect_ratio=decrease:force_divisible_by=2"
|
||
),
|
||
# 只写出被选中的帧，避免 CFR 补帧产生重复文件。
|
||
"-vsync",
|
||
"vfr",
|
||
str(frames_dir / "frame_%04d.png"),
|
||
],
|
||
stdout=subprocess.PIPE,
|
||
stderr=subprocess.PIPE,
|
||
text=True,
|
||
)
|
||
extract_started = time.monotonic()
|
||
last_logged = 0
|
||
progress_step = max(1, expected_frames // 10)
|
||
for line in process.stdout:
|
||
done = _parse_progress_line(line)
|
||
if done is None:
|
||
continue
|
||
# 每写出约 10 个检查点打一次日志，避免刷屏；同时给出处理速度。
|
||
if done - last_logged >= progress_step:
|
||
elapsed = time.monotonic() - extract_started
|
||
logger.info(
|
||
"抽帧进度 %d/%d 帧 (%.1f 帧/s)",
|
||
min(done, expected_frames), expected_frames,
|
||
done / elapsed if elapsed > 0 else 0.0,
|
||
)
|
||
last_logged = done
|
||
process.wait()
|
||
stderr = process.stderr.read()
|
||
if process.returncode != 0:
|
||
return InvokeResponse(status="failed", error=stderr[-500:] or "ffmpeg failed")
|
||
|
||
# 第 k 个输出文件对应原始帧号 k×step，时间 = 帧号 / fps（帧精确，无累计偏差）。
|
||
files = _sorted_frame_files(frames_dir)
|
||
manifest = [
|
||
{"time": round((index * step) / fps, 3), "image_uri": str(path)}
|
||
for index, path in enumerate(files)
|
||
]
|
||
manifest_path = output_dir / "frames.json"
|
||
manifest_path.write_text(json.dumps(manifest, ensure_ascii=False), encoding="utf-8")
|
||
logger.info("抽帧完成: %d 帧, 帧间隔 %d, 裁切 %dx%d+%d+%d", len(files), step, w_px, h_px, x_px, y_px)
|
||
return InvokeResponse(
|
||
status="completed",
|
||
outputs={"frames_manifest": str(manifest_path), "frame_count": len(files)},
|
||
)
|