151 lines
5.9 KiB
Python
151 lines
5.9 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
下载一个视频 → 抽帧 → 用这些【真实帧】对 OCR 端点做压测。
|
||
比随机生成的图更贴近线上负载（真实字幕内容，OCR 干的活和实际一致）。
|
||
|
||
用法（项目根目录运行）:
|
||
python temp/video_bench.py <视频...> --url http://IP:18081 [选项]
|
||
|
||
<视频...> 一个或多个，cos://bucket/path/x.mp4 或本地 ./video/x.mp4
|
||
多个视频时各抽一部分帧合并成图池(max-frames 在视频间均分)
|
||
也可用 --list 文件 从文件读视频列表(每行一个)
|
||
|
||
选项:
|
||
--url URL OCR 端点(可多次，soak 只用第一个)
|
||
--concurrency N 并发(默认 100；单台拐点约 80)
|
||
--mode soak|sweep soak=持续满速(默认)；sweep=自动加压找极限
|
||
--fps F 抽帧频率(默认 1，压测用不需要 3fps，省时省内存)
|
||
--max-frames N 最多用多少帧做图池(默认 2000)
|
||
--interval N soak 打印间隔秒(默认 5)
|
||
--ffmpeg PATH ffmpeg 路径(默认 ./bin/ffmpeg)
|
||
--coscli PATH coscli 路径(默认 ./coscli-linux-amd64)
|
||
--cos-config PATH coscli 配置(默认 .cos.yaml)
|
||
--keep 跑完不删临时视频/帧
|
||
"""
|
||
import argparse
|
||
import glob
|
||
import os
|
||
import shutil
|
||
import subprocess
|
||
import sys
|
||
import tempfile
|
||
|
||
# 复用同目录 bench 脚本的压测逻辑
|
||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||
import bench_ocr_endpoints as bench # noqa: E402
|
||
|
||
|
||
def download_video(src, coscli, cos_config, dst):
|
||
"""cos:// 路径用 coscli 下载到 dst；本地路径直接返回。"""
|
||
if not src.startswith("cos://"):
|
||
if not os.path.exists(src):
|
||
print(f"本地视频不存在: {src}")
|
||
return None
|
||
return src
|
||
print(f"下载视频: {src}")
|
||
r = subprocess.run([coscli, "-c", cos_config, "cp", src, dst],
|
||
capture_output=True, text=True)
|
||
if r.returncode != 0 or not os.path.exists(dst):
|
||
print(f"下载失败: {r.stderr.strip()[:300]}")
|
||
return None
|
||
print(f" → {dst} ({os.path.getsize(dst)/1024/1024:.0f} MB)")
|
||
return dst
|
||
|
||
|
||
def extract_frames(ffmpeg, video, frame_dir, fps):
|
||
"""用 ffmpeg 抽帧到 frame_dir，返回帧文件列表。"""
|
||
os.makedirs(frame_dir, exist_ok=True)
|
||
print(f"抽帧 (fps={fps}) ...")
|
||
cmd = [ffmpeg, "-i", video, "-vf", f"fps={fps}", "-q:v", "2",
|
||
os.path.join(frame_dir, "f_%06d.jpg"), "-loglevel", "error", "-y"]
|
||
r = subprocess.run(cmd, capture_output=True, text=True)
|
||
if r.returncode != 0:
|
||
print(f"抽帧失败: {r.stderr.strip()[:400]}")
|
||
return []
|
||
frames = sorted(glob.glob(os.path.join(frame_dir, "*.jpg")))
|
||
print(f" 抽出 {len(frames)} 帧")
|
||
return frames
|
||
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument("videos", nargs="*", default=[], help="一个或多个视频(cos://... 或本地路径)")
|
||
ap.add_argument("--list", default=None, help="从文件读视频列表(每行一个，可代替直接传)")
|
||
ap.add_argument("--url", action="append", default=None, help="OCR 端点(可多次)")
|
||
ap.add_argument("--concurrency", type=int, default=100)
|
||
ap.add_argument("--mode", choices=["soak", "sweep"], default="soak")
|
||
ap.add_argument("--fps", type=float, default=1.0)
|
||
ap.add_argument("--max-frames", type=int, default=2000)
|
||
ap.add_argument("--interval", type=int, default=5)
|
||
ap.add_argument("--ffmpeg", default="./bin/ffmpeg")
|
||
ap.add_argument("--coscli", default="./coscli-linux-amd64")
|
||
ap.add_argument("--cos-config", default=".cos.yaml")
|
||
ap.add_argument("--keep", action="store_true")
|
||
args = ap.parse_args()
|
||
|
||
endpoints = bench.get_endpoints(args.url)
|
||
if not endpoints:
|
||
print("没有 OCR 端点，用 --url 指定")
|
||
return
|
||
|
||
# 收集视频列表：命令行参数 + --list 文件
|
||
videos = list(args.videos)
|
||
if args.list and os.path.exists(args.list):
|
||
with open(args.list, encoding="utf-8") as f:
|
||
videos += [ln.strip() for ln in f if ln.strip()]
|
||
if not videos:
|
||
print("没有视频，传一个或多个视频路径，或用 --list 文件")
|
||
return
|
||
# 每个视频均分帧额度，避免一个长视频占满图池
|
||
per_video_cap = max(1, args.max_frames // len(videos))
|
||
|
||
work = tempfile.mkdtemp(prefix="video_bench_")
|
||
try:
|
||
all_frames = []
|
||
for i, src in enumerate(videos):
|
||
print(f"\n[{i+1}/{len(videos)}] {src}")
|
||
vlocal = os.path.join(work, f"video_{i}.mp4")
|
||
fdir = os.path.join(work, f"frames_{i}")
|
||
v = download_video(src, args.coscli, args.cos_config, vlocal)
|
||
if not v:
|
||
continue
|
||
fr = extract_frames(args.ffmpeg, v, fdir, args.fps)
|
||
if len(fr) > per_video_cap:
|
||
fr = fr[:per_video_cap]
|
||
print(f" 取前 {per_video_cap} 帧")
|
||
all_frames += fr
|
||
# 处理完即删该视频，省空间
|
||
if v == vlocal:
|
||
try:
|
||
os.remove(vlocal)
|
||
except OSError:
|
||
pass
|
||
|
||
if not all_frames:
|
||
print("没抽到任何帧")
|
||
return
|
||
|
||
# 帧 → base64 图池
|
||
print(f"\n加载 {len(all_frames)} 帧为图池 ...")
|
||
pool = [bench.img2base64(f) for f in all_frames]
|
||
print("=" * 60)
|
||
print(f"图池: {len(pool)} 张真实帧(来自 {len(videos)} 个视频) | "
|
||
f"端点: {endpoints[0]} | 并发: {args.concurrency}")
|
||
print("=" * 60)
|
||
|
||
# 4) 压测
|
||
if args.mode == "sweep":
|
||
bench.sweep_endpoint(endpoints[0], pool[0], 10, 320, 6)
|
||
else:
|
||
bench.soak_endpoint(endpoints[0], pool, args.concurrency, args.interval)
|
||
|
||
finally:
|
||
if args.keep:
|
||
print(f"\n临时文件保留在: {work}")
|
||
else:
|
||
shutil.rmtree(work, ignore_errors=True)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|