Files
ocr_data/tool/VideoProcessing/multithread_algo.py
2026-09-28 16:09:31 +08:00

481 lines
21 KiB
Python
Raw Permalink 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 tool.log.log_manager import logger
from config.config_utils import config, get_ffmpeg, get_ffprobe
import shutil
import subprocess
import json
import cv2
import os
import itertools
import threading
import time
def extract_frames(
video_path,
video_name,
start_frame,
end_frame,
output_dir,
frames_to_extract,
frame_prefix=""
):
"""
进行抽帧
:param video_path: 视频地址
:param video_name: 视频名称
:param start_frame: 开始的帧
:param end_frame: 结束的帧
:param output_dir: 输出地址
:param frames_to_extract: 帧时间列表
:param frame_prefix:
:return: 抽取的帧数
"""
cap = cv2.VideoCapture(video_path)
try:
extracted_count = 0
for current_time in frames_to_extract:
# 只抽取在当前线程范围内且预定义的帧
if start_frame <= current_time < end_frame:
current_ms = 1000 * current_time
cap.set(cv2.CAP_PROP_POS_MSEC, current_ms)
ret, frame = cap.read()
if not ret:
print(f"{current_time}无ret")
break
frame_filename = os.path.join(
output_dir,
f"{video_name}{frame_prefix}_{3 * int(current_time) + int((current_time - int(current_time) + 0.1) / 0.33)}.jpg"
)
cv2.imwrite(frame_filename, frame)
extracted_count += 1
logger.log(f"线程 {start_frame}-{end_frame}: 抽取了{extracted_count}帧", program="split_frame")
return extracted_count
except Exception as e:
logger.log(f"线程 {start_frame}-{end_frame} 抽帧异常: {e}", level="error", program="split_frame")
raise
finally:
cap.release()
def _run_extract_frames(result_container, index, *args, **kwargs):
"""包装函数,捕获子线程异常并存储到共享容器"""
try:
result_container[index] = extract_frames(*args, **kwargs)
except Exception as e:
result_container[index] = e
def split_video_get_frame_images(
video_path,
video_name,
num_segments,
output_dir,
interval_seconds=1.0
):
"""
拆分视频,获取视频的帧图片
:param video_path:
:param video_name:
:param num_segments:
:param output_dir:
:param interval_seconds:
:return:
"""
cap = cv2.VideoCapture(video_path)
total_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT)
cap_fps = cap.get(cv2.CAP_PROP_FPS)
print("total frames:", total_frames)
print("fps:", cap_fps)
video_duration = total_frames / cap_fps
cap.release()
frames_to_extract = []
if interval_seconds == 0.33:
total_seconds = int(video_duration)
for seconds in range(total_seconds + 1):
for i in range(3):
target_time = seconds + i * 0.33
if target_time < total_frames:
frames_to_extract.append(target_time)
# 去重排序
frames_to_extract = sorted(set(frames_to_extract))
else:
t = 0.0
while t < video_duration:
frames_to_extract.append(t)
t += interval_seconds
frames_to_extract = sorted(set(frames_to_extract))
frames_to_extract = [tt for tt in frames_to_extract if tt < video_duration]
expected_frames = len(frames_to_extract)
logger.log(f"【视频统计】 视频时长:{video_duration:.2f}秒, 统一计算抽帧位置:{expected_frames}", program="split_frame")
# 分段处理
frames_per_segment = video_duration // num_segments
threads = []
thread_results = {}
for i in range(num_segments):
start_frame = frames_per_segment * i
end_frame = (i + 1) * frames_per_segment if i < num_segments - 1 else video_duration
thread = threading.Thread(
target=_run_extract_frames,
args=(thread_results, i, video_path, video_name, start_frame, end_frame, output_dir, frames_to_extract, "_first")
)
threads.append(thread)
thread.start()
for thread in threads:
thread.join()
# 检查子线程异常
for i, result in thread_results.items():
if isinstance(result, Exception):
raise RuntimeError(f"线程 {i} 抽帧失败: {result}") from result
return expected_frames, video_duration
def get_video_codec(video_path: str) -> str:
"""通过 ffprobe 获取视频编码格式"""
cmd = [
get_ffprobe(), '-v', 'error', '-select_streams', 'v:0',
'-show_entries', 'stream=codec_name',
'-of', 'default=noprint_wrappers=1:nokey=1', video_path
]
result = subprocess.run(cmd, capture_output=True, text=True)
return result.stdout.strip().lower()
# ==================== GPU 硬解检测(结果缓存,只查一次)====================
_gpu_lock = threading.Lock()
_gpu_checked = False
_gpu_available = False
_gpu_count = 0 # 检测到的 GPU 张数(多卡轮询用)
_gpu_rr = itertools.count() # 多卡轮询计数器(next() 在 CPython 下原子)
# GPU 无法解码的编码:首次失败后记下,后续该编码直接走 CPU。
# 假设多卡同型号(同样支持/不支持某编码),故按编码全局标记。
_gpu_bad_codecs = set()
def gpu_decode_available() -> bool:
"""检测本机是否有可用的 NVIDIA GPU 硬解(nvidia-smi 可用 + ffmpeg 支持 cuda),
并数出 GPU 张数。结果缓存,整个进程只真正检测一次。"""
global _gpu_checked, _gpu_available, _gpu_count
if _gpu_checked:
return _gpu_available
with _gpu_lock:
if _gpu_checked:
return _gpu_available
avail = False
count = 0
try:
r = subprocess.run(['nvidia-smi', '-L'], capture_output=True, text=True, timeout=10)
if r.returncode == 0:
count = sum(1 for ln in r.stdout.splitlines() if ln.strip().startswith('GPU '))
r2 = subprocess.run([get_ffmpeg(), '-hide_banner', '-hwaccels'],
capture_output=True, text=True, timeout=10)
avail = count > 0 and 'cuda' in (r2.stdout or '').lower()
except (FileNotFoundError, OSError, subprocess.TimeoutExpired):
avail, count = False, 0
_gpu_available = avail
_gpu_count = count
_gpu_checked = True
logger.log(
f"GPU 硬解检测: {'可用,抽帧优先用 GPU(cuda)' if avail else '不可用,使用 CPU 软解'}"
f"(检测到 {count} 张卡)", program="split_frame")
return _gpu_available
# —— 自动探测空闲 GPU(避开被 co-located OCR 服务占用的卡)——
_gpu_free_cache = {"ts": 0.0, "free": None} # free=空闲卡索引集合; None=查不到(让上层回退)
_gpu_free_lock = threading.Lock()
def _query_gpu_mem_used() -> dict:
"""nvidia-smi 查每张卡【已用显存(MB)】, 返回 {index: mem_mb}; 失败返回 {}。"""
try:
r = subprocess.run(
['nvidia-smi', '--query-gpu=index,memory.used',
'--format=csv,noheader,nounits'],
capture_output=True, text=True, timeout=5)
out = {}
for ln in r.stdout.splitlines():
parts = [p.strip() for p in ln.split(',')]
if len(parts) >= 2 and parts[0].isdigit():
try:
out[int(parts[0])] = float(parts[1])
except ValueError:
pass
return out
except Exception:
return {}
def _get_free_gpus(busy_mem_mb: float, ttl: float = 20.0):
"""返回【空闲】(已用显存 < busy_mem_mb)的 GPU 索引集合; 带 TTL 缓存, 不频繁调 nvidia-smi。
返回 None 表示查不到(nvidia-smi 失败), 让上层回退到原有逻辑。"""
now = time.time()
with _gpu_free_lock:
if _gpu_free_cache["free"] is not None and (now - _gpu_free_cache["ts"]) < ttl:
return _gpu_free_cache["free"]
mem = _query_gpu_mem_used()
free = {i for i, m in mem.items() if m < busy_mem_mb} if mem else None
_gpu_free_cache["free"] = free
_gpu_free_cache["ts"] = now
return free
def _next_gpu_device() -> int:
"""多卡轮询取下一张 GPU 索引;单卡或未知时返回 0。
- gpu_reserve_last_for_whisper: 留最后一张给 whisper(抽帧只用前 N-1 张)
- gpu_skip_indices: 手动【跳过】这些卡
- gpu_auto_select: 自动探测, 只往【空闲卡】(显存占用 < gpu_busy_mem_mb)上分配,
自动避开被 OCR 服务占着、抽帧慢25倍的卡。查不到/全忙时回退正常轮询。"""
if _gpu_count <= 1:
return 0
n = _gpu_count
skip = set()
auto = False
busy_mem = 4096.0
try:
cfg = config.get_splitframe_config()
if cfg.get('gpu_reserve_last_for_whisper', True):
n = _gpu_count - 1 # 留最后一张给 whisper
for i in (cfg.get('gpu_skip_indices') or []):
try:
skip.add(int(i))
except (TypeError, ValueError):
pass
auto = bool(cfg.get('gpu_auto_select', False))
busy_mem = float(cfg.get('gpu_busy_mem_mb', 4096))
except Exception:
pass
# 候选 = 前 n 张里去掉手动跳过的卡
candidates = [i for i in range(max(1, n)) if i not in skip]
# 自动探测: 进一步只保留空闲卡(若探测到至少一张空闲); 全忙或查不到则维持原候选
if auto:
free = _get_free_gpus(busy_mem)
if free:
free_cands = [i for i in candidates if i in free]
if free_cands:
candidates = free_cands
if not candidates:
candidates = list(range(max(1, n))) # 兜底, 不至于不抽帧
return candidates[next(_gpu_rr) % len(candidates)]
def _clean_temp(temp_pattern: str):
"""清理某次抽帧产生的临时帧(GPU 失败回退前用)。"""
d = os.path.dirname(temp_pattern)
prefix = os.path.basename(temp_pattern).split('%')[0]
try:
for f in os.listdir(d):
if f.startswith(prefix):
try:
os.remove(os.path.join(d, f))
except OSError:
pass
except OSError:
pass
def _resolve_ffmpeg_threads(raw) -> int:
"""splitframe.ffmpeg_threads 支持字面数字,或 "auto":按 CPU核数/并发进程数 自动摊分,
避免 30 台核数不一样的机器都得手改这个数。并发进程数从环境变量 OCR_MAX_WORKERS 读
(main.py/run_worker.py 启动时按实际生效的 --max-workers 写入,进程池 forkserver 子进程会继承);
读不到(如手动 import 调试)按 4 估算。0/未配置 = 不限制(交给 ffmpeg 默认行为)。"""
if isinstance(raw, str) and raw.strip().lower() == 'auto':
cpu = os.cpu_count() or 4
try:
workers = int(os.environ.get('OCR_MAX_WORKERS', '0') or 0)
except ValueError:
workers = 0
workers = workers if workers > 0 else 4
return max(1, cpu // workers)
try:
return int(raw)
except (TypeError, ValueError):
return 0
def _run_ffmpeg_extract(video_path: str, fps: float, temp_pattern: str, codec: str) -> str:
"""执行抽帧:有 GPU 且该编码可硬解则用 GPU,失败/不支持自动回退 CPU。返回 'gpu' 或 'cpu'。"""
cfg = config.get_splitframe_config()
use_gpu = cfg.get('use_gpu', True)
# 明确不让 GPU 解的编码(如 V100 解不了 AV1 → 配 ["av1"] 直接走 CPU,省掉一次失败探测)
skip_codecs = {str(c).lower() for c in (cfg.get('gpu_skip_codecs') or [])}
ffmpeg_threads = _resolve_ffmpeg_threads(cfg.get('ffmpeg_threads', 0))
# 字幕固定在下半屏时,抽帧阶段就裁掉上半部,帧变小 → 下游每帧处理(merge读帧/CV/OCR)成本按比例下降。
# crop_bottom_ratio=保留底部比例(如 0.55=保留下 55%);null/缺省=不裁(全帧)。
vf = f'fps={fps:.6f}'
crop_r = cfg.get('crop_bottom_ratio')
try:
crop_r = float(crop_r) if crop_r is not None else None
except (TypeError, ValueError):
crop_r = None
if crop_r is not None and 0.0 < crop_r < 1.0:
vf = f'{vf},crop=iw:ih*{crop_r:.4f}:0:ih*{1.0 - crop_r:.4f}'
out_opts = ['-vf', vf, '-q:v', '2', '-y', temp_pattern, '-loglevel', 'error']
# ① GPU 硬解:支持该编码的卡上用 GPU 解码(多卡轮询分配到不同卡)。
# 解码在 GPU、自动下载到内存,fps/JPEG 编码在 CPU。
# 跳过条件:配置黑名单(gpu_skip_codecs,如 av1) 或 之前解失败过的编码(_gpu_bad_codecs)。
# 其他能解的编码(如 H264/HEVC)走 GPU;某编码首次解失败也会记下、之后转 CPU。
# ffmpeg 抽帧超时(秒): 防止坏视频/网络读卡死时 worker 永久阻塞。默认1小时,够超长视频抽帧。
try:
ff_timeout = int(cfg.get('ffmpeg_timeout', 3600))
except (TypeError, ValueError):
ff_timeout = 3600
if (use_gpu and codec not in skip_codecs and codec not in _gpu_bad_codecs
and gpu_decode_available()):
dev = _next_gpu_device()
gpu_cmd = [get_ffmpeg()]
if ffmpeg_threads > 0:
gpu_cmd += ['-threads', str(ffmpeg_threads)]
gpu_cmd += ['-hwaccel', 'cuda', '-hwaccel_device', str(dev),
'-i', video_path] + out_opts
try:
subprocess.run(gpu_cmd, check=True, capture_output=True, timeout=ff_timeout)
return f'gpu{dev}'
except subprocess.TimeoutExpired:
logger.log(f"GPU 抽帧超时(>{ff_timeout}s),回退 CPU: {os.path.basename(video_path)}",
level="warning", program="split_frame")
_clean_temp(temp_pattern)
except subprocess.CalledProcessError as e:
with _gpu_lock:
_gpu_bad_codecs.add(codec)
err = (e.stderr or b'')
err = err.decode(errors='ignore') if isinstance(err, bytes) else str(err)
logger.log(f"GPU 硬解失败(编码 {codec or '未知'}),该编码后续改用 CPU: {err[:200]}",
level="warning", program="split_frame")
_clean_temp(temp_pattern) # 清掉半成品再回退
# ② CPU 软解(限解码线程,避免多视频并发互抢)
cmd = [get_ffmpeg()]
if ffmpeg_threads > 0:
cmd += ['-threads', str(ffmpeg_threads)]
cmd += ['-i', video_path] + out_opts
# 捕获 stderr:失败时把 ffmpeg 真实报错带出来(否则只看到 exit 255 不知道原因)
try:
r = subprocess.run(cmd, capture_output=True, text=True, timeout=ff_timeout)
except subprocess.TimeoutExpired:
_clean_temp(temp_pattern)
raise RuntimeError(f"ffmpeg(cpu) 抽帧超时(>{ff_timeout}s): {os.path.basename(video_path)}")
if r.returncode != 0:
raise RuntimeError(
f"ffmpeg(cpu) 退出码 {r.returncode}: {(r.stderr or '').strip()[-600:]}")
return 'cpu'
def split_frame_ffmpeg(video_path: str, video_name: str, output_dir: str,
interval_seconds: float) -> tuple:
"""
使用 FFmpeg 提取帧(有 GPU 则硬解加速,否则软解;用于 AV1 等 OpenCV 无法解码的格式)。
输出文件命名与 OpenCV 方案保持一致。
"""
# 获取视频时长与编码格式
probe = subprocess.run([
get_ffprobe(), '-v', 'error', '-select_streams', 'v:0',
'-show_entries', 'stream=duration,codec_name',
'-of', 'json', video_path
], capture_output=True, text=True)
# ffprobe 失败(返回码非0或无输出)时给清楚的错,避免后面 json.loads("") 报看不懂的
# "Expecting value"。常见原因:bin/ffprobe 是 GPU 版但本机 glibc 太老(<2.28)跑不起来。
if probe.returncode != 0 or not (probe.stdout or "").strip():
raise RuntimeError(
f"ffprobe 执行失败(返回码 {probe.returncode});"
f"常见原因:bin/ffprobe 与本机不兼容(如 GPU 版 ffmpeg 需 glibc>=2.28)。"
f"stderr: {(probe.stderr or '').strip()[:300]}")
info = json.loads(probe.stdout)
stream0 = info['streams'][0] if info.get('streams') else {}
video_duration = float(stream0.get('duration', 0) or 0)
codec = (stream0.get('codec_name') or '').lower()
# 流时长缺失(某些封装(尤其AV1/部分mp4)只在容器层有duration, 流里没有)时兜底查容器时长——
# 不然 video_duration=0 会导致 frames_to_extract 算成空列表(下面 for 循环整个不执行), 但
# 实际 ffmpeg 硬解还是会正常跑完整个视频(耗时不受这个影响), 抽完的帧因为"期望帧数=0"在
# 后面的重命名循环里被当成多余的全部删掉 —— 表现为"抽帧完成(0帧, 耗时正常, 视频0s)",
# 花了几分钟GPU时间白干。跟 tool/process.py 的 _probe_duration() 是同一个坑、同一个修法。
if video_duration <= 0:
try:
fmt_probe = subprocess.run([
get_ffprobe(), '-v', 'error',
'-show_entries', 'format=duration', '-of', 'json', video_path
], capture_output=True, text=True)
fmt_info = json.loads(fmt_probe.stdout or '{}')
video_duration = float((fmt_info.get('format') or {}).get('duration', 0) or 0)
except Exception:
pass
# 计算目标时间点(与 OpenCV 方案一致)
frames_to_extract = []
if interval_seconds == 0.33:
for sec in range(int(video_duration) + 1):
for i in range(3):
t = sec + i * 0.33
if t < video_duration:
frames_to_extract.append(t)
else:
t = 0.0
while t < video_duration:
frames_to_extract.append(t)
t += interval_seconds
frames_to_extract = sorted(set(frames_to_extract))
expected = len(frames_to_extract)
# 一次性用 FFmpeg 提取所有帧(fps 滤镜)
# 注意:interval 0.33 是"每秒3帧"的约定值,真正频率应为 3.0。
# 若用 1/0.33=3.0303,会比真正 3fps 快约1%,导致第N帧实际抽自 N×0.33 秒、
# 而 SRT 按 N/3≈N×0.3333 秒标注,产生随帧号累积的时间漂移(18分钟处可差~10秒)。
fps = 3.0 if interval_seconds == 0.33 else 1.0 / interval_seconds
temp_pattern = os.path.join(output_dir, f"__tmp_{video_name}_%08d.jpg")
# 有 GPU 用 GPU 硬解,否则 CPU 软解(CPU 路径按 ffmpeg_threads 限解码线程)
backend = _run_ffmpeg_extract(video_path, fps, temp_pattern, codec)
# 按期望时间点重命名
temp_files = sorted(
f for f in os.listdir(output_dir) if f.startswith(f"__tmp_{video_name}_")
)
extracted = 0
for i, tmp in enumerate(temp_files):
if i >= len(frames_to_extract):
os.remove(os.path.join(output_dir, tmp))
continue
ct = frames_to_extract[i]
frame_idx = 3 * int(ct) + int((ct - int(ct) + 0.1) / 0.33)
new_name = os.path.join(output_dir, f"{video_name}_first_{frame_idx}.jpg")
os.rename(os.path.join(output_dir, tmp), new_name)
extracted += 1
logger.log(f"FFmpeg 抽帧完成({backend}): 应抽{expected}帧, 实际{extracted}帧", program="split_frame")
return expected, video_duration, backend
def split_frame(video_path, extract_interval=0.33):
"""
进行视频分割
:param video_path: 视频本地位置
:param extract_interval: 分割间隔,默认0.33间隔,1秒3张图
:return:
"""
logger.start_program("split_frame")
video_name = os.path.basename(video_path).partition(".")[0]
to_save_dir = os.path.join(os.getcwd(), "split", video_name)
os.makedirs(to_save_dir, exist_ok=True)
logger.log(f"【 正在抽帧处理 】{video_name}, 抽帧间隔:{extract_interval}秒", program="split_frame")
start = time.time()
# 统一使用 FFmpeg 单次顺序解码(比 OpenCV 逐帧 seek 快 3-5 倍)
expected_frames, video_duration, backend = split_frame_ffmpeg(
video_path, video_name, to_save_dir, extract_interval)
end = time.time()
actual_frames = len([f for f in os.listdir(to_save_dir) if f.endswith('.jpg')])
logger.log(f"【抽帧完成】: 应抽{expected_frames}帧, 实际抽取{actual_frames}帧, 耗时{end - start:.2f}秒", program="split_frame")
logger.end_program("split_frame")
return video_duration, expected_frames, actual_frames, backend