首次提交
This commit is contained in:
0
tool/Ocr/__init__.py
Normal file
0
tool/Ocr/__init__.py
Normal file
372
tool/Ocr/tencentcloud_ocr.py
Normal file
372
tool/Ocr/tencentcloud_ocr.py
Normal file
@@ -0,0 +1,372 @@
|
||||
# -*- encoding: utf-8 -*-
|
||||
'''
|
||||
@File : tencentcloud_ocr.py
|
||||
@Time : @
|
||||
@Author : eason
|
||||
@Version : 1.0
|
||||
@Contact :
|
||||
@License : (C)Copyright 2021-2024
|
||||
@Desc : 将图片转换为base64编码，并使用内部Youtu OCR服务进行文字识别（app_id自动生成）
|
||||
'''
|
||||
|
||||
import json
|
||||
import os
|
||||
import base64
|
||||
import random
|
||||
import threading
|
||||
import time
|
||||
import requests
|
||||
from config.config_utils import config # 如果不再需要配置可以移除，但保留其他配置可能使用
|
||||
|
||||
|
||||
# ---- OCR 端点（多 IP 负载均衡）配置 ----
|
||||
OCR_PATH = "/youtu/ocrapi/hpgeneralocr_v3"
|
||||
OCR_DEFAULT_PORT = 18081 # config 里只填 IP 时自动补的端口
|
||||
DEFAULT_ENDPOINTS = ["http://asr.ocr.general.smart.l20.production.polaris:18081"]
|
||||
WEIGHTS_FILE = "./config/ocr_weights.json"
|
||||
|
||||
|
||||
def normalize_endpoint(e):
|
||||
"""把 'IP' / 'IP:port' / 'http://IP:port' 统一成 'http://IP:port'。
|
||||
config 里只填裸 IP 时，自动补 http:// 和默认端口。"""
|
||||
if not e:
|
||||
return None
|
||||
e = str(e).strip().rstrip("/")
|
||||
if not e:
|
||||
return None
|
||||
if not (e.startswith("http://") or e.startswith("https://")):
|
||||
e = "http://" + e
|
||||
host = e.split("://", 1)[1].split("/", 1)[0] # 取 host[:port] 部分
|
||||
if ":" not in host: # 没写端口就补默认端口
|
||||
e = e + f":{OCR_DEFAULT_PORT}"
|
||||
return e
|
||||
|
||||
|
||||
def img2base64(img_path: str) -> str:
|
||||
"""将图片转换为base64编码"""
|
||||
with open(img_path, 'rb') as f:
|
||||
image_base64 = base64.b64encode(f.read()).decode()
|
||||
return image_base64
|
||||
|
||||
|
||||
def get_ocr_endpoints():
|
||||
"""从配置读取 OCR 端点列表，自动归一化（只填 IP 也行）；缺省回退到内置默认。"""
|
||||
ocr_cfg = config.get_tencent_ocr_config() or {}
|
||||
eps = ocr_cfg.get("endpoints")
|
||||
if isinstance(eps, str):
|
||||
eps = [eps]
|
||||
if not eps:
|
||||
eps = DEFAULT_ENDPOINTS
|
||||
return [u for u in (normalize_endpoint(e) for e in eps) if u]
|
||||
|
||||
|
||||
def load_endpoint_weights():
|
||||
"""读取压测产出的 QPS 权重文件 {base_url: qps}；不存在则返回 {}。"""
|
||||
if not os.path.exists(WEIGHTS_FILE):
|
||||
return {}
|
||||
try:
|
||||
with open(WEIGHTS_FILE, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
return {k.rstrip("/"): float(v) for k, v in data.items() if float(v) > 0}
|
||||
except (json.JSONDecodeError, ValueError, TypeError, OSError):
|
||||
return {}
|
||||
|
||||
|
||||
class OcrDispatcher:
|
||||
"""线程安全的 OCR 端点分发器：用 P-of-d-Choices(默认d=3)做【负载/延迟感知】派发——
|
||||
随机取 sample_size 个端点, 把请求发给其中【在途更少 × 更快】的那个。慢/积压端点自动少分。
|
||||
每个进程独立按后端真实快慢反应 → 无需任何协调/中心LB, 也能在 120 台 OCR 上涌现出全局均衡。
|
||||
d 从 2 提到 3(甚至更多): 端点数越多, 随机抽2个都恰好较差的概率越高, 抽样数多一个能明显
|
||||
收窄这个方差, 之后收益递减, 不用无限往上加。
|
||||
失败端点冷却摘除(冷却时长带随机抖动, 防止多个进程对同一端点同时解冷却又一起打回去,
|
||||
造成"整体一起失败-一起恢复-一起打爆"的同步震荡)。多个并发 OCR 线程共享一个实例。"""
|
||||
|
||||
def __init__(self, endpoints, weights=None, cooldown=30.0, sample_size=3):
|
||||
self.endpoints = list(endpoints)
|
||||
# PID 打乱(历史原因保留): 让各进程初始顺序不同; 现在 P-of-d 是随机取样, 影响不大但无害。
|
||||
try:
|
||||
random.Random(os.getpid()).shuffle(self.endpoints)
|
||||
except Exception:
|
||||
pass
|
||||
weights = weights or {}
|
||||
# 压测权重(QPS): 作为"处理能力", 强端点 score 更低 → 多分。无权重默认 1.0。
|
||||
self.weights = {e: max(float(weights.get(e, 1.0)), 0.0001) for e in self.endpoints}
|
||||
self.cooldown = cooldown
|
||||
self.sample_size = max(2, int(sample_size)) # 每次候选抽样个数(P-of-d 的 d)
|
||||
self._lock = threading.Lock()
|
||||
self._inflight = {e: 0 for e in self.endpoints} # 当前在途请求数(本进程视角)
|
||||
self._lat = {e: 0.05 for e in self.endpoints} # 延迟 EWMA(秒), 初始乐观→新端点也会被试
|
||||
self._cooldown_until = {e: 0.0 for e in self.endpoints}
|
||||
self._rng = random.Random(os.getpid() ^ 0x9E3779B1)
|
||||
# —— 派发计数(验证分发是否均匀): 每端点累计派发数, 每 _dump_every 次落盘一次本进程文件 ——
|
||||
self._count = {e: 0 for e in self.endpoints}
|
||||
self._since_dump = 0
|
||||
try:
|
||||
self._dump_every = int((config.get_tencent_ocr_config() or {}).get('dispatch_log_every', 0) or 0)
|
||||
except (ValueError, TypeError, AttributeError):
|
||||
self._dump_every = 0
|
||||
|
||||
def _score(self, e):
|
||||
"""负载分: (在途+1) × 延迟EWMA / 权重。越低越优先。"""
|
||||
return (self._inflight[e] + 1) * self._lat[e] / self.weights[e]
|
||||
|
||||
def pick(self):
|
||||
"""P-of-d: 随机不重复抽 sample_size(默认3) 个可用端点, 选负载分最低的;
|
||||
选中即占一个在途名额(必须配对 done())。全部冷却时退化为全部可用。"""
|
||||
now = time.time()
|
||||
snapshot = None
|
||||
with self._lock:
|
||||
available = [e for e in self.endpoints if self._cooldown_until[e] <= now]
|
||||
if not available:
|
||||
available = list(self.endpoints) # 全挂了也得发, 给一次机会
|
||||
if len(available) == 1:
|
||||
chosen = available[0]
|
||||
else:
|
||||
k = min(self.sample_size, len(available))
|
||||
cands = self._rng.sample(available, k) # 不重复抽样, 比各自 choice() 更有效(不会抽到同一个)
|
||||
chosen = min(cands, key=self._score)
|
||||
self._inflight[chosen] = self._inflight.get(chosen, 0) + 1
|
||||
self._count[chosen] = self._count.get(chosen, 0) + 1
|
||||
if self._dump_every > 0:
|
||||
self._since_dump += 1
|
||||
if self._since_dump >= self._dump_every:
|
||||
self._since_dump = 0
|
||||
snapshot = (dict(self._count),
|
||||
{e: round(self._lat[e], 4) for e in self._lat})
|
||||
if snapshot is not None: # 落盘在锁外, 不阻塞其它线程派发
|
||||
self._write_dump(snapshot[0], snapshot[1])
|
||||
return chosen
|
||||
|
||||
def _write_dump(self, counts, lats):
|
||||
"""把本进程的累计派发计数写到 ./ocr_dispatch/<pid>.json(覆盖式, 单进程单文件无竞争)。"""
|
||||
try:
|
||||
d = "./ocr_dispatch"
|
||||
os.makedirs(d, exist_ok=True)
|
||||
path = os.path.join(d, f"{os.getpid()}.json")
|
||||
data = {"pid": os.getpid(), "ts": time.time(),
|
||||
"total": sum(counts.values()), "counts": counts, "lat_ewma": lats}
|
||||
tmp = path + ".tmp"
|
||||
with open(tmp, "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, ensure_ascii=False)
|
||||
os.replace(tmp, path)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def done(self, endpoint, latency=None, ok=True):
|
||||
"""请求结束回调: 释放在途名额; 成功则更新延迟EWMA, 失败则冷却摘除(冷却时长带±30%抖动,
|
||||
避免同一时刻失败的多个进程在同一时刻解冷却、一起冲回去打同一个刚失败的端点)。"""
|
||||
with self._lock:
|
||||
if self._inflight.get(endpoint, 0) > 0:
|
||||
self._inflight[endpoint] -= 1
|
||||
if endpoint not in self._cooldown_until:
|
||||
return
|
||||
if ok:
|
||||
if latency is not None:
|
||||
self._lat[endpoint] = 0.7 * self._lat[endpoint] + 0.3 * max(latency, 0.001)
|
||||
self._cooldown_until[endpoint] = 0.0
|
||||
else:
|
||||
jitter = self._rng.uniform(0.7, 1.3)
|
||||
self._cooldown_until[endpoint] = time.time() + self.cooldown * jitter
|
||||
|
||||
# —— 兼容旧接口(内部都走 done) ——
|
||||
def report_failure(self, endpoint):
|
||||
self.done(endpoint, ok=False)
|
||||
|
||||
def report_success(self, endpoint):
|
||||
self.done(endpoint, ok=True)
|
||||
|
||||
|
||||
_dispatcher = None
|
||||
_dispatcher_lock = threading.Lock()
|
||||
|
||||
|
||||
def get_dispatcher():
|
||||
"""惰性单例分发器（首次调用时读取端点+权重）。"""
|
||||
global _dispatcher
|
||||
if _dispatcher is None:
|
||||
with _dispatcher_lock:
|
||||
if _dispatcher is None:
|
||||
eps = get_ocr_endpoints()
|
||||
_dispatcher = OcrDispatcher(eps, load_endpoint_weights())
|
||||
return _dispatcher
|
||||
|
||||
|
||||
def _build_payload(image_base64: str) -> dict:
|
||||
timestamp = int(time.time() * 1000)
|
||||
rand_suffix = random.randint(100000, 999999)
|
||||
app_id = f"{timestamp}{rand_suffix}"
|
||||
return {
|
||||
"app_id": app_id,
|
||||
"image": image_base64,
|
||||
"ocr_template": "OCR",
|
||||
"options": {"pdf_page_index": 0, "detect_text": True},
|
||||
}
|
||||
|
||||
|
||||
# 跨进程【全局 OCR 并发上限】信号量。进程池下由 main 在父进程创建、经 initializer 注入每个 worker;
|
||||
# 不管多少进程×多少线程,全局同时在飞的 OCR 请求数被这个信号量硬性卡住,防止打爆端点。
|
||||
# 为 None 时不限流(线程池模式 / 未注入)。
|
||||
_ocr_semaphore = None
|
||||
|
||||
|
||||
def set_ocr_semaphore(sem):
|
||||
"""由 worker 进程初始化时调用,注入全局并发信号量。"""
|
||||
global _ocr_semaphore
|
||||
_ocr_semaphore = sem
|
||||
|
||||
|
||||
# 每个进程一个持久 Session(连接池/keep-alive 复用连接)。
|
||||
# 原来用 requests.post() 每次新建+关闭连接 → 高并发下海量 TIME_WAIT → 本地临时端口耗尽
|
||||
# → "Cannot assign requested address" / 连接超时(表现成"所有端点超时", 其实是本机端口用完)。
|
||||
# 复用连接后, 到每个端点只保持少量长连接, 端口不再churn。惰性创建 → 进程池下各进程独立 Session。
|
||||
_session = None
|
||||
_session_lock = threading.Lock()
|
||||
|
||||
|
||||
def _get_session():
|
||||
global _session
|
||||
if _session is None:
|
||||
with _session_lock:
|
||||
if _session is None:
|
||||
s = requests.Session()
|
||||
adapter = requests.adapters.HTTPAdapter(
|
||||
pool_connections=256, # 连接池数(≥端点数, 每端点一个池)
|
||||
pool_maxsize=64, # 每端点最大并发连接(复用)
|
||||
max_retries=0,
|
||||
pool_block=False)
|
||||
s.mount("http://", adapter)
|
||||
s.mount("https://", adapter)
|
||||
_session = s
|
||||
return _session
|
||||
|
||||
|
||||
def ocr_request_once(base_url: str, image_base64: str, timeout: int = 30) -> str:
|
||||
"""向指定端点发一次 OCR 请求，返回原始响应文本；失败抛异常。
|
||||
受全局信号量限流;用持久 Session 复用连接(防端口耗尽)。"""
|
||||
url = base_url.rstrip("/") + OCR_PATH
|
||||
headers = {"Content-Type": "application/json"}
|
||||
sem = _ocr_semaphore
|
||||
if sem is not None:
|
||||
sem.acquire()
|
||||
try:
|
||||
resp = _get_session().post(url, headers=headers, json=_build_payload(image_base64),
|
||||
timeout=timeout)
|
||||
resp.raise_for_status()
|
||||
return resp.text
|
||||
finally:
|
||||
if sem is not None:
|
||||
sem.release()
|
||||
|
||||
|
||||
def _normalize_youtu_to_tencent(raw: str) -> str:
|
||||
"""将 Youtu OCR 返回格式转为项目其他模块期望的腾讯云 OCR 格式。"""
|
||||
try:
|
||||
data = json.loads(raw)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return raw
|
||||
|
||||
# 已经是腾讯云格式则直接返回
|
||||
if "TextDetections" in data:
|
||||
return raw
|
||||
|
||||
if data.get("error_code") != 0:
|
||||
return json.dumps({"TextDetections": [], "Language": "zh",
|
||||
"RequestId": data.get("session_id", "")}, ensure_ascii=False)
|
||||
|
||||
text_detections = []
|
||||
for recognize_item in data.get("recognize_list", []):
|
||||
item_list = recognize_item.get("item_content", {}).get("item_list", [])
|
||||
for item in item_list:
|
||||
content = item.get("content", "")
|
||||
confidence = int(item.get("conf", 0) * 100) # 0-1 → 0-100
|
||||
|
||||
vertex = item.get("vertex", {})
|
||||
if vertex:
|
||||
polygon = [
|
||||
{"X": vertex.get("x1", 0), "Y": vertex.get("y1", 0)},
|
||||
{"X": vertex.get("x2", 0), "Y": vertex.get("y2", 0)},
|
||||
{"X": vertex.get("x3", 0), "Y": vertex.get("y3", 0)},
|
||||
{"X": vertex.get("x4", 0), "Y": vertex.get("y4", 0)},
|
||||
]
|
||||
else:
|
||||
coord = item.get("coord", {})
|
||||
x, y = coord.get("x", 0), coord.get("y", 0)
|
||||
w, h = coord.get("width", 0), coord.get("height", 0)
|
||||
polygon = [
|
||||
{"X": x, "Y": y},
|
||||
{"X": x + w, "Y": y},
|
||||
{"X": x + w, "Y": y + h},
|
||||
{"X": x, "Y": y + h},
|
||||
]
|
||||
|
||||
text_detections.append({
|
||||
"DetectedText": content,
|
||||
"Confidence": confidence,
|
||||
"Polygon": polygon,
|
||||
})
|
||||
|
||||
return json.dumps({
|
||||
"TextDetections": text_detections,
|
||||
"Language": "zh",
|
||||
"RequestId": data.get("session_id", ""),
|
||||
}, ensure_ascii=False)
|
||||
|
||||
|
||||
def _get_ocr_call_params():
|
||||
"""从 config 读【单次OCR的重试上限/超时】, 带默认。
|
||||
封顶重试是为了防"卡1小时": 端点都超时时, 一张图最坏耗时 = max_attempts × timeout,
|
||||
而不是原来的 len(endpoints)(119) × 30s ≈ 3570s。"""
|
||||
try:
|
||||
cfg = config.get_tencent_ocr_config() or {}
|
||||
max_attempts = int(cfg.get("max_attempts", 4) or 4)
|
||||
timeout = int(cfg.get("timeout", 12) or 12)
|
||||
except (ValueError, TypeError, AttributeError):
|
||||
max_attempts, timeout = 4, 12
|
||||
return max(1, max_attempts), max(1, timeout)
|
||||
|
||||
|
||||
def youtu_ocr_api(image_base64: str) -> str:
|
||||
"""
|
||||
调用内部 Youtu OCR 服务进行文字识别。
|
||||
按 QPS 权重自动在多个 OCR IP 间分发；某端点失败时摘除冷却并换下一个端点重试。
|
||||
重试次数封顶(config: ocr.max_attempts), 单次超时可配(ocr.timeout), 避免端点全超时时拖死 worker。
|
||||
"""
|
||||
dispatcher = get_dispatcher()
|
||||
max_attempts, timeout = _get_ocr_call_params()
|
||||
# 最多试 max_attempts 个端点(且不超过端点总数), 而不是把 119 个全试一遍
|
||||
attempts = min(max(1, len(dispatcher.endpoints)), max_attempts)
|
||||
last_err = None
|
||||
for _ in range(attempts):
|
||||
url = dispatcher.pick()
|
||||
t0 = time.time()
|
||||
try:
|
||||
raw = ocr_request_once(url, image_base64, timeout=timeout)
|
||||
# 把真实延迟反馈给分发器 → 慢端点延迟EWMA升高 → 后续自动少分(负载感知核心)
|
||||
dispatcher.done(url, latency=time.time() - t0, ok=True)
|
||||
return _normalize_youtu_to_tencent(raw)
|
||||
except requests.exceptions.RequestException as e:
|
||||
last_err = e
|
||||
dispatcher.done(url, latency=time.time() - t0, ok=False)
|
||||
continue
|
||||
print(f"Youtu OCR 请求失败（已尝试 {attempts} 个端点, 单次超时 {timeout}s）: {last_err}")
|
||||
return ""
|
||||
|
||||
|
||||
# 为了保持原函数名兼容，将新接口赋值给原名称
|
||||
tencentcloud_ocr_api = youtu_ocr_api
|
||||
|
||||
|
||||
def tencent_ocr_video_all_frames(pic_dir_path: str,
|
||||
jsonl_save_path: str = 'frames_ocr.jsonl') -> list:
|
||||
"""
|
||||
对一个视频抽帧后的全部图片进行 ocr处理并保存起来
|
||||
"""
|
||||
with open(jsonl_save_path, "w", encoding='utf-8') as f:
|
||||
for root, _, files in os.walk(pic_dir_path):
|
||||
for pic_name in sorted(files):
|
||||
if pic_name.endswith(('jpg')):
|
||||
pic_path = os.path.join(root, pic_name)
|
||||
img_base64 = img2base64(img_path=pic_path)
|
||||
json_string = tencentcloud_ocr_api(image_base64=img_base64)
|
||||
f.write("pic_path: {}\n".format(pic_path))
|
||||
f.write(json_string + "\n")
|
||||
0
tool/VideoProcessing/__init__.py
Normal file
0
tool/VideoProcessing/__init__.py
Normal file
480
tool/VideoProcessing/multithread_algo.py
Normal file
480
tool/VideoProcessing/multithread_algo.py
Normal file
@@ -0,0 +1,480 @@
|
||||
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
|
||||
|
||||
0
tool/__init__.py
Normal file
0
tool/__init__.py
Normal file
BIN
tool/__pycache__/__init__.cpython-311.pyc
Normal file
BIN
tool/__pycache__/__init__.cpython-311.pyc
Normal file
Binary file not shown.
BIN
tool/__pycache__/job_cos.cpython-311.pyc
Normal file
BIN
tool/__pycache__/job_cos.cpython-311.pyc
Normal file
Binary file not shown.
2024
tool/getsrt.py
Normal file
2024
tool/getsrt.py
Normal file
File diff suppressed because it is too large
Load Diff
814
tool/job_cos.py
Normal file
814
tool/job_cos.py
Normal file
@@ -0,0 +1,814 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""COS 任务分发 / 成功记录采集。
|
||||
|
||||
核心原语: COS PutObject 带 ForbidOverwrite(x-cos-forbid-overwrite:true) —— 同名对象
|
||||
只能被创建一次, 并发写同一个锁文件只有第一个成功(其余 409)。以此实现跨机器原子认领,
|
||||
无需任何额外服务(Redis/DB)。
|
||||
|
||||
策略: 粘性认领(sticky)。认领某分片的机器是它的唯一所有者;机器重启后只续做自己名下
|
||||
未完成的分片, 别的机器永不抢占 -> 零重复。机器永久报废时, 人工跑 tools/release_machine.py
|
||||
把它名下的任务放回队列。
|
||||
|
||||
COS 布局(jobs_root 下):
|
||||
todo/<shard>.txt 待处理分片(每行一个视频 COS 路径)。完成才删, 是"任务存在性"的唯一凭证。
|
||||
claimed/<shard>.lock 认领锁, 内容 JSON {machine_id, host, claim_ts, hb_ts}。
|
||||
done/<shard>.done 完成标记, 内容 JSON {machine_id, stats, end_ts}。
|
||||
独立位置(各自 config 字段, 可指向别处):
|
||||
daily_backup_path/<日期>/数据_<机器号>.jsonl 每日 数据.jsonl 全量快照备份。
|
||||
daily_backup_path/<日期>/已处理_<机器号>.txt 上面jsonl的轻量版(只有video_cos_path,
|
||||
一行一个), 供 global_processed_set() 全局去重用。
|
||||
long_list_backup_path/todo/超长待重跑_<机器号>_<日期>_<时分秒>.txt 超长视频重跑队列(本身就是另一套 jobs_root)。
|
||||
"""
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import socket
|
||||
import subprocess
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
from qcloud_cos import CosConfig, CosS3Client
|
||||
from qcloud_cos.cos_exception import CosServiceError
|
||||
|
||||
from config.config_utils import config
|
||||
|
||||
# 老版本 cos-python-sdk-v5 的 put_object 不认识 ForbidOverwrite 参数(其 maplist 没有该映射),
|
||||
# 直接传会报 "No Parameter Named ForbidOverwrite"。这里把映射补进去, 使 ForbidOverwrite='true'
|
||||
# 注入 x-cos-forbid-overwrite:true 请求头 —— 认领锁的原子性(同名对象只能创建一次)全靠它。
|
||||
# 不升级 SDK(线上离线机不便升级), 版本无关, 幂等。
|
||||
try:
|
||||
from qcloud_cos import cos_comm as _cos_comm
|
||||
_cos_comm.maplist.setdefault("ForbidOverwrite", "x-cos-forbid-overwrite")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# ---------------------------------------------------------------- 客户端 ----
|
||||
|
||||
_CLIENT = None
|
||||
|
||||
|
||||
def client():
|
||||
"""模块级单例 CosS3Client(复用项目现有 COS 凭据)。"""
|
||||
global _CLIENT
|
||||
if _CLIENT is None:
|
||||
c = config.get_cos_config()
|
||||
_CLIENT = CosS3Client(CosConfig(
|
||||
Region=c["region"], SecretId=c["secret_id"], SecretKey=c["secret_key"],
|
||||
))
|
||||
return _CLIENT
|
||||
|
||||
|
||||
def parse_cos_url(url: str):
|
||||
"""cos://bucket/a/b -> ('bucket', 'a/b')。无 cos:// 前缀时按 bucket 配置补全。"""
|
||||
url = url.strip()
|
||||
if url.startswith("cos://"):
|
||||
rest = url[len("cos://"):]
|
||||
bucket, _, key = rest.partition("/")
|
||||
return bucket, key.strip("/")
|
||||
return config.get_cos_config()["bucket"], url.strip("/")
|
||||
|
||||
|
||||
# ----------------------------------------------------------------- 配置 ----
|
||||
|
||||
def _cfg():
|
||||
return config.get_dispatch_config() or {}
|
||||
|
||||
|
||||
def _detect_ip() -> str:
|
||||
"""取本机对外默认路由的网卡IP:用 UDP socket connect 一个外部地址触发内核选路由
|
||||
(不会真的发包), 拿到的本地地址天然跳过 docker 网桥(如172.17.0.1)/回环这类非默认路由
|
||||
接口 —— 通常就是 SSH 连进来那张网卡的IP(前提:机器只有一条默认路由出口, 常见的
|
||||
'一张主网卡+docker网桥'部署; 若有多张真实网卡且SSH走的不是默认路由那张, 不适用,
|
||||
建议改用 `ip route get 8.8.8.8` 在目标机上核实后再决定)。
|
||||
取不到再退化到 `hostname -I` 第一个; 都拿不到返回空串, 上层回退主机名。
|
||||
仅在 IP 固定(非 DHCP)的机器上用它兜底身份才安全。"""
|
||||
try:
|
||||
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||||
try:
|
||||
s.connect(("8.8.8.8", 80))
|
||||
ip = s.getsockname()[0]
|
||||
finally:
|
||||
s.close()
|
||||
if ip and not ip.startswith("127."):
|
||||
return ip
|
||||
except OSError:
|
||||
pass
|
||||
try:
|
||||
r = subprocess.run(['hostname', '-I'], capture_output=True, text=True, timeout=5)
|
||||
ips = (r.stdout or '').split()
|
||||
return ips[0] if ips else ''
|
||||
except (OSError, subprocess.SubprocessError):
|
||||
return ''
|
||||
|
||||
|
||||
def get_machine_id() -> str:
|
||||
"""机器ID 取值优先级: config.dispatch.machine_id > 环境变量 MACHINE_ID >
|
||||
本机固定IP(默认路由网卡, 见 _detect_ip) > 主机名。
|
||||
必须【唯一且重启后稳定】—— 锁靠它判断'这片是不是我的'。克隆机场景下主机名常常相同,
|
||||
固定IP各机器天然不同,可以直接兜底;但如果 IP 是 DHCP 会变的,不要依赖这个,
|
||||
换IP等于换身份,旧身份名下未完成的锁会变孤儿——那种环境请在 config 里显式填 machine_id。"""
|
||||
mid = (_cfg().get("machine_id") or "").strip()
|
||||
if mid:
|
||||
return mid
|
||||
mid = (os.environ.get("MACHINE_ID") or "").strip()
|
||||
if mid:
|
||||
return mid
|
||||
ip = _detect_ip()
|
||||
if ip:
|
||||
return ip
|
||||
return socket.gethostname()
|
||||
|
||||
|
||||
def is_ocr_colocated() -> bool:
|
||||
"""本机是不是 OCR 服务共置机: 本机默认路由IP是否出现在 config.tencentcloud.ocr.endpoints
|
||||
列表里。这批端点IP本来就是各台跑OCR服务的机器, 如果本机自己也在这个列表里, 说明这台机器
|
||||
除了跑抽帧/OCR客户端之外, 本身还扛着一个OCR服务, 需要多留CPU给它, 不能被抽帧吃满。
|
||||
只覆盖"OCR端点=本机IP"这一种共置方式; 其它不出现在这个列表里的共享方式(比如机器上还跑
|
||||
了别的无关服务)探测不到, 那种情况还是得手动传小一点的 --max-workers。"""
|
||||
try:
|
||||
endpoints = set((config.get_tencent_ocr_config() or {}).get("endpoints") or [])
|
||||
except Exception:
|
||||
return False
|
||||
if not endpoints:
|
||||
return False
|
||||
ip = _detect_ip()
|
||||
return bool(ip) and ip in endpoints
|
||||
|
||||
|
||||
MAX_AUTO_WORKERS = 30 # auto_max_workers() 的硬上限, 见函数docstring③
|
||||
|
||||
|
||||
def auto_max_workers(cpu: int, disk_guard_gb: float = 0.0) -> int:
|
||||
"""--max-workers 未指定时的自动选择:
|
||||
① 默认按CPU核数; 如果 is_ocr_colocated() 判断本机还跑着OCR服务, 打对折让出一半CPU。
|
||||
② 再按本地磁盘算一次上限, 取①②里更小的那个 —— 核多盘小的机器(比如高核数但磁盘只有
|
||||
100多G)按纯CPU数选出的并发数,同时下载/抽帧/拼合成图会把磁盘写爆;`--disk-guard-gb`只挡
|
||||
"新启动"挡不住"已经在跑的", 真正决定磁盘峰值占用的是并发数本身, 所以从源头把并发数卡住:
|
||||
磁盘上限 = (可用磁盘 - disk_guard_gb这层安全垫) ÷ 每视频预估峰值占用(GB)
|
||||
每视频预估占用读 config.dispatch.est_disk_gb_per_video, 默认 2.0G —— 按 subtitle 配置常见的
|
||||
30~2400秒短视频估的(2400秒=3fps≈7200帧, 裁剪后的帧+合成图+源视频下载, 峰值同时存在时
|
||||
量级在1~2G上下; 视频普遍更短/分辨率更低可以调小, 更长/更高清要调大, 保守起见默认按上限估)。
|
||||
disk_guard_gb 传0(不用磁盘保护)时不再额外扣这层安全垫, 直接用当前可用磁盘算。
|
||||
③ 最后再套一个硬上限 MAX_AUTO_WORKERS(30) —— 超高核数机器(如64/128核)按纯CPU数选出的
|
||||
并发数会远超实际有意义的范围(单机瓶颈通常先卡在OCR接口/磁盘IO/内存, 而不是CPU), 且并发
|
||||
数太高会让 ffmpeg_threads="auto" 摊出来的每进程线程数过低(见 multithread_algo.py)。
|
||||
cpu<=0(测不到核数, 比如非Linux)时按10兜底。"""
|
||||
cpu = cpu if cpu and cpu > 0 else 10
|
||||
workers = cpu
|
||||
if is_ocr_colocated():
|
||||
workers = max(1, cpu // 2)
|
||||
logging.info(f"[资源探测] 本机IP在 tencentcloud.ocr.endpoints 里, 判断为OCR共置机, "
|
||||
f"max-workers 自动打对折让资源给OCR服务: {cpu}核 → {workers}")
|
||||
|
||||
try:
|
||||
free_gb = shutil.disk_usage(".").free / 1024 / 1024 / 1024
|
||||
budget_gb = free_gb - max(0.0, disk_guard_gb)
|
||||
est_gb = float(_cfg().get("est_disk_gb_per_video", 2.0) or 2.0)
|
||||
if budget_gb <= 0:
|
||||
logging.warning(f"[资源探测] 可用磁盘({free_gb:.1f}G)已经不够留给磁盘保护垫"
|
||||
f"({disk_guard_gb}G), max-workers 收紧到1, 建议先清理磁盘")
|
||||
workers = 1
|
||||
elif est_gb > 0:
|
||||
disk_cap = max(1, int(budget_gb / est_gb))
|
||||
if disk_cap < workers:
|
||||
logging.info(f"[资源探测] 按磁盘算的并发上限更小, max-workers 收紧: "
|
||||
f"{workers} → {disk_cap}(可用磁盘{free_gb:.1f}G - 保护垫{disk_guard_gb}G = "
|
||||
f"{budget_gb:.1f}G预算 ÷ 每视频预估{est_gb}G)")
|
||||
workers = disk_cap
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
if workers > MAX_AUTO_WORKERS:
|
||||
logging.info(f"[资源探测] max-workers 超过硬上限, 收紧: {workers} → {MAX_AUTO_WORKERS}")
|
||||
workers = MAX_AUTO_WORKERS
|
||||
return workers
|
||||
|
||||
|
||||
_root_override = {"value": None} # run_worker.py 用: 主队列空闲时临时切去认领长视频队列
|
||||
|
||||
|
||||
def set_active_root_override(root):
|
||||
"""临时覆盖 _active_jobs_root() 的返回值(仅当前进程内生效, 不改config/不用重启)。
|
||||
传 None 清除覆盖、恢复按 config 判断。给 run_worker.py 做'主队列没活时顺手兼跑长视频队列'用:
|
||||
覆盖生效期间, claim/list_todo/download_shard/mark_done/release 等全部原语都会转而操作
|
||||
覆盖指向的这个队列根, 跟专职长视频机(run_long_video=true)走的是同一套代码路径, 只是切换是
|
||||
动态的、按需的, 而不是启动时定死。"""
|
||||
_root_override["value"] = root
|
||||
|
||||
|
||||
def _active_jobs_root():
|
||||
"""本 worker 实际认领的队列根, 优先级:
|
||||
_root_override(运行时临时切换) > run_long_video=true(专跑超长视频, 用 long_list_backup_path)
|
||||
> jobs_root(默认主队列)。"""
|
||||
if _root_override["value"]:
|
||||
return _root_override["value"]
|
||||
c = _cfg()
|
||||
if c.get("run_long_video"):
|
||||
root = c.get("long_list_backup_path")
|
||||
if root:
|
||||
return root
|
||||
return c.get("jobs_root") or ""
|
||||
|
||||
|
||||
def _bucket_base():
|
||||
"""从【当前生效的队列根】解析 (bucket, base_prefix)。base 末尾无 '/'。
|
||||
只给 _key()(todo/claimed/done 分片路径)用——这个"当前生效"是有意跟着
|
||||
run_long_video/_root_override 变的(分片本来就该落在当前认领的那个队列下)。"""
|
||||
root = _active_jobs_root()
|
||||
if not root:
|
||||
raise RuntimeError("config.dispatch 未配置队列地址(jobs_root / long_list_backup_path)")
|
||||
bucket, key = parse_cos_url(root)
|
||||
return bucket, key.strip("/")
|
||||
|
||||
|
||||
def _primary_bucket_base():
|
||||
"""固定解析自 jobs_root(主队列), 不受 run_long_video / 运行时 _root_override 影响 ——
|
||||
daily_backup_path / monitor_root 没显式配置时的兜底基准必须用这个, 不能用 _bucket_base()
|
||||
(它会跟 _active_jobs_root() 一起变): 专职长视频机(run_long_video=true)跑起来后
|
||||
_active_jobs_root() 长期指向 long_list_backup_path, 如果兜底基准也跟着它变, 算出来的
|
||||
monitor_root/daily_backup_path 就会落到 .../long_video/monitor、.../long_video/daily 这种
|
||||
跟其它机器不一样的路径下 —— fleet_monitor.py 和 global_processed_set() 只认标准的
|
||||
monitor_root/daily_backup_path, 读不到这些"另起炉灶"的路径, 表现为这台机器的小时摘要/
|
||||
去重快照对外"消失了"(实际数据都在, 只是传错地方)。"""
|
||||
root = _cfg().get("jobs_root") or ""
|
||||
if not root:
|
||||
raise RuntimeError("config.dispatch.jobs_root 未配置")
|
||||
bucket, key = parse_cos_url(root)
|
||||
return bucket, key.strip("/")
|
||||
|
||||
|
||||
def _key(kind: str, shard: str = "") -> str:
|
||||
"""拼某类对象的 key。kind ∈ {todo, claimed, done}。"""
|
||||
_, base = _bucket_base()
|
||||
ext = {"todo": ".txt", "claimed": ".lock", "done": ".done"}[kind]
|
||||
if shard:
|
||||
return f"{base}/{kind}/{shard}{ext}"
|
||||
return f"{base}/{kind}/"
|
||||
|
||||
|
||||
# --------------------------------------------------------------- COS 原语 ----
|
||||
|
||||
def _put(key: str, data: bytes, forbid_overwrite: bool = False):
|
||||
bkt, _ = _bucket_base()
|
||||
kwargs = {}
|
||||
if forbid_overwrite:
|
||||
kwargs["ForbidOverwrite"] = "true" # 经 maplist 注入 x-cos-forbid-overwrite 头
|
||||
client().put_object(Bucket=bkt, Body=data, Key=key, **kwargs)
|
||||
|
||||
|
||||
def _get_bytes(key: str):
|
||||
bkt, _ = _bucket_base()
|
||||
resp = client().get_object(Bucket=bkt, Key=key)
|
||||
return resp["Body"].get_raw_stream().read()
|
||||
|
||||
|
||||
def _exists(key: str) -> bool:
|
||||
bkt, _ = _bucket_base()
|
||||
return client().object_exists(Bucket=bkt, Key=key)
|
||||
|
||||
|
||||
def _delete(key: str):
|
||||
bkt, _ = _bucket_base()
|
||||
try:
|
||||
client().delete_object(Bucket=bkt, Key=key)
|
||||
except CosServiceError as e:
|
||||
if e.get_status_code() != 404:
|
||||
raise
|
||||
|
||||
|
||||
def _list_shards(kind: str):
|
||||
"""列出某前缀下所有分片名(去掉前缀和扩展名)。"""
|
||||
bkt, _ = _bucket_base()
|
||||
prefix = _key(kind)
|
||||
ext = {"todo": ".txt", "claimed": ".lock", "done": ".done"}[kind]
|
||||
out, marker = [], ""
|
||||
while True:
|
||||
resp = client().list_objects(Bucket=bkt, Prefix=prefix, Marker=marker, MaxKeys=1000)
|
||||
for c in resp.get("Contents", []) or []:
|
||||
k = c["Key"]
|
||||
if k.endswith(ext):
|
||||
out.append(k[len(prefix):-len(ext)])
|
||||
if resp.get("IsTruncated") in ("true", True):
|
||||
marker = resp.get("NextMarker") or (resp["Contents"][-1]["Key"] if resp.get("Contents") else "")
|
||||
if not marker:
|
||||
break
|
||||
else:
|
||||
break
|
||||
return out
|
||||
|
||||
|
||||
# ----------------------------------------------------------------- 接口 ----
|
||||
|
||||
def list_todo():
|
||||
return _list_shards("todo")
|
||||
|
||||
|
||||
def shard_matches(name: str, patterns) -> bool:
|
||||
"""分片名是否匹配 patterns(glob 通配, 如 'cos_urls_8*')。
|
||||
patterns 为空/None -> 认领全部(返回 True)。可传单个字符串或列表。
|
||||
匹配时对 '<名>' 和 '<名>.txt' 都试一遍, 写不写 .txt 都行。"""
|
||||
import fnmatch
|
||||
if not patterns:
|
||||
return True
|
||||
if isinstance(patterns, str):
|
||||
patterns = [patterns]
|
||||
cands = (name, name + ".txt")
|
||||
return any(fnmatch.fnmatch(c, p) for p in patterns for c in cands)
|
||||
|
||||
|
||||
def is_done(shard: str) -> bool:
|
||||
return _exists(_key("done", shard))
|
||||
|
||||
|
||||
def read_lock(shard: str):
|
||||
"""读认领锁内容(dict)。不存在返回 None。"""
|
||||
try:
|
||||
return json.loads(_get_bytes(_key("claimed", shard)).decode("utf-8"))
|
||||
except CosServiceError as e:
|
||||
if e.get_status_code() == 404:
|
||||
return None
|
||||
raise
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def claim(shard: str, machine_id: str) -> bool:
|
||||
"""原子认领: 写 claimed/<shard>.lock(禁止覆盖)。成功=True 表示这片归我;
|
||||
409(已被别人锁) -> False。"""
|
||||
now = datetime.now()
|
||||
body = json.dumps({
|
||||
"machine_id": machine_id, "host": socket.gethostname(),
|
||||
"claim_ts": now.isoformat(), "hb_ts": now.isoformat(),
|
||||
}, ensure_ascii=False).encode("utf-8")
|
||||
try:
|
||||
_put(_key("claimed", shard), body, forbid_overwrite=True)
|
||||
return True
|
||||
except CosServiceError as e:
|
||||
if e.get_status_code() == 409:
|
||||
return False # 已被别的机器抢到
|
||||
raise
|
||||
|
||||
|
||||
def heartbeat(shard: str, machine_id: str):
|
||||
"""刷新锁的心跳时间(允许覆盖, 仍写回 machine_id 以保持所有权信息)。"""
|
||||
lk = read_lock(shard) or {}
|
||||
lk.update({"machine_id": machine_id, "host": socket.gethostname(),
|
||||
"hb_ts": datetime.now().isoformat()})
|
||||
lk.setdefault("claim_ts", lk["hb_ts"])
|
||||
_put(_key("claimed", shard), json.dumps(lk, ensure_ascii=False).encode("utf-8"))
|
||||
|
||||
|
||||
def download_shard(shard: str):
|
||||
"""读 todo/<shard>.txt, 返回视频 COS 路径列表。"""
|
||||
text = _get_bytes(_key("todo", shard)).decode("utf-8")
|
||||
return [ln.strip() for ln in text.splitlines() if ln.strip()]
|
||||
|
||||
|
||||
def mark_done(shard: str, machine_id: str, stats: dict):
|
||||
_put(_key("done", shard), json.dumps({
|
||||
"machine_id": machine_id, "stats": stats,
|
||||
"end_ts": datetime.now().isoformat(),
|
||||
}, ensure_ascii=False).encode("utf-8"))
|
||||
|
||||
|
||||
def release(shard: str):
|
||||
"""跑完后清理: 删 todo 和 lock(保留 done 作为完成账)。"""
|
||||
_delete(_key("todo", shard))
|
||||
_delete(_key("claimed", shard))
|
||||
|
||||
|
||||
def my_open_claims(machine_id: str):
|
||||
"""我名下、尚未完成(无 done)的分片 —— 重启后续做用。"""
|
||||
out = []
|
||||
for shard in _list_shards("claimed"):
|
||||
lk = read_lock(shard)
|
||||
if lk and lk.get("machine_id") == machine_id and not is_done(shard):
|
||||
out.append(shard)
|
||||
return out
|
||||
|
||||
|
||||
# ------------------------------------------------------------- 优雅停机 ----
|
||||
|
||||
def stop_requested() -> bool:
|
||||
"""COS 上存在 stop_flag_key(默认 _jobs/STOP)则请求优雅停机。"""
|
||||
flag = _cfg().get("stop_flag_key")
|
||||
if not flag:
|
||||
_, base = _bucket_base()
|
||||
flag = f"{base}/STOP"
|
||||
bkt, _ = _bucket_base()
|
||||
try:
|
||||
return client().object_exists(Bucket=bkt, Key=flag.strip("/"))
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
# ----------------------------------------------------- 每日 数据.jsonl 备份 ----
|
||||
|
||||
# put_object 是"一次性读进内存 + 单个 HTTP 请求"上传, 底层 send_request 默认超时 30 秒 ——
|
||||
# 数据.jsonl 是全量快照, 产能高的机器跑几天就上百 MB/上 GB, 30 秒内传不完就
|
||||
# "The write operation timed out", 然后每几分钟重试又超时, 无限循环从没成功(实测 21.13.80.229)。
|
||||
# 超过这个大小的文件改用 coscli 上传(它支持分块/断点续传, 几 GB 也没问题); 小文件(如几KB的
|
||||
# 小时摘要)仍走 put_object —— 起 coscli 子进程有开销, 小文件没必要。
|
||||
_COSCLI_UPLOAD_THRESHOLD = 50 * 1024 * 1024 # 50MB
|
||||
|
||||
|
||||
def _coscli_put(local_path: str, cos_url: str) -> bool:
|
||||
"""用 coscli 把 local_path 上传到 cos_url(完整 cos:// 路径)。成功 True, 失败 False。
|
||||
coscli 支持分块上传, 大文件不像 put_object 那样受单请求超时限制。
|
||||
coscli 路径/配置从 splitframe 配置读(跟主流程上传 SRT 用的是同一个二进制和配置)。"""
|
||||
try:
|
||||
sf = config.get_splitframe_config() or {}
|
||||
coscli = sf.get("coscli_path", "./coscli-linux-amd64")
|
||||
coscfg = sf.get("cosconfig_path", ".cos.yaml")
|
||||
cmd = [coscli, "-c", coscfg, "cp", local_path, cos_url]
|
||||
r = subprocess.run(cmd, capture_output=True, text=True, timeout=1800)
|
||||
if r.returncode == 0:
|
||||
return True
|
||||
logging.warning(f"[备份] coscli 上传失败(将回退 put_object): {(r.stderr or '').strip()[:200]}")
|
||||
return False
|
||||
except Exception as e:
|
||||
logging.warning(f"[备份] coscli 上传异常(将回退 put_object): {type(e).__name__}: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def _upload_dated_file(local_path: str, dest_root: str, fname: str) -> str:
|
||||
"""把 local_path 上传到 dest_root/<日期>/fname。
|
||||
local_path 不存在 / dest_root 为空 -> 返回 ''。返回上传后的 cos:// 路径。
|
||||
大文件(> _COSCLI_UPLOAD_THRESHOLD)走 coscli 分块上传, 避免 put_object 30 秒超时;
|
||||
coscli 失败(二进制不在/权限等)自动回退 put_object 再试一次 —— 保证引入 coscli 不会
|
||||
比原来更差(最坏情况就是退回原行为)。"""
|
||||
if not local_path or not dest_root or not os.path.exists(local_path):
|
||||
return ""
|
||||
bkt, base_key = parse_cos_url(dest_root)
|
||||
date = datetime.now().strftime("%Y%m%d") # 按日期建文件夹
|
||||
key = f"{base_key.strip('/')}/{date}/{fname}"
|
||||
cos_url = f"cos://{bkt}/{key}"
|
||||
|
||||
size = os.path.getsize(local_path)
|
||||
if size > _COSCLI_UPLOAD_THRESHOLD:
|
||||
logging.info(f"[备份] {fname} 有 {size / 1024 / 1024:.0f}MB, 超过 "
|
||||
f"{_COSCLI_UPLOAD_THRESHOLD // 1024 // 1024}MB, 用 coscli 分块上传")
|
||||
if _coscli_put(local_path, cos_url):
|
||||
return cos_url
|
||||
# coscli 没成: 回退 put_object(大文件很可能还是超时, 但至少不比原来差, 且日志会记清楚)
|
||||
|
||||
with open(local_path, "rb") as f:
|
||||
data = f.read()
|
||||
client().put_object(Bucket=bkt, Body=data, Key=key)
|
||||
return cos_url
|
||||
|
||||
|
||||
def upload_daily_backup(jsonl_path: str, machine_id: str) -> str:
|
||||
"""把本地 数据.jsonl 上传到 dispatch.daily_backup_path/<日期>/数据_<机器号>.jsonl，
|
||||
顺带生成并上传一份轻量的"已处理清单"(见 upload_processed_list) 供 global_processed_set()
|
||||
用——两者是同一次备份动作的副产物，调这一个函数就够了。"""
|
||||
dest_root = _cfg().get("daily_backup_path")
|
||||
if not dest_root:
|
||||
bucket, base = _bucket_base()
|
||||
dest_root = f"cos://{bucket}/{base}/daily"
|
||||
dst = _upload_dated_file(jsonl_path, dest_root, f"数据_{machine_id}.jsonl")
|
||||
if dst:
|
||||
try:
|
||||
upload_processed_list(jsonl_path, machine_id)
|
||||
except Exception as e:
|
||||
logging.warning(f"[每日备份] 已处理清单上传失败(不影响 数据.jsonl 备份): {e}")
|
||||
return dst
|
||||
|
||||
|
||||
def upload_longlist_backup(machine_id: str) -> str:
|
||||
"""把本地超长待重跑列表(subtitle.long_video_list)上传到
|
||||
dispatch.long_list_backup_path/todo/<原文件名>_<机器号>_<日期>_<时分秒>.txt。
|
||||
放进 todo/ 子目录, 使其直接成为可认领的任务分片 —— 把别的机器的 jobs_root
|
||||
指向 long_list_backup_path, 它们就能自动认领这些超长视频重跑。
|
||||
文件名带【机器号+日期+时分秒】: 机器号保证各机器不互相覆盖; 时分秒是因为本函数是
|
||||
"传完即清空本地已传部分"的增量式设计(见下), 一天内不管自动备份触发几次、或人工额外
|
||||
手动跑几次(如 tools/upload_daily.py / apply_update.sh 更新前那一下), 每次快照的都只是
|
||||
"上次传完之后新增的那一小段" —— 如果 key 只精确到天, 同一天第二次上传会用 put_object
|
||||
直接覆盖同名对象, 把第一次已传但可能还没被任何机器认领处理的内容连本地带COS一起冲没。
|
||||
精确到秒后每次上传各自独立成对象, 不会互相覆盖, 只是分片数量会随调用次数变多(不影响使用,
|
||||
worker 一样能各自认领)。"""
|
||||
# 专跑长视频的机器(run_long_video)不再把超长列表传回同一队列, 防回环。
|
||||
if _cfg().get("run_long_video"):
|
||||
return ""
|
||||
local = (config.get_subtitle_config() or {}).get("long_video_list", "./超长待重跑.txt")
|
||||
if not local or not os.path.exists(local):
|
||||
return ""
|
||||
# 快照当前内容(文件 append-only); 全空则不传
|
||||
with open(local, "r", encoding="utf-8") as f:
|
||||
snapshot = f.readlines()
|
||||
if not any(ln.strip() for ln in snapshot):
|
||||
return ""
|
||||
dest_root = _cfg().get("long_list_backup_path")
|
||||
if not dest_root:
|
||||
bucket, base = _bucket_base()
|
||||
dest_root = f"cos://{bucket}/{base}/long_video"
|
||||
bkt, base_key = parse_cos_url(dest_root)
|
||||
stem = os.path.splitext(os.path.basename(local))[0] # 超长待重跑
|
||||
ts = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
key = f"{base_key.strip('/')}/todo/{stem}_{machine_id}_{ts}.txt"
|
||||
client().put_object(Bucket=bkt, Body="".join(snapshot).encode("utf-8"), Key=key)
|
||||
# 上传成功后清空已上传部分: 只删前 len(snapshot) 行, 保留上传期间新追加的(append 在末尾),
|
||||
# 这样每天只传增量、跑完即了, 不会和已完成分片重叠。开关 long_list_clear_after_upload(默认开)。
|
||||
if _cfg().get("long_list_clear_after_upload", True):
|
||||
try:
|
||||
with open(local, "r", encoding="utf-8") as f:
|
||||
cur = f.readlines()
|
||||
with open(local, "w", encoding="utf-8") as f:
|
||||
f.writelines(cur[len(snapshot):])
|
||||
except OSError:
|
||||
pass
|
||||
return f"cos://{bkt}/{key}"
|
||||
|
||||
|
||||
def _daily_backup_bucket_base():
|
||||
"""daily_backup_path 配置解析为 (bucket, base_key)，缺省时与 upload_daily_backup 保持一致。
|
||||
兜底基准用 _primary_bucket_base()(固定挂 jobs_root), 不能用 _bucket_base() —— 见后者docstring
|
||||
里的解释, 否则专职长视频机会把每日备份传到 long_video 队列自己底下, 全局去重读不到。"""
|
||||
dest_root = _cfg().get("daily_backup_path")
|
||||
if not dest_root:
|
||||
bucket, base = _primary_bucket_base()
|
||||
dest_root = f"cos://{bucket}/{base}/daily"
|
||||
bucket, base = parse_cos_url(dest_root)
|
||||
return bucket, base.strip("/")
|
||||
|
||||
|
||||
def _extract_processed_paths(jsonl_path: str) -> list:
|
||||
"""从本地 数据.jsonl 逐行解析出 video_cos_path 列表(顺序不去重, 上传前用)。"""
|
||||
paths = []
|
||||
try:
|
||||
with open(jsonl_path, "r", encoding="utf-8") as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
vid = json.loads(line).get("video_cos_path")
|
||||
if vid:
|
||||
paths.append(vid)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
except OSError:
|
||||
pass
|
||||
return paths
|
||||
|
||||
|
||||
def upload_processed_list(jsonl_path: str, machine_id: str) -> str:
|
||||
"""从本地 数据.jsonl 提取 video_cos_path, 生成一份纯文本清单(一行一个路径, 不含状态/耗时等
|
||||
其它字段)上传到 daily_backup_path/<日期>/已处理_<机器号>.txt —— 跟 upload_daily_backup 上传的
|
||||
完整 jsonl 是同一份数据的两种形态, 前者给 global_processed_set() 用: 只关心"这个视频是否
|
||||
处理过"时, 拉这个纯清单比拉完整jsonl再挨行解析JSON 小得多、快得多(jsonl里每条记录还带
|
||||
timestamp/duration/image_counts/stage_durations 等一堆跟去重无关的字段)。"""
|
||||
dest_root = _cfg().get("daily_backup_path")
|
||||
if not dest_root:
|
||||
bucket, base = _bucket_base()
|
||||
dest_root = f"cos://{bucket}/{base}/daily"
|
||||
paths = _extract_processed_paths(jsonl_path)
|
||||
if not paths:
|
||||
return ""
|
||||
bkt, base_key = parse_cos_url(dest_root)
|
||||
date = datetime.now().strftime("%Y%m%d")
|
||||
key = f"{base_key.strip('/')}/{date}/已处理_{machine_id}.txt"
|
||||
client().put_object(Bucket=bkt, Body=("\n".join(paths) + "\n").encode("utf-8"), Key=key)
|
||||
return f"cos://{bkt}/{key}"
|
||||
|
||||
|
||||
def list_processed_list_files():
|
||||
"""列出 daily_backup_path 下所有 <日期>/已处理_<机器号>.txt，返回 [(日期, 机器号, key), ...]。"""
|
||||
bkt, base = _daily_backup_bucket_base()
|
||||
prefix = base + "/"
|
||||
out, marker = [], ""
|
||||
while True:
|
||||
resp = client().list_objects(Bucket=bkt, Prefix=prefix, Marker=marker, MaxKeys=1000)
|
||||
for c in resp.get("Contents", []) or []:
|
||||
rel = c["Key"][len(prefix):] # <日期>/已处理_<机器号>.txt
|
||||
parts = rel.split("/")
|
||||
if len(parts) == 2 and parts[1].startswith("已处理_") and parts[1].endswith(".txt"):
|
||||
out.append((parts[0], parts[1][len("已处理_"):-len(".txt")], c["Key"]))
|
||||
if resp.get("IsTruncated") in ("true", True):
|
||||
marker = resp.get("NextMarker") or (resp["Contents"][-1]["Key"] if resp.get("Contents") else "")
|
||||
if not marker:
|
||||
break
|
||||
else:
|
||||
break
|
||||
return out
|
||||
|
||||
|
||||
def global_processed_set(verbose: bool = False, all_dates: bool = False) -> set:
|
||||
"""聚合"已处理清单"(已处理_<机器号>.txt, 由 upload_processed_list 在每次备份时跟 数据.jsonl
|
||||
一起生成), 合并成全局已处理 video_cos_path 集合。只拉轻量txt清单、不下载/解析完整jsonl,
|
||||
省流量也省CPU。语义与本机 load_processed() 一致：不分 success/no_subtitle/error，只要出现过
|
||||
就算处理过。用于灌任务前过滤，避免同一视频被 30 台机器重复处理。
|
||||
|
||||
all_dates=False(默认, 日常灌任务用): 每台机器只取【最新一天】——数据.jsonl 是全量快照(本地
|
||||
从不清空), 取最新一份即覆盖该机器全部历史, 省得把每天的备份都下载一遍; 不同机器最新日期
|
||||
可能不同(比如有机器当天没开机备份), 各自独立取最新, 不要求同一天。
|
||||
|
||||
all_dates=True(机群已停机/一次性核对时用): 不做"每机器只取最新"这个假设, 把【所有日期】
|
||||
的清单全部下载合并——机群都停了就没有"之后还会补传更晚日期"这回事, 更保险的做法是把能拿到
|
||||
的所有历史清单全并起来判断"处理过没有", 而不是赌某一天的快照就是最全的(万一那台机器最新
|
||||
一天的快照因为某种原因不完整/被截断过, 只取最新反而会漏)。代价是下载的文件数从"机器数"
|
||||
变成"机器数×历史天数", 视频清单量小(轻量txt)时这个代价可以接受。"""
|
||||
files = list_processed_list_files()
|
||||
if not files:
|
||||
return set()
|
||||
|
||||
if all_dates:
|
||||
targets = [(date_str, machine_id, key) for date_str, machine_id, key in files]
|
||||
else:
|
||||
latest = {} # machine_id -> (日期, key)
|
||||
for date_str, machine_id, key in files:
|
||||
if machine_id not in latest or date_str > latest[machine_id][0]:
|
||||
latest[machine_id] = (date_str, key)
|
||||
targets = [(date_str, machine_id, key) for machine_id, (date_str, key) in latest.items()]
|
||||
|
||||
bkt, _ = _daily_backup_bucket_base()
|
||||
processed = set()
|
||||
for date_str, machine_id, key in targets:
|
||||
try:
|
||||
data = client().get_object(Bucket=bkt, Key=key)["Body"].get_raw_stream().read()
|
||||
except Exception as e:
|
||||
logging.warning(f"[全局去重] 读取 {key} 失败, 跳过: {e}")
|
||||
continue
|
||||
n = 0
|
||||
for line in data.decode("utf-8", errors="ignore").splitlines():
|
||||
line = line.strip()
|
||||
if line:
|
||||
processed.add(line)
|
||||
n += 1
|
||||
if verbose:
|
||||
logging.info(f"[全局去重] 机器 {machine_id}({date_str}): {n} 条")
|
||||
return processed
|
||||
|
||||
|
||||
def run_daily_backup_loop(machine_id: str, hour: int, stop, jsonl_path: str = "数据.jsonl",
|
||||
check_interval: int = 120):
|
||||
"""后台线程: 每天到 hour 点把 数据.jsonl + 超长列表上传 COS, 每天只传一次。
|
||||
main.py(本地模式) 和 run_worker.py(COS模式) 通用。stop 为 threading.Event, set() 即退出。"""
|
||||
last_date = None
|
||||
while not stop.wait(check_interval): # 每 check_interval 秒检查一次
|
||||
now = datetime.now()
|
||||
today = now.strftime("%Y%m%d")
|
||||
if now.hour >= hour and last_date != today:
|
||||
try:
|
||||
dst = upload_daily_backup(jsonl_path, machine_id)
|
||||
if dst:
|
||||
logging.info(f"[每日备份] 数据.jsonl 已上传 -> {dst}")
|
||||
else:
|
||||
logging.info("[每日备份] 本地无 数据.jsonl, 跳过")
|
||||
ll = upload_longlist_backup(machine_id)
|
||||
if ll:
|
||||
logging.info(f"[每日备份] 超长待重跑列表已上传 -> {ll}")
|
||||
last_date = today # 当天标记完成, 避免反复触发
|
||||
except Exception as e:
|
||||
logging.error(f"[每日备份] 上传失败(稍后重试): {e}")
|
||||
|
||||
|
||||
# ----------------------------------------------------- 每小时监控摘要 ----
|
||||
|
||||
def _monitor_bucket_base():
|
||||
"""monitor_root 配置解析为 (bucket, base_key)，缺省时 = jobs_root 的同级 monitor 目录
|
||||
(跟 _daily_backup_bucket_base 同一个派生模式)。兜底基准同样必须用 _primary_bucket_base()
|
||||
而不是 _bucket_base() —— 这就是本函数曾经的实际bug: 专职长视频机(run_long_video=true)
|
||||
没显式配 monitor_root 时, 旧版用 _bucket_base() 兜底会解析到【当前生效队列根】
|
||||
(_active_jobs_root(), 长视频机上长期等于 long_list_backup_path), 算出
|
||||
.../long_video/monitor/... 这种独立于其它机器的路径, fleet_monitor.py 只读标准
|
||||
monitor_root, 找不到这些机器的小时摘要, 表现为"全队机器数对不上"(实际数据都传了,
|
||||
只是传错地方, 数据没丢, 改这个函数后新上传的会去对的地方, 旧的仍留在错的路径下)。"""
|
||||
dest_root = _cfg().get("monitor_root")
|
||||
if not dest_root:
|
||||
bucket, base = _primary_bucket_base()
|
||||
dest_root = f"cos://{bucket}/{base}/monitor"
|
||||
bucket, base = parse_cos_url(dest_root)
|
||||
return bucket, base.strip("/")
|
||||
|
||||
|
||||
def compute_hourly_digest(jsonl_path: str, hour_dt: datetime):
|
||||
"""从本地 数据.jsonl 里筛出 process_start(缺失回退timestamp) 落在 hour_dt 所在整点小时的记录,
|
||||
算出这个小时的摘要。hour_dt 的分钟/秒会被忽略, 只取年月日时。
|
||||
本地jsonl不存在/读取失败 -> 返回 None(调用方应跳过上传)。
|
||||
这一小时确实一条记录都没处理(count==0, 比如COS任务队列暂时被抢空、worker在idle轮询)
|
||||
-> 仍然返回一份全0的摘要, 不当 None 处理 —— 队列见底idle是这套分布式调度系统的正常状态
|
||||
(见 run_worker.py 的 keep_alive/idle_poll_seconds), 如果这种情况也不传摘要, fleet_monitor.py
|
||||
的"机器掉线"检测会把"确认活着只是没活干"误判成"可能停机", 制造假警报。"""
|
||||
if not os.path.exists(jsonl_path):
|
||||
return None
|
||||
hour_prefix = hour_dt.strftime("%Y-%m-%dT%H") # 匹配 ISO 格式 process_start 的前13位
|
||||
count = success = no_subtitle = error = 0
|
||||
success_hours = 0.0
|
||||
try:
|
||||
with open(jsonl_path, "r", encoding="utf-8") as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
rec = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
ts = rec.get("process_start") or rec.get("timestamp") or ""
|
||||
if ts[:13] != hour_prefix:
|
||||
continue
|
||||
count += 1
|
||||
status = rec.get("status")
|
||||
if status == "success":
|
||||
success += 1
|
||||
try:
|
||||
success_hours += float(rec.get("duration") or 0) / 3600.0
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
elif status == "no_subtitle":
|
||||
no_subtitle += 1
|
||||
else:
|
||||
error += 1
|
||||
except OSError:
|
||||
return None
|
||||
return {
|
||||
"machine_id": None, # upload_hourly_digest 里填
|
||||
"hour": hour_dt.strftime("%Y-%m-%dT%H"),
|
||||
"count": count, "success": success, "no_subtitle": no_subtitle, "error": error,
|
||||
"success_hours": round(success_hours, 2),
|
||||
"uploaded_at": datetime.now().isoformat(),
|
||||
}
|
||||
|
||||
|
||||
def upload_hourly_digest(machine_id: str, jsonl_path: str, hour_dt: datetime) -> str:
|
||||
"""算好 hour_dt 那个整点小时的摘要并上传到
|
||||
monitor_root/hourly/<日期YYYYMMDD>/<小时HH>/<机器号>.json。
|
||||
即使这一小时count==0(队列见底idle)也会正常上传一份全0摘要——只有本地 数据.jsonl
|
||||
根本不存在/读取失败(digest为None, 比如worker刚起来还没写过jsonl)才跳过, 返回''。"""
|
||||
digest = compute_hourly_digest(jsonl_path, hour_dt)
|
||||
if digest is None:
|
||||
return ""
|
||||
digest["machine_id"] = machine_id
|
||||
bkt, base = _monitor_bucket_base()
|
||||
date_str = hour_dt.strftime("%Y%m%d")
|
||||
hour_str = hour_dt.strftime("%H")
|
||||
key = f"{base}/hourly/{date_str}/{hour_str}/{machine_id}.json"
|
||||
client().put_object(Bucket=bkt, Body=json.dumps(digest, ensure_ascii=False).encode("utf-8"), Key=key)
|
||||
return f"cos://{bkt}/{key}"
|
||||
|
||||
|
||||
def list_hourly_digest_files_for_hour(hour_dt: datetime) -> list:
|
||||
"""列出 monitor_root/hourly/<日期>/<小时>/ 下所有机器的摘要文件, 返回 [(机器号, key), ...]。"""
|
||||
bkt, base = _monitor_bucket_base()
|
||||
date_str = hour_dt.strftime("%Y%m%d")
|
||||
hour_str = hour_dt.strftime("%H")
|
||||
prefix = f"{base}/hourly/{date_str}/{hour_str}/"
|
||||
out, marker = [], ""
|
||||
while True:
|
||||
resp = client().list_objects(Bucket=bkt, Prefix=prefix, Marker=marker, MaxKeys=1000)
|
||||
for c in resp.get("Contents", []) or []:
|
||||
rel = c["Key"][len(prefix):]
|
||||
if rel.endswith(".json"):
|
||||
out.append((rel[:-len(".json")], c["Key"]))
|
||||
if resp.get("IsTruncated") in ("true", True):
|
||||
marker = resp.get("NextMarker") or (resp["Contents"][-1]["Key"] if resp.get("Contents") else "")
|
||||
if not marker:
|
||||
break
|
||||
else:
|
||||
break
|
||||
return out
|
||||
|
||||
|
||||
def download_hourly_digest(key: str) -> dict:
|
||||
"""下载解析单份小时摘要json, 返回dict。失败抛异常, 调用方负责捕获。"""
|
||||
bkt, _ = _monitor_bucket_base()
|
||||
data = client().get_object(Bucket=bkt, Key=key)["Body"].get_raw_stream().read()
|
||||
return json.loads(data.decode("utf-8"))
|
||||
|
||||
|
||||
def upload_monitor_report(local_path: str, hour_dt: datetime) -> str:
|
||||
"""把生成好的监控报告文件上传到 monitor_root/reports/<日期YYYYMMDD>_<小时HH>.md。"""
|
||||
if not os.path.exists(local_path):
|
||||
return ""
|
||||
bkt, base = _monitor_bucket_base()
|
||||
date_str = hour_dt.strftime("%Y%m%d")
|
||||
hour_str = hour_dt.strftime("%H")
|
||||
key = f"{base}/reports/{date_str}_{hour_str}.md"
|
||||
with open(local_path, "rb") as f:
|
||||
data = f.read()
|
||||
client().put_object(Bucket=bkt, Body=data, Key=key)
|
||||
return f"cos://{bkt}/{key}"
|
||||
|
||||
|
||||
def run_hourly_digest_loop(machine_id: str, stop, jsonl_path: str = "数据.jsonl",
|
||||
check_interval: int = 120):
|
||||
"""后台线程: 每小时把刚结束的那个整点小时的处理摘要上传COS, 每小时只传一次。
|
||||
main.py(本地模式) 不需要接入(一次性批处理没有"常驻小时"的概念, 见设计文档)。
|
||||
stop 为 threading.Event, set() 即退出。"""
|
||||
last_uploaded_hour = None
|
||||
while not stop.wait(check_interval):
|
||||
now = datetime.now()
|
||||
prev_hour_dt = now.replace(minute=0, second=0, microsecond=0) - timedelta(hours=1)
|
||||
prev_hour_key = prev_hour_dt.strftime("%Y%m%d%H")
|
||||
if last_uploaded_hour == prev_hour_key:
|
||||
continue
|
||||
try:
|
||||
dst = upload_hourly_digest(machine_id, jsonl_path, prev_hour_dt)
|
||||
if dst:
|
||||
logging.info(f"[小时摘要] {prev_hour_key} 已上传 -> {dst}")
|
||||
else:
|
||||
logging.info(f"[小时摘要] {prev_hour_key} 本地 数据.jsonl 尚不存在, 跳过上传")
|
||||
last_uploaded_hour = prev_hour_key
|
||||
except Exception as e:
|
||||
logging.error(f"[小时摘要] {prev_hour_key} 上传失败(下次继续重试): {e}")
|
||||
0
tool/log/__init__.py
Normal file
0
tool/log/__init__.py
Normal file
218
tool/log/log_manager.py
Normal file
218
tool/log/log_manager.py
Normal file
@@ -0,0 +1,218 @@
|
||||
# log_manager.py
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import threading
|
||||
from collections import deque
|
||||
from datetime import datetime
|
||||
from typing import Dict, List, Optional, Union
|
||||
|
||||
|
||||
class ProgramFilter(logging.Filter):
|
||||
"""自定义日志过滤器，为日志记录添加程序名称属性"""
|
||||
|
||||
def filter(self, record: logging.LogRecord) -> bool:
|
||||
if not hasattr(record, 'program'):
|
||||
record.program = "UNKNOWN"
|
||||
return True
|
||||
|
||||
|
||||
class UnifiedLogger:
|
||||
"""统一日志管理器（线程安全单例），支持多程序日志收集"""
|
||||
|
||||
_instance: Optional['UnifiedLogger'] = None
|
||||
_lock = threading.Lock()
|
||||
|
||||
def __new__(cls) -> 'UnifiedLogger':
|
||||
with cls._lock:
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
return cls._instance
|
||||
|
||||
def __init__(self) -> None:
|
||||
# 避免重复初始化
|
||||
if hasattr(self, '_initialized') and self._initialized:
|
||||
return
|
||||
|
||||
# 每视频状态使用线程本地存储：批处理多线程并发时，每个视频在各自线程内
|
||||
# 独立记录日志/计时/统计，互不串扰。（嵌套子线程如OCR并发worker不应直接
|
||||
# 调用本logger的状态方法——它们会落到worker线程的独立状态而丢失。）
|
||||
self._local = threading.local()
|
||||
|
||||
# 是否打印详细日志到控制台。默认 False：控制台只显示"视频→步骤"里程碑，
|
||||
# 详细日志仍完整存进每视频 JSON(log/<video>.json)。设 LOG_VERBOSE=1 恢复全部。
|
||||
self.verbose = os.environ.get("LOG_VERBOSE", "0") == "1"
|
||||
|
||||
# 自定义 logger（避免污染根日志器）
|
||||
self.logger = logging.getLogger("UnifiedLogger")
|
||||
self.logger.setLevel(logging.DEBUG)
|
||||
self.logger.propagate = False # 不传播到 root，避免每条日志重复打印
|
||||
self.logger.addFilter(ProgramFilter())
|
||||
|
||||
# 控制台 Handler：非 verbose 时只放行 WARNING+，把 INFO 细节挡在控制台外(仍进 JSON)
|
||||
console_handler = logging.StreamHandler()
|
||||
console_handler.setLevel(logging.DEBUG if self.verbose else logging.WARNING)
|
||||
console_formatter = logging.Formatter(
|
||||
'%(asctime)s [%(levelname)s] [%(program)s] %(message)s'
|
||||
)
|
||||
console_handler.setFormatter(console_formatter)
|
||||
self.logger.addHandler(console_handler)
|
||||
|
||||
# 第三方库降噪：jieba 的 DEBUG("Building prefix dict...")太吵
|
||||
logging.getLogger("jieba").setLevel(logging.WARNING)
|
||||
# qcloud_cos SDK 噪音大: 每次 get/put/head object 都打一条 INFO(带完整url/headers),
|
||||
# 而且它把"对象不存在"(HEAD 404, 本项目里大量用来判断分片是否已完成/锁是否存在,
|
||||
# 属于正常控制流)也当 WARNING 打一条 —— 两种都没信息量, 真正的失败走异常
|
||||
# (CosServiceError), 本项目自己的 try/except 会记录, 不受这里调级影响。
|
||||
logging.getLogger("qcloud_cos").setLevel(logging.ERROR)
|
||||
|
||||
self._initialized = True
|
||||
|
||||
def step(self, video: str, step: str, extra: str = "") -> None:
|
||||
"""打印一条简洁的"视频 → 步骤"里程碑（始终显示，不受 verbose 影响）。"""
|
||||
line = f"[{datetime.now().strftime('%H:%M:%S')}] {video} → {step}"
|
||||
if extra:
|
||||
line += f" ({extra})"
|
||||
print(line, flush=True)
|
||||
try:
|
||||
self._state()["log_buffer"].append({
|
||||
"timestamp": datetime.now().isoformat(), "level": "STEP",
|
||||
"program": "step", "message": f"{step} {extra}".strip(),
|
||||
})
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _state(self) -> dict:
|
||||
"""获取当前线程的日志状态，惰性初始化。"""
|
||||
st = getattr(self._local, 'state', None)
|
||||
if st is None:
|
||||
st = {
|
||||
"log_buffer": deque(maxlen=10000),
|
||||
"log_file": "unified_log.json",
|
||||
"program_sequence": [],
|
||||
"program_start_times": {},
|
||||
"program_end_times": {},
|
||||
"custom_stats": {},
|
||||
}
|
||||
self._local.state = st
|
||||
return st
|
||||
|
||||
def init_logging(self, log_file: str = "unified_log.json") -> None:
|
||||
"""初始化日志系统（可重新指定日志文件路径）"""
|
||||
self._state()["log_file"] = log_file
|
||||
log_dir = os.path.dirname(log_file)
|
||||
if log_dir:
|
||||
os.makedirs(log_dir, exist_ok=True)
|
||||
|
||||
def reset(self) -> None:
|
||||
"""重置当前线程的日志状态（用于多线程场景下每个视频独立记录）"""
|
||||
st = self._state()
|
||||
st["log_buffer"].clear()
|
||||
st["program_sequence"].clear()
|
||||
st["program_start_times"].clear()
|
||||
st["program_end_times"].clear()
|
||||
st["custom_stats"].clear()
|
||||
|
||||
def set_stat(self, key: str, value: object) -> None:
|
||||
"""记录一项自定义统计，会写入日志 metadata.custom_stats"""
|
||||
self._state()["custom_stats"][key] = value
|
||||
|
||||
def _get_level_int(self, level: Union[str, int]) -> int:
|
||||
"""将日志级别转换为 logging 模块的整数常量"""
|
||||
if isinstance(level, int):
|
||||
return level
|
||||
level_map = {
|
||||
"DEBUG": logging.DEBUG,
|
||||
"INFO": logging.INFO,
|
||||
"WARNING": logging.WARNING,
|
||||
"ERROR": logging.ERROR,
|
||||
"CRITICAL": logging.CRITICAL
|
||||
}
|
||||
return level_map.get(level.upper(), logging.INFO)
|
||||
|
||||
def log(self, message: str, level: Union[str, int] = "info", program: str = "A") -> None:
|
||||
"""记录一条日志"""
|
||||
timestamp = datetime.now().isoformat()
|
||||
level_str = level if isinstance(level, str) else logging.getLevelName(level)
|
||||
log_entry = {
|
||||
"timestamp": timestamp,
|
||||
"level": level_str,
|
||||
"program": program,
|
||||
"message": message
|
||||
}
|
||||
|
||||
self._state()["log_buffer"].append(log_entry)
|
||||
|
||||
# 输出到自定义 logger
|
||||
level_int = self._get_level_int(level)
|
||||
self.logger.log(level_int, message, extra={"program": program})
|
||||
|
||||
def start_program(self, program_name: str) -> None:
|
||||
"""记录程序开始执行"""
|
||||
st = self._state()
|
||||
if program_name in st["program_start_times"]:
|
||||
self.logger.warning(f"程序 {program_name} 已开始，忽略重复 start 调用", extra={"program": "SYSTEM"})
|
||||
return
|
||||
|
||||
st["program_start_times"][program_name] = time.time()
|
||||
st["program_sequence"].append(program_name)
|
||||
|
||||
self.log(f"程序 {program_name} 开始执行", program=program_name)
|
||||
|
||||
def end_program(self, program_name: str) -> None:
|
||||
"""记录程序执行结束"""
|
||||
st = self._state()
|
||||
if program_name not in st["program_start_times"]:
|
||||
self.logger.error(f"程序 {program_name} 未调用 start_program，无法记录结束", extra={"program": "SYSTEM"})
|
||||
return
|
||||
|
||||
end_time = time.time()
|
||||
st["program_end_times"][program_name] = end_time
|
||||
duration = end_time - st["program_start_times"][program_name]
|
||||
|
||||
self.log(f"程序 {program_name} 执行完成 (耗时: {duration:.2f}s)", program=program_name)
|
||||
|
||||
def save_logs(self) -> None:
|
||||
"""将当前线程的所有日志（含元数据）保存到 JSON 文件"""
|
||||
st = self._state()
|
||||
logs_copy = list(st["log_buffer"])
|
||||
start_times_copy = dict(st["program_start_times"])
|
||||
end_times_copy = dict(st["program_end_times"])
|
||||
sequence_copy = list(st["program_sequence"])
|
||||
stats_copy = dict(st["custom_stats"])
|
||||
log_file = st["log_file"]
|
||||
|
||||
# 计算程序执行耗时
|
||||
durations = {}
|
||||
for prog, start in start_times_copy.items():
|
||||
end = end_times_copy.get(prog, time.time())
|
||||
durations[prog] = end - start
|
||||
|
||||
metadata = {
|
||||
"log_created": datetime.now().isoformat(),
|
||||
"program_sequence": sequence_copy,
|
||||
"total_entries": len(logs_copy),
|
||||
"program_durations": durations,
|
||||
"custom_stats": stats_copy
|
||||
}
|
||||
|
||||
full_log = {
|
||||
"metadata": metadata,
|
||||
"logs": logs_copy
|
||||
}
|
||||
|
||||
try:
|
||||
with open(log_file, 'w', encoding='utf-8') as f:
|
||||
json.dump(full_log, f, ensure_ascii=False, indent=2)
|
||||
self.logger.info(f"所有日志已保存到: {log_file}", extra={"program": "SYSTEM"})
|
||||
except (OSError, IOError) as e:
|
||||
self.logger.error(f"保存日志失败: {e}", extra={"program": "SYSTEM"})
|
||||
|
||||
def get_logs(self) -> List[dict]:
|
||||
"""获取当前线程缓冲区中的所有日志副本"""
|
||||
return list(self._state()["log_buffer"])
|
||||
|
||||
|
||||
# 全局单例实例
|
||||
logger = UnifiedLogger()
|
||||
885
tool/merge.py
Normal file
885
tool/merge.py
Normal file
@@ -0,0 +1,885 @@
|
||||
import os
|
||||
import json
|
||||
import cv2
|
||||
import numpy as np
|
||||
import random
|
||||
import re
|
||||
import string
|
||||
import math
|
||||
import tempfile
|
||||
from collections import defaultdict
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from difflib import SequenceMatcher
|
||||
from PIL import Image
|
||||
from typing import Optional, List, Tuple, Dict, Any, Union
|
||||
# 假设这些外部模块已正确导入
|
||||
from tool.Ocr.tencentcloud_ocr import tencentcloud_ocr_api, img2base64
|
||||
from tool.whisper_manager import whisper_manager
|
||||
from tool.log.log_manager import logger
|
||||
from config.config_utils import config
|
||||
|
||||
|
||||
# ==================== 辅助函数 ====================
|
||||
def natural_sort_key(filename: str) -> List[Union[int, str]]:
|
||||
"""
|
||||
自然排序键函数，将文件名中的数字部分转换为整数进行排序。
|
||||
"""
|
||||
return [int(part) if part.isdigit() else part.lower()
|
||||
for part in re.split(r'(\d+)', filename)]
|
||||
|
||||
|
||||
def has_chinese(text: str) -> bool:
|
||||
"""文本是否含中文字符（CJK 统一表意文字）。"""
|
||||
return any('一' <= ch <= '鿿' for ch in (text or ""))
|
||||
|
||||
|
||||
def calculate_english_ratio(text: str) -> float:
|
||||
"""计算文本中英文字符（字母、数字、常见标点）的占比。"""
|
||||
if not text:
|
||||
return 0.0
|
||||
total_chars = len(text)
|
||||
if total_chars == 0:
|
||||
return 0.0
|
||||
|
||||
english_chars = 0
|
||||
for ch in text:
|
||||
if ch.isalpha() and ch.isascii():
|
||||
english_chars += 1
|
||||
elif ch.isdigit():
|
||||
english_chars += 1
|
||||
elif ch in ' .,!?;:\'"-()[]{}':
|
||||
english_chars += 1
|
||||
return english_chars / total_chars
|
||||
|
||||
|
||||
def check_subtitle_quality_enhanced(best_subtitle: str, strict_mode: bool = True) -> bool:
|
||||
"""
|
||||
增强版字幕质量检测。
|
||||
|
||||
:param best_subtitle: 待检测的文本
|
||||
:param strict_mode: 是否使用严格模式
|
||||
:return: 若文本质量合格则返回 True
|
||||
"""
|
||||
if not best_subtitle or not best_subtitle.strip():
|
||||
return True
|
||||
|
||||
subtitle = best_subtitle.strip()
|
||||
total_chars = len(subtitle)
|
||||
if total_chars == 0:
|
||||
return True
|
||||
|
||||
digit_count = sum(1 for ch in subtitle if ch.isdigit())
|
||||
symbol_count = sum(1 for ch in subtitle if ch in string.punctuation)
|
||||
special_ratio = (digit_count + symbol_count) / total_chars
|
||||
english_ratio = calculate_english_ratio(subtitle)
|
||||
|
||||
if strict_mode:
|
||||
digit_ratio = digit_count / total_chars
|
||||
if total_chars >= 10 and digit_ratio <= 0.9 and special_ratio <= 0.9:
|
||||
return True
|
||||
# 数字占比 ≤ 70% 且特殊字符占比 ≤ 80% 且英文占比 ≤ 90%
|
||||
return digit_ratio <= 0.7 and special_ratio <= 0.8 and english_ratio <= 0.9
|
||||
else:
|
||||
# 普通模式仅检查特殊字符占比
|
||||
return special_ratio <= 0.5
|
||||
|
||||
|
||||
def detect_main_region(image_files: List[str],
|
||||
input_dir: str,
|
||||
sample_size: int = 10,
|
||||
motion_threshold: int = 30) -> Tuple[int, int, int, int]:
|
||||
"""
|
||||
通过运动检测确定主画面区域。
|
||||
|
||||
:return: 边界框 (x, y, width, height)
|
||||
"""
|
||||
if not image_files:
|
||||
raise ValueError("图片列表为空")
|
||||
|
||||
first_img_path = os.path.join(input_dir, image_files[0])
|
||||
prev_frame = cv2.imread(first_img_path)
|
||||
if prev_frame is None:
|
||||
raise ValueError(f"无法读取第一帧: {first_img_path}")
|
||||
prev_gray = cv2.cvtColor(prev_frame, cv2.COLOR_BGR2GRAY)
|
||||
height, width = prev_gray.shape
|
||||
|
||||
# 随机采样中间帧（避开开头和结尾）
|
||||
mid_start = len(image_files) // 4
|
||||
mid_end = 3 * len(image_files) // 4
|
||||
sample_indices = random.sample(range(mid_start, mid_end),
|
||||
min(sample_size, mid_end - mid_start))
|
||||
|
||||
motion_accumulator = np.zeros((height, width), dtype=np.float32)
|
||||
|
||||
for idx in sample_indices:
|
||||
img_path = os.path.join(input_dir, image_files[idx])
|
||||
curr_frame = cv2.imread(img_path)
|
||||
if curr_frame is None:
|
||||
continue
|
||||
curr_gray = cv2.cvtColor(curr_frame, cv2.COLOR_BGR2GRAY)
|
||||
frame_diff = cv2.absdiff(curr_gray, prev_gray)
|
||||
_, thresh = cv2.threshold(frame_diff, motion_threshold, 255, cv2.THRESH_BINARY)
|
||||
motion_accumulator += thresh.astype(np.float32)
|
||||
prev_gray = curr_gray
|
||||
|
||||
motion_accumulator = cv2.normalize(motion_accumulator, None, 0, 255, cv2.NORM_MINMAX)
|
||||
_, motion_mask = cv2.threshold(motion_accumulator.astype(np.uint8), 100, 255, cv2.THRESH_BINARY)
|
||||
|
||||
contours, _ = cv2.findContours(motion_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||||
if not contours:
|
||||
# 无运动区域时返回中心偏下区域
|
||||
return (0, height // 4, width, height // 2)
|
||||
|
||||
main_contour = max(contours, key=cv2.contourArea)
|
||||
x, y, w, h = cv2.boundingRect(main_contour)
|
||||
|
||||
# 合理性检查
|
||||
if w < width * 0.1 or h < height * 0.1:
|
||||
return (0, height // 4, width, height // 2)
|
||||
|
||||
# 考虑下方空间
|
||||
bottom_space = height - (y + h)
|
||||
if bottom_space > height * 0.1:
|
||||
extend = int(0.8 * bottom_space)
|
||||
h += extend
|
||||
# 调整 y 使区域居中于原运动区域下方
|
||||
y = y + h // 2
|
||||
h //= 2
|
||||
else:
|
||||
h += bottom_space
|
||||
|
||||
return (x, y, w, h)
|
||||
|
||||
|
||||
# ==================== 字幕区域检测 ====================
|
||||
def _sample_indices_uniform(total: int, sample_size: int) -> List[int]:
|
||||
"""均匀采样索引，确保帧间有间隔。"""
|
||||
print(total, sample_size)
|
||||
if sample_size <= 1:
|
||||
return [total // 2]
|
||||
step = total // (sample_size - 1)
|
||||
indices = [i * step for i in range(sample_size)]
|
||||
# 确保不超出范围
|
||||
indices = [min(i, total - 1) for i in indices]
|
||||
return indices
|
||||
|
||||
|
||||
def _ocr_frames_concurrent(image_files: List[str], input_dir: str,
|
||||
indices: List[int], workers: int) -> Dict[int, str]:
|
||||
"""并发对指定帧索引做 OCR，返回 {idx: 原始响应字符串}。
|
||||
在视频主线程内开子线程池，子线程只做网络请求、不写日志（满足日志线程隔离要求）。"""
|
||||
results: Dict[int, str] = {}
|
||||
if not indices:
|
||||
return results
|
||||
|
||||
def _one(idx):
|
||||
try:
|
||||
img_path = os.path.join(input_dir, image_files[idx])
|
||||
return idx, tencentcloud_ocr_api(img2base64(img_path=img_path))
|
||||
except Exception:
|
||||
return idx, None
|
||||
|
||||
with ThreadPoolExecutor(max_workers=max(1, min(workers, len(indices)))) as ex:
|
||||
for idx, res in ex.map(_one, indices):
|
||||
results[idx] = res
|
||||
return results
|
||||
|
||||
|
||||
def cv_locate_subtitle(image_files: List[str], input_dir: str,
|
||||
height: int, width: int, sample: int = 80):
|
||||
"""纯 CV(零 OCR)定位硬字幕带。读已抽好的 JPEG 帧(cv2.imread,与 AV1 解码无关)。
|
||||
返回 (verdict, region, info):
|
||||
'subtitle' (x,y,w,h) —— 找到下部动态字幕带
|
||||
'none' None —— 整体静止+无动态文字带 = PPT/文档/截图,可直接跳过
|
||||
'uncertain'None —— 拿不准,回退原 OCR 检测(不改变原行为)
|
||||
安全准则: 只有【整体静止且无任何动态文字带】才判 none,绝不误杀真字幕视频。"""
|
||||
n = len(image_files)
|
||||
if n < 10:
|
||||
return "uncertain", None, {"reason": "帧太少"}
|
||||
idxs = np.linspace(n * 0.1, n * 0.9, min(sample, n)).astype(int)
|
||||
grays = []
|
||||
for i in idxs:
|
||||
g = cv2.imread(os.path.join(input_dir, image_files[int(i)]), cv2.IMREAD_GRAYSCALE)
|
||||
if g is not None:
|
||||
grays.append(g)
|
||||
if len(grays) < 8:
|
||||
return "uncertain", None, {"reason": "读帧失败"}
|
||||
H, W = grays[0].shape
|
||||
x0, x1 = int(0.15 * W), int(0.85 * W) # 只看横向中间70%,避开角落logo
|
||||
|
||||
edge_acc = np.zeros(H); present = np.zeros(H)
|
||||
for g in grays:
|
||||
e = cv2.Canny(g, 80, 200)
|
||||
row = e[:, x0:x1].sum(axis=1) / 255.0
|
||||
edge_acc += row
|
||||
present += (row > (x1 - x0) * 0.02)
|
||||
edge_profile = edge_acc / len(grays)
|
||||
presence = present / len(grays)
|
||||
diff_acc = np.zeros(H)
|
||||
for k in range(1, len(grays)):
|
||||
diff_acc += cv2.absdiff(grays[k], grays[k - 1])[:, x0:x1].mean(axis=1)
|
||||
temporal = diff_acc / max(1, len(grays) - 1)
|
||||
|
||||
TEMP_HI = 8.0 # 带变化≥此 = 字幕在翻动(动态)
|
||||
TEMP_LO = 4.0 # 文字带变化<此 = 静止
|
||||
STATIC_MED = 3.0 # 全画面 temporal 中位数<此 = 整体静止(PPT/文档)
|
||||
PRES = 0.4
|
||||
med = float(np.median(temporal))
|
||||
|
||||
cand = (presence > PRES) & (edge_profile > edge_profile.max() * 0.25)
|
||||
rows = np.where(cand)[0]
|
||||
if len(rows) == 0:
|
||||
if med < STATIC_MED:
|
||||
return "none", None, {"reason": "无文字带且画面静止", "med": round(med, 1)}
|
||||
return "uncertain", None, {"reason": "无稳定文字带"}
|
||||
|
||||
bands = []
|
||||
s = p = rows[0]
|
||||
for r in rows[1:]:
|
||||
if r - p > H * 0.03:
|
||||
bands.append((s, p)); s = r
|
||||
p = r
|
||||
bands.append((s, p))
|
||||
MIN_H = int(H * 0.012)
|
||||
bands = [b for b in bands if (b[1] - b[0]) >= MIN_H]
|
||||
|
||||
def bchg(b):
|
||||
return float(temporal[b[0]:b[1] + 1].mean())
|
||||
|
||||
dyn = [b for b in bands if bchg(b) >= TEMP_HI] if bands else []
|
||||
if not dyn:
|
||||
# 无动态字幕带: 整体静止 或 所有文字带都静止 → PPT跳过; 否则交OCR
|
||||
if med < STATIC_MED:
|
||||
return "none", None, {"reason": "无动态字幕带且画面静止(PPT)", "med": round(med, 1)}
|
||||
if bands and max(bchg(b) for b in bands) < TEMP_LO:
|
||||
return "none", None, {"reason": "所有文字带静止", "max_chg": round(max(bchg(b) for b in bands), 1)}
|
||||
return "uncertain", None, {"reason": "拿不准(交OCR)", "med": round(med, 1)}
|
||||
|
||||
dyn.sort(key=lambda b: presence[b[0]:b[1] + 1].mean() * (0.4 + (b[0] + b[1]) / 2 / H), reverse=True)
|
||||
ys, ye = dyn[0]
|
||||
band_h = (ye - ys) / H; cy = (ys + ye) / 2 / H
|
||||
# 若抽帧已裁到下半屏(crop_bottom_ratio),帧内"偏上"不再代表标题,放宽位置判定。
|
||||
crop_r = config.get_splitframe_config().get('crop_bottom_ratio')
|
||||
try:
|
||||
crop_r = float(crop_r) if crop_r is not None else None
|
||||
except (TypeError, ValueError):
|
||||
crop_r = None
|
||||
min_cy = 0.15 if (crop_r and 0.0 < crop_r < 1.0) else 0.45
|
||||
# 有动态带=【不是PPT】。带太宽/偏上只代表不敢直接给区域 → 回退OCR,绝不判none。
|
||||
if band_h > 0.35:
|
||||
return "uncertain", None, {"reason": "动态带太宽,交OCR", "band_h": round(band_h, 2)}
|
||||
if cy < min_cy:
|
||||
return "uncertain", None, {"reason": "带偏上,交OCR", "cy": round(cy, 2)}
|
||||
pad = int(H * 0.012)
|
||||
y = max(0, ys - pad); h = min(H, ye + pad) - y
|
||||
return "subtitle", (0, int(y), int(W), int(h)), {"cy": round(cy, 2), "chg": round(bchg((ys, ye)), 1)}
|
||||
|
||||
|
||||
def detect_subtitle_region(image_files: List[str],
|
||||
input_dir: str,
|
||||
sample_size: int = 15,
|
||||
keep_num: int = 9) -> Tuple[Optional[Tuple[int, int, int, int]],
|
||||
List[int],
|
||||
List[str],
|
||||
List[Dict]]:
|
||||
"""
|
||||
使用 OCR 检测字幕区域。
|
||||
|
||||
:return: (subtitle_region, sample_indices, watermark_logs, candidate_regions)
|
||||
subtitle_region 为 (x, y, width, height)；若无字幕则返回 None
|
||||
"""
|
||||
_sub_cfg = config.get_subtitle_config()
|
||||
require_chinese = _sub_cfg.get('require_chinese', True) # 只要中文字幕
|
||||
bands_thr = int(_sub_cfg.get('redundant_bands', 12)) # 冗余阈值(垂直位置簇数)，越大越宽松
|
||||
|
||||
# 读取第一张图片获取尺寸
|
||||
first_img_path = os.path.join(input_dir, image_files[0])
|
||||
img = cv2.imread(first_img_path)
|
||||
if img is None:
|
||||
raise ValueError(f"无法读取图片: {first_img_path}")
|
||||
height, width = img.shape[:2]
|
||||
|
||||
# === CV 快速闸门(零 OCR,~1.4s)===
|
||||
# 'none' → 静态/PPT,直接跳过(省掉对这类视频的几千次 OCR);
|
||||
# 'subtitle' 且开启 cv_use_region → 直接用 CV 区域(跳过 OCR 检测);
|
||||
# 否则 → 继续走下面的 OCR 检测(行为不变)。
|
||||
if _sub_cfg.get('cv_detect', True):
|
||||
try:
|
||||
cv_verdict, cv_region, cv_info = cv_locate_subtitle(image_files, input_dir, height, width)
|
||||
except Exception as _e:
|
||||
cv_verdict, cv_region, cv_info = "uncertain", None, {"reason": f"CV异常:{_e}"}
|
||||
if cv_verdict == "none":
|
||||
logger.log(f"CV定位: 无动态字幕(静态/PPT),跳过 [{cv_info}]", level="warning", program="merge")
|
||||
return None, [], [], []
|
||||
if cv_verdict == "subtitle" and _sub_cfg.get('cv_use_region', False):
|
||||
logger.log(f"CV定位: 字幕区域 {cv_region} (零OCR) [{cv_info}]", program="merge")
|
||||
x, y, w, h = cv_region
|
||||
cand = [{"x_min": x, "y_min": y, "x_max": x + w, "y_max": y + h, "width": w, "height": h}]
|
||||
return cv_region, [], [], cand
|
||||
logger.log(f"CV定位: {cv_verdict} [{cv_info}] → 走OCR检测", program="merge")
|
||||
|
||||
# 均匀采样帧（扩大到 10%-90% 区间，提高命中字幕帧的概率）
|
||||
mid_start = len(image_files) // 10
|
||||
mid_end = 9 * len(image_files) // 10
|
||||
sample_range = mid_end - mid_start
|
||||
sample_indices = [mid_start + idx for idx in _sample_indices_uniform(sample_range, sample_size)]
|
||||
|
||||
# 存储候选区域和水印信息
|
||||
candidate_regions = [] # 每个元素为 dict: {x_min, y_min, x_max, y_max, width, height}
|
||||
text_positions = defaultdict(list) # {text: [{'x':..., 'y':..., 'idx':...}]}
|
||||
watermark_logs = []
|
||||
|
||||
# 并发预取采样帧的 OCR（原为逐帧串行，~15帧×2s 全在等网络）
|
||||
ocr_workers = config.get_subtitle_config().get('ocr_workers', 10)
|
||||
ocr_cache = _ocr_frames_concurrent(image_files, input_dir, sample_indices, ocr_workers)
|
||||
|
||||
for idx in sample_indices:
|
||||
ocr_result = ocr_cache.get(idx)
|
||||
if not ocr_result:
|
||||
continue
|
||||
ocr_data = json.loads(ocr_result)
|
||||
if "TextDetections" not in ocr_data:
|
||||
continue
|
||||
|
||||
for detection in ocr_data["TextDetections"]:
|
||||
if "Polygon" not in detection:
|
||||
continue
|
||||
polygon = detection["Polygon"]
|
||||
xs = [point["X"] for point in polygon]
|
||||
ys = [point["Y"] for point in polygon]
|
||||
x_min, x_max = min(xs), max(xs)
|
||||
y_min, y_max = min(ys), max(ys)
|
||||
w = x_max - x_min
|
||||
h = y_max - y_min
|
||||
center_x = (x_min + x_max) // 2
|
||||
center_y = (y_min + y_max) // 2
|
||||
|
||||
# 过滤左右边缘区域的文字（可能是水印或logo）
|
||||
if center_x > width * 0.85 or center_x < width * 0.15:
|
||||
continue
|
||||
if detection.get('Confidence', 0) <= 75:
|
||||
continue
|
||||
|
||||
detected_text = detection.get("DetectedText", "")
|
||||
if not check_subtitle_quality_enhanced(detected_text) and detection.get('Confidence', 0) <= 90:
|
||||
continue
|
||||
if len(detected_text) < 5:
|
||||
continue
|
||||
# 只要中文字幕：没有中文字符的(纯英文/数字)不当作候选字幕区域
|
||||
if require_chinese and not has_chinese(detected_text):
|
||||
continue
|
||||
|
||||
# 水印检测：文本相似度>0.7且位置相近
|
||||
is_watermark = False
|
||||
for existing_text, pos_list in text_positions.items():
|
||||
similarity = SequenceMatcher(None, detected_text, existing_text).ratio()
|
||||
if similarity > 0.7:
|
||||
# 检查位置是否相近
|
||||
for pos in pos_list:
|
||||
if pos["idx"] != idx and abs(pos["y"] - center_y) < h:
|
||||
is_watermark = True
|
||||
watermark_logs.append(
|
||||
f"第{idx}帧可能的水印: {detected_text} (位置: {center_x}, {center_y})")
|
||||
logger.log(f"检测到可能的水印: {detected_text} (位置: {center_x}, {center_y})",
|
||||
program="merge")
|
||||
# 标记原文本为水印
|
||||
pos["is_watermark"] = True
|
||||
break
|
||||
if is_watermark:
|
||||
break
|
||||
|
||||
if is_watermark:
|
||||
continue
|
||||
|
||||
# 记录文本位置
|
||||
text_positions[detected_text].append({
|
||||
"text": detected_text,
|
||||
"x": center_x,
|
||||
"y": center_y,
|
||||
"is_watermark": False,
|
||||
"idx": idx
|
||||
})
|
||||
|
||||
# 存储候选区域
|
||||
candidate_regions.append({
|
||||
"x_min": x_min,
|
||||
"y_min": y_min,
|
||||
"x_max": x_max,
|
||||
"y_max": y_max,
|
||||
"width": w,
|
||||
"height": h
|
||||
})
|
||||
|
||||
# 如果没有候选区域，检查首尾帧是否有可能存在字幕
|
||||
if not candidate_regions:
|
||||
# 检查前 keep_num 张和后 keep_num 张是否有文字（并发）
|
||||
keep_num = 9
|
||||
n = len(image_files)
|
||||
check_indices = sorted(set(
|
||||
list(range(min(keep_num, n))) + list(range(max(0, n - keep_num), n))
|
||||
))
|
||||
fb_cache = _ocr_frames_concurrent(image_files, input_dir, check_indices, ocr_workers)
|
||||
has_text = False
|
||||
for idx in check_indices:
|
||||
ocr_result = fb_cache.get(idx)
|
||||
if ocr_result:
|
||||
ocr_data = json.loads(ocr_result)
|
||||
if "TextDetections" in ocr_data and ocr_data["TextDetections"]:
|
||||
has_text = True
|
||||
break
|
||||
|
||||
if not has_text:
|
||||
logger.log("保留的前后9张图片均未检测到任何文字，认为该视频没有字幕", level="warning", program="merge")
|
||||
return None, sample_indices, watermark_logs, candidate_regions
|
||||
|
||||
# 无候选区域但有文字：回退到主区域检测
|
||||
main_region = detect_main_region(image_files, input_dir)
|
||||
x, y, w, h = main_region
|
||||
# 只取主区域下半部分
|
||||
return (0, y + h // 2, width, h // 2), sample_indices, watermark_logs, candidate_regions
|
||||
|
||||
# 逆向清理：只清理画面上半部分（y < height*0.5）的水印候选区域
|
||||
# 下半部分即便被误判为水印也保留，避免把字幕字幕一并删掉
|
||||
confirmed_watermark_ys = {
|
||||
pos["y"]
|
||||
for pos_list in text_positions.values()
|
||||
for pos in pos_list
|
||||
if pos.get("is_watermark")
|
||||
}
|
||||
top_watermark_ys = {wy for wy in confirmed_watermark_ys if wy < height * 0.5}
|
||||
if top_watermark_ys:
|
||||
before = len(candidate_regions)
|
||||
candidate_regions = [
|
||||
r for r in candidate_regions
|
||||
if not any(
|
||||
abs((r["y_min"] + r["y_max"]) // 2 - wy) < 80
|
||||
for wy in top_watermark_ys
|
||||
)
|
||||
]
|
||||
removed = before - len(candidate_regions)
|
||||
if removed:
|
||||
logger.log(f"逆向移除 {removed} 个上半区水印候选区域", program="merge")
|
||||
|
||||
# 过滤过小的区域（高度<15）
|
||||
candidate_regions = [r for r in candidate_regions if r["height"] >= 15]
|
||||
|
||||
# 水印移除/高度过滤后可能把候选全部清空 → 判无字幕(否则下面 min() 会因空序列崩成"拼接失败")
|
||||
if not candidate_regions:
|
||||
logger.log("候选字幕区域被过滤后为空(高度过小/全是水印)，判无字幕", level="warning", program="merge")
|
||||
return None, sample_indices, watermark_logs, candidate_regions
|
||||
|
||||
# 冗余判定：用"不同垂直位置簇"的数量，而非候选框总数。
|
||||
# 正常字幕会在多帧的同一位置重复出现，框数随采样帧数增长但集中在少数垂直位置；
|
||||
# 纯文本/水印过多的视频才会把文字散布到很多不同的垂直位置。
|
||||
region_cys = sorted((r["y_min"] + r["y_max"]) // 2 for r in candidate_regions)
|
||||
distinct_bands = 1 if region_cys else 0
|
||||
for _i in range(1, len(region_cys)):
|
||||
if region_cys[_i] - region_cys[_i - 1] > 80:
|
||||
distinct_bands += 1
|
||||
if distinct_bands >= bands_thr:
|
||||
logger.log(f"冗余文字过多（文字分散在{distinct_bands}个垂直位置，共{len(candidate_regions)}个候选区域），视频将被过滤",
|
||||
level="warning", program="merge")
|
||||
return None, sample_indices, watermark_logs, candidate_regions
|
||||
|
||||
# 合并所有候选区域
|
||||
all_y_min = min(r["y_min"] for r in candidate_regions)
|
||||
all_y_max = max(r["y_max"] for r in candidate_regions)
|
||||
merged_h = all_y_max - all_y_min
|
||||
# 扩展边界
|
||||
expanded_y_min = max(0, int(all_y_min - 65))
|
||||
expanded_y_max = min(height, int(all_y_max + 65))
|
||||
final_h = expanded_y_max - expanded_y_min
|
||||
|
||||
# 如果合并区域过高（>60% 画面高度），用聚类找最大候选簇，避免退回全帧
|
||||
if final_h >= height * 0.6 and len(candidate_regions) >= 4:
|
||||
logger.log("合并区域过高，聚类寻找最大字幕区域簇", program="merge")
|
||||
best_cluster = []
|
||||
for region in candidate_regions:
|
||||
cy = (region["y_min"] + region["y_max"]) // 2
|
||||
cluster = [r for r in candidate_regions
|
||||
if abs((r["y_min"] + r["y_max"]) // 2 - cy) < 200]
|
||||
# 优先更大的簇；同等大小时选中心 y 更大的（靠近底部的字幕）
|
||||
if len(cluster) > len(best_cluster) or (
|
||||
len(cluster) == len(best_cluster) and best_cluster and
|
||||
sum((r["y_min"] + r["y_max"]) // 2 for r in cluster) >
|
||||
sum((r["y_min"] + r["y_max"]) // 2 for r in best_cluster)):
|
||||
best_cluster = cluster
|
||||
|
||||
if best_cluster:
|
||||
c_y_min = min(r["y_min"] for r in best_cluster)
|
||||
c_y_max = max(r["y_max"] for r in best_cluster)
|
||||
exp_y_min = max(0, c_y_min - 65)
|
||||
exp_y_max = min(height, c_y_max + 65)
|
||||
logger.log(f"聚类选定字幕区域: y={exp_y_min}-{exp_y_max}", program="merge")
|
||||
return (0, exp_y_min, width, exp_y_max - exp_y_min), sample_indices, watermark_logs, candidate_regions
|
||||
|
||||
main_region = detect_main_region(image_files, input_dir)
|
||||
x, y, w, h = main_region
|
||||
return (0, y, width, h), sample_indices, watermark_logs, candidate_regions
|
||||
|
||||
return (0, expanded_y_min, width, final_h), sample_indices, watermark_logs, candidate_regions
|
||||
|
||||
|
||||
# ==================== 拼图配置计算 ====================
|
||||
def calculate_optimal_tile_configuration(subtitle_region: Tuple[int, int, int, int],
|
||||
tile_width: int,
|
||||
tile_height: int,
|
||||
first_img_path: str,
|
||||
logger,
|
||||
max_base64_size: int = 10 * 1024 * 1024) -> Tuple[int, int]:
|
||||
"""
|
||||
计算最优的拼图配置（横向和纵向数量）。
|
||||
|
||||
:return: (tile_width, tile_height)
|
||||
"""
|
||||
x, y, w, h = subtitle_region
|
||||
BORDER_SIZE = 40
|
||||
tile_width_px = w + 2 * BORDER_SIZE
|
||||
tile_height_px = h + 2 * BORDER_SIZE
|
||||
|
||||
# 估算单个图块的大小（基于第一张图片）
|
||||
img = cv2.imread(first_img_path)
|
||||
if img is None:
|
||||
raise ValueError(f"无法读取图片: {first_img_path}")
|
||||
img_h, img_w = img.shape[:2]
|
||||
img_size = os.path.getsize(first_img_path)
|
||||
single_tile_size = (tile_width_px * tile_height_px) / (img_w * img_h) * img_size * 3 # RGB估算
|
||||
single_tile_base64 = single_tile_size * 4 // 3 # base64 增加 33%
|
||||
|
||||
# 每张合成图最多拼几块：硬上限(可配置) 与 base64 10MB 限制 取小。
|
||||
# 加大上限 → OCR 请求数下降(每张装更多帧)，但单次 OCR 图更大、且文字条更密，
|
||||
# 过大有 OCR 漏识别/超分辨率上限的风险，调大后务必抽查 SRT 质量。
|
||||
hard_cap = config.get_subtitle_config().get('max_tiles_per_composite', 40)
|
||||
try:
|
||||
hard_cap = int(hard_cap)
|
||||
except (TypeError, ValueError):
|
||||
hard_cap = 40
|
||||
max_tiles = min(hard_cap, int(max_base64_size // single_tile_base64))
|
||||
max_tiles = max(1, max_tiles)
|
||||
# 横向块数上限随总块数放宽(块多时仍想接近 0.5 的宽高比，避免拼成超细长条)
|
||||
w_cap = max(7, int(round((max_tiles * 0.5) ** 0.5)) + 1)
|
||||
logger.log(f"单个图块 base64 大小: {single_tile_base64 / 1024:.1f} KB", program="merge")
|
||||
logger.log(f"最大允许图块数量: {max_tiles} (硬上限 {hard_cap})", program="merge")
|
||||
|
||||
# 寻找最佳宽高比（接近 0.5，即横向多于纵向）
|
||||
best_config = (tile_width, tile_height) # 默认
|
||||
best_ratio_diff = float('inf')
|
||||
for w_tiles in range(1, min(max_tiles, w_cap) + 1):
|
||||
h_tiles = int(max_tiles / w_tiles)
|
||||
aspect = w_tiles / h_tiles
|
||||
diff = abs(aspect - 0.5)
|
||||
if diff < best_ratio_diff:
|
||||
best_ratio_diff = diff
|
||||
best_config = (w_tiles, h_tiles)
|
||||
|
||||
logger.log(f"推荐拼图配置: {best_config[0]} x {best_config[1]}", program="merge")
|
||||
return best_config
|
||||
|
||||
|
||||
# ==================== 音频提取（可复用） ====================
|
||||
def extract_audio_text(video_path: str, logger) -> Dict[str, str]:
|
||||
"""
|
||||
从视频中提取开头和结尾的音频文字。
|
||||
|
||||
:return: {'first_3s': str, 'last_3s': str}
|
||||
"""
|
||||
audio_text = {}
|
||||
if not video_path or not os.path.exists(video_path):
|
||||
return audio_text
|
||||
|
||||
logger.log("开始提取视频中的音频文字...", program="merge")
|
||||
# 从配置读取 whisper 模型大小（默认 small）
|
||||
model_size = config.get_audio_config().get('model_size', 'small')
|
||||
whisper_manager.load_model(model_size)
|
||||
|
||||
temp_audio_path = os.path.join(tempfile.gettempdir(),
|
||||
f"temp_audio_{os.path.basename(video_path)}.wav")
|
||||
try:
|
||||
total_duration = whisper_manager.extract_audio_from_video(video_path, temp_audio_path, logger)
|
||||
if total_duration is None:
|
||||
logger.log("音频提取失败，跳过音频文字识别", program="merge")
|
||||
return audio_text
|
||||
|
||||
if total_duration >= 5:
|
||||
first_5s_text = whisper_manager.extract_audio_text_with_whisper(
|
||||
temp_audio_path, 0, min(5, total_duration), logger
|
||||
)
|
||||
audio_text["first_3s"] = first_5s_text
|
||||
logger.log(f"前五秒音频文字: {first_5s_text}", program="merge")
|
||||
|
||||
if total_duration >= 10:
|
||||
last_5s_start = int(total_duration - 5)
|
||||
last_5s_text = whisper_manager.extract_audio_text_with_whisper(
|
||||
temp_audio_path, last_5s_start, total_duration, logger
|
||||
)
|
||||
audio_text["last_3s"] = last_5s_text
|
||||
logger.log(f"后五秒音频文字: {last_5s_text}", program="merge")
|
||||
finally:
|
||||
try:
|
||||
os.unlink(temp_audio_path)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return audio_text
|
||||
|
||||
|
||||
# ==================== 主处理函数 ====================
|
||||
def process_images(input_dir, output_dir, logger, tile_width=3, tile_height=6, keep_num=9, interval_seconds=None,
|
||||
video_path=None):
|
||||
"""
|
||||
处理图片拼接并生成元数据，自动计算最优拼图配置
|
||||
|
||||
参数:
|
||||
input_dir: 输入图片目录
|
||||
output_dir: 输出目录
|
||||
keep_num: 保留的首尾图片数量
|
||||
"""
|
||||
# 创建输出目录
|
||||
logger.start_program("merge")
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
composite_dir = os.path.join(output_dir, "composite")
|
||||
os.makedirs(composite_dir, exist_ok=True)
|
||||
|
||||
# 获取所有图片文件并排序
|
||||
image_files = [f for f in os.listdir(input_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
|
||||
image_files.sort(key=natural_sort_key)
|
||||
total_images = len(image_files)
|
||||
duration = total_images / 3
|
||||
# 确定字幕区域（通过OCR检测）
|
||||
subtitle_region, sample_indices, watermark, candidate_regions = detect_subtitle_region(image_files, input_dir)
|
||||
|
||||
if subtitle_region is None:
|
||||
logger.end_program("merge")
|
||||
return duration, None, None, None, None, None, sample_indices, watermark, candidate_regions
|
||||
logger.log(f"通过OCR检测确定字幕区域: {subtitle_region}", program="merge")
|
||||
first_img_path = os.path.join(input_dir, image_files[0])
|
||||
|
||||
# 提取音频文字（如果提供了视频路径）
|
||||
audio_text = {}
|
||||
if video_path and os.path.exists(video_path):
|
||||
logger.log("开始提取视频中的音频文字...", program="merge")
|
||||
model_size = config.get_audio_config().get('model_size', 'small')
|
||||
whisper_manager.load_model(model_size)
|
||||
# 创建临时音频文件
|
||||
temp_audio_path = os.path.join(tempfile.gettempdir(), f"temp_audio_{os.path.basename(video_path)}.wav")
|
||||
total_duration = whisper_manager.extract_audio_from_video(video_path, temp_audio_path, logger)
|
||||
# 从视频提取音频
|
||||
if total_duration is not None:
|
||||
|
||||
# 提取前五秒文字
|
||||
if total_duration >= 5:
|
||||
first_5s_text = whisper_manager.extract_audio_text_with_whisper(temp_audio_path, 0,
|
||||
min(5, total_duration), logger)
|
||||
audio_text["first_3s"] = first_5s_text
|
||||
logger.log(f"前五秒音频文字: {first_5s_text}", program="merge")
|
||||
|
||||
# 提取后五秒文字
|
||||
if total_duration >= 10:
|
||||
last_5s_start = int(total_duration - 5)
|
||||
last_5s_text = whisper_manager.extract_audio_text_with_whisper(temp_audio_path, last_5s_start,
|
||||
total_duration, logger)
|
||||
audio_text["last_3s"] = last_5s_text
|
||||
logger.log(f"后五秒音频文字: {last_5s_text}", program="merge")
|
||||
|
||||
# 清理临时音频文件
|
||||
try:
|
||||
os.unlink(temp_audio_path)
|
||||
except OSError:
|
||||
pass
|
||||
else:
|
||||
logger.log("音频提取失败，跳过音频文字识别", program="merge")
|
||||
# 边框设置
|
||||
BORDER_SIZE = 40
|
||||
|
||||
# 保留前keep_num张和最后keep_num张图片
|
||||
kept_first = image_files[:keep_num]
|
||||
kept_last = image_files[-keep_num:] if len(image_files) > 2 * keep_num else []
|
||||
|
||||
# 需要处理的图片（中间部分）
|
||||
process_files = image_files[keep_num:-keep_num] if len(image_files) > 2 * keep_num else image_files[keep_num:]
|
||||
|
||||
# 归一化图块像素预算：高分辨率帧或过高的字幕区域会让单个图块的 base64 过大，
|
||||
# 撑到 10MB 上限后每张合成图能塞的块数骤减（从~32 掉到 5~8），导致合成图（=OCR请求）
|
||||
# 数量激增 3~6 倍。按基准像素量（约 1080p 单字幕条）等比缩小裁剪后的字幕条，使每张
|
||||
# 合成图的块数稳定在~32，OCR 成本与源分辨率/区域高度无关；文字分辨率仍足够 OCR。
|
||||
x0, y0, w0, h0 = subtitle_region
|
||||
no_composite = config.get_subtitle_config().get('no_composite', False)
|
||||
if no_composite:
|
||||
# 不拼接：每个字幕条单独成图(tile 1×1)、不降采样(单条 base64 很小，全分辨率最准)。
|
||||
# 下游 OCR 映射/网格搜索/去重完全复用，只是每张图=1 条、OCR 次数=帧数。
|
||||
tile_scale = 1.0
|
||||
w_scaled, h_scaled = w0, h0
|
||||
tile_width, tile_height = 1, 1
|
||||
logger.log("不拼接模式：每个字幕条单独 OCR(tile 1×1，全分辨率)", program="merge")
|
||||
else:
|
||||
# 归一化图块像素预算：高分辨率帧或过高的字幕区域会让单个图块的 base64 过大，
|
||||
# 撑到 10MB 上限后每张合成图能塞的块数骤减（从~32 掉到 5~8），导致合成图（=OCR请求）
|
||||
# 数量激增 3~6 倍。按基准像素量（约 1080p 单字幕条）等比缩小裁剪后的字幕条，使每张
|
||||
# 合成图的块数稳定在~32，OCR 成本与源分辨率/区域高度无关；文字分辨率仍足够 OCR。
|
||||
TARGET_TILE_AREA = 1920 * 240
|
||||
tile_scale = min(1.0, math.sqrt(TARGET_TILE_AREA / float(max(1, w0 * h0))))
|
||||
w_scaled = max(1, int(round(w0 * tile_scale)))
|
||||
h_scaled = max(1, int(round(h0 * tile_scale)))
|
||||
if tile_scale < 1.0:
|
||||
logger.log(f"字幕条降采样: {w0}x{h0} → {w_scaled}x{h_scaled} (scale={tile_scale:.2f})", program="merge")
|
||||
tile_width, tile_height = calculate_optimal_tile_configuration(
|
||||
(x0, y0, w_scaled, h_scaled), tile_width, tile_height, first_img_path, logger)
|
||||
|
||||
# 初始化元数据
|
||||
metadata = {
|
||||
"tile_width": tile_width,
|
||||
"tile_height": tile_height,
|
||||
"subtitle_region": subtitle_region,
|
||||
"border_size": BORDER_SIZE,
|
||||
"interval_seconds": interval_seconds,
|
||||
"audio_text": audio_text,
|
||||
"kept_first": [],
|
||||
"kept_last": [],
|
||||
"composites": []
|
||||
}
|
||||
|
||||
# 复制保留的图片到composite文件夹
|
||||
kept_first_new_names = []
|
||||
for i, img in enumerate(kept_first):
|
||||
src = os.path.join(input_dir, img)
|
||||
dst_filename = f"001_{i + 1:02d}_original_{img}"
|
||||
dst = os.path.join(composite_dir, dst_filename)
|
||||
|
||||
if not os.path.exists(dst):
|
||||
with Image.open(src) as im:
|
||||
bordered_img = Image.new('RGB',
|
||||
(im.width + 2 * BORDER_SIZE, im.height + 2 * BORDER_SIZE),
|
||||
(255, 255, 255))
|
||||
bordered_img.paste(im, (BORDER_SIZE, BORDER_SIZE))
|
||||
|
||||
if img.lower().endswith('.png'):
|
||||
bordered_img.save(dst, format='PNG')
|
||||
else:
|
||||
bordered_img.save(dst, format='JPEG', quality=95)
|
||||
|
||||
kept_first_new_names.append(dst_filename)
|
||||
|
||||
metadata["kept_first"] = kept_first_new_names
|
||||
|
||||
# 处理中间图片进行拼接
|
||||
if process_files:
|
||||
x, y, w, h = subtitle_region # 原始帧上的裁剪坐标
|
||||
ws, hs = w_scaled, h_scaled # 降采样后的图块尺寸（拼接/坐标用此）
|
||||
|
||||
# 计算图块和合成图尺寸（基于降采样后的尺寸）
|
||||
tile_img_width = ws + 2 * BORDER_SIZE
|
||||
tile_img_height = hs + 2 * BORDER_SIZE
|
||||
composite_width = tile_width * tile_img_width
|
||||
composite_height = tile_height * tile_img_height
|
||||
|
||||
# 分批：每批 = 一张合成图
|
||||
tiles_per_composite = tile_width * tile_height
|
||||
batches = []
|
||||
for i in range(0, len(process_files), tiles_per_composite):
|
||||
batches.append((len(batches) + 1, process_files[i:i + tiles_per_composite]))
|
||||
|
||||
def _build_composite(batch):
|
||||
"""处理一张合成图：读源帧→裁剪→贴块→写盘。返回 (序号, composite_data)。
|
||||
热点：逐帧读 1080p JPEG + 裁剪 + 写盘，原为单线程；这里并行化（PIL 编解码
|
||||
在 C 层释放 GIL，线程池能真正并行，按核数提速）。"""
|
||||
composite_num, batch_files = batch
|
||||
composite_img = Image.new('RGB', (composite_width, composite_height), (255, 255, 255))
|
||||
composite_filename = f"002_{composite_num:04d}_composite.jpg"
|
||||
composite_data = {
|
||||
"composite_id": f"composite_{composite_num:04d}",
|
||||
"file_name": composite_filename,
|
||||
"tiles": []
|
||||
}
|
||||
for pos in range(tiles_per_composite):
|
||||
row = pos // tile_width
|
||||
col = pos % tile_width
|
||||
x_pos = col * tile_img_width
|
||||
y_pos = row * tile_img_height
|
||||
|
||||
if pos < len(batch_files):
|
||||
img_name = batch_files[pos]
|
||||
img_path = os.path.join(input_dir, img_name)
|
||||
with Image.open(img_path) as img:
|
||||
subtitle_img = img.crop((x, y, x + w, y + h))
|
||||
if tile_scale < 1.0:
|
||||
subtitle_img = subtitle_img.resize((ws, hs))
|
||||
# 合成图背景已是白色，直接把字幕条贴到带边框的偏移位置
|
||||
composite_img.paste(subtitle_img, (x_pos + BORDER_SIZE, y_pos + BORDER_SIZE))
|
||||
composite_data["tiles"].append({
|
||||
"position": pos,
|
||||
"x": x_pos + BORDER_SIZE,
|
||||
"y": y_pos + BORDER_SIZE,
|
||||
"width": ws,
|
||||
"height": hs,
|
||||
"source_image": img_name,
|
||||
"is_blank": False
|
||||
})
|
||||
else:
|
||||
# 空白块：合成图背景本就是白色，无需粘贴
|
||||
composite_data["tiles"].append({
|
||||
"position": pos,
|
||||
"x": x_pos + BORDER_SIZE,
|
||||
"y": y_pos + BORDER_SIZE,
|
||||
"width": ws,
|
||||
"height": hs,
|
||||
"source_image": None,
|
||||
"is_blank": True
|
||||
})
|
||||
|
||||
output_path = os.path.join(composite_dir, composite_filename)
|
||||
composite_img.save(output_path, format='JPEG', quality=95)
|
||||
return composite_num, composite_data
|
||||
|
||||
# 并行处理所有合成图，结果按序号排序回填（顺序不能乱，元数据要对得上时间轴）
|
||||
merge_workers = config.get_subtitle_config().get('merge_workers') or (os.cpu_count() or 8)
|
||||
merge_workers = max(1, min(int(merge_workers), len(batches)))
|
||||
results = []
|
||||
with ThreadPoolExecutor(max_workers=merge_workers) as executor:
|
||||
for r in executor.map(_build_composite, batches):
|
||||
results.append(r)
|
||||
results.sort(key=lambda t: t[0])
|
||||
metadata["composites"] = [cd for _, cd in results]
|
||||
|
||||
# 处理最后保留的图片
|
||||
kept_last_new_names = []
|
||||
for i, img in enumerate(kept_last):
|
||||
src = os.path.join(input_dir, img)
|
||||
dst_filename = f"003_{i + 1:02d}_original_{img}"
|
||||
dst = os.path.join(composite_dir, dst_filename)
|
||||
|
||||
if not os.path.exists(dst):
|
||||
with Image.open(src) as im:
|
||||
bordered_img = Image.new('RGB',
|
||||
(im.width + 2 * BORDER_SIZE, im.height + 2 * BORDER_SIZE),
|
||||
(255, 255, 255))
|
||||
bordered_img.paste(im, (BORDER_SIZE, BORDER_SIZE))
|
||||
|
||||
if img.lower().endswith('.png'):
|
||||
bordered_img.save(dst, format='PNG')
|
||||
else:
|
||||
bordered_img.save(dst, format='JPEG', quality=95)
|
||||
|
||||
kept_last_new_names.append(dst_filename)
|
||||
|
||||
metadata["kept_last"] = kept_last_new_names
|
||||
|
||||
# 保存元数据（保存到 composite_dir，避免被清理时删除）
|
||||
metadata_path = os.path.join(composite_dir, "metadata.json")
|
||||
with open(metadata_path, 'w') as f:
|
||||
json.dump(metadata, f, indent=2)
|
||||
|
||||
logger.log(f"处理完成，拼图配置: {tile_width}x{tile_height}", program="merge")
|
||||
logger.log(f"共生成 {len(metadata['composites'])} 张合成图", program="merge")
|
||||
logger.log(f"元数据已保存到: {metadata_path}", program="merge")
|
||||
logger.end_program("merge")
|
||||
return duration, subtitle_region, tile_width, tile_height, len(
|
||||
metadata['composites']), audio_text, sample_indices, watermark, candidate_regions
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# 示例用法
|
||||
input_directory = "input_images"
|
||||
output_directory = "output"
|
||||
from log.log_manager import logger # 假设已有全局 logger
|
||||
|
||||
result = process_images(input_directory, output_directory, logger)
|
||||
471
tool/process.py
Normal file
471
tool/process.py
Normal file
@@ -0,0 +1,471 @@
|
||||
from config.config_utils import config, get_ffprobe
|
||||
import argparse
|
||||
import hashlib
|
||||
import logging
|
||||
import subprocess
|
||||
import shutil
|
||||
import sys
|
||||
import os
|
||||
import json
|
||||
import time
|
||||
from tool.log.log_manager import UnifiedLogger
|
||||
from tool.VideoProcessing.multithread_algo import split_frame
|
||||
|
||||
|
||||
def derive_video_name(cos_path: str) -> str:
|
||||
"""把 COS 完整路径变成本地工作目录/文件名用的唯一标识:
|
||||
<原文件名(不含扩展名)>_<完整路径MD5前8位>。
|
||||
|
||||
不能只用文件名 —— 源数据里大量不同父目录复用同样的数字文件名(如
|
||||
cos://bkt/0021/17688/00046.mp4 和 cos://bkt/0021/17690/00046.mp4), 纯文件名当本地
|
||||
标识会导致两个不同视频被并发处理到时互相踩踏同一份本地文件(下载/抽帧/拼接的中间文件
|
||||
全挤在同一个 ./split/00046 目录), 引发 "Directory not empty" 崩溃和一连串连锁失败
|
||||
(实测占某台机器全部 error 记录的 100%)。
|
||||
|
||||
保留原文件名前缀是为了日志/目录名还能看出是哪个视频, 不至于变成纯哈希、丧失可读性；
|
||||
main.py 里 process_video_task() 必须调用同一个函数(而不是各自按文件名重算), 两处
|
||||
对同一个 cos_path 才会算出同一个 video_name。"""
|
||||
filename = os.path.basename(cos_path.strip())
|
||||
stem = os.path.splitext(filename)[0]
|
||||
digest = hashlib.md5(cos_path.strip().encode("utf-8")).hexdigest()[:8]
|
||||
return f"{stem}_{digest}"
|
||||
|
||||
|
||||
def _probe_duration(video_path):
|
||||
"""用 ffprobe 取视频时长(秒)；失败返回 None。
|
||||
先取流时长, 再取容器(format)时长——有些封装(尤其AV1/部分mp4)流里没duration、只在format里。"""
|
||||
# ① 流时长
|
||||
try:
|
||||
r = subprocess.run([
|
||||
get_ffprobe(), '-v', 'error', '-select_streams', 'v:0',
|
||||
'-show_entries', 'stream=duration', '-of', 'json', video_path
|
||||
], capture_output=True, text=True)
|
||||
d = float(json.loads(r.stdout)['streams'][0].get('duration', 0) or 0)
|
||||
if d > 0:
|
||||
return d
|
||||
except Exception:
|
||||
pass
|
||||
# ② 容器时长(兜底)
|
||||
try:
|
||||
r = subprocess.run([
|
||||
get_ffprobe(), '-v', 'error',
|
||||
'-show_entries', 'format=duration', '-of', 'json', video_path
|
||||
], capture_output=True, text=True)
|
||||
d = float(json.loads(r.stdout).get('format', {}).get('duration', 0) or 0)
|
||||
if d > 0:
|
||||
return d
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def _append_long_video(cos_path, list_path):
|
||||
"""把被"视频过长"过滤掉的视频 COS 路径追加到一份 txt，供以后(GPU机器/放大上限)补跑。
|
||||
进程池下多个 worker 进程会并发写同一文件：用 O_APPEND 单行追加(POSIX 下小写原子)，
|
||||
失败也只记日志、绝不影响主流程。"""
|
||||
if not cos_path or not list_path:
|
||||
return
|
||||
try:
|
||||
with open(list_path, 'a', encoding='utf-8') as f:
|
||||
f.write(cos_path.rstrip('\n') + '\n')
|
||||
except Exception:
|
||||
pass
|
||||
from tool.merge import process_images
|
||||
from tool.getsrt import processsrt, OcrIncompleteError
|
||||
|
||||
# 配置日志
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(levelname)s - %(message)s',
|
||||
handlers=[
|
||||
logging.FileHandler('video_processing.log'),
|
||||
logging.StreamHandler(sys.stdout)
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def delete_local_folder(folder_path):
|
||||
"""
|
||||
删除本地文件夹，但保留子文件夹composite及其内容
|
||||
:param folder_path: 文件夹路径
|
||||
:return:
|
||||
"""
|
||||
try:
|
||||
if not os.path.exists(folder_path):
|
||||
logging.warning(f"文件夹不存在，无需删除: {folder_path}")
|
||||
return True
|
||||
|
||||
# 定义要保留的子文件夹路径
|
||||
composite_path = os.path.join(folder_path, 'composite')
|
||||
|
||||
# 如果composite子文件夹存在，先临时保存它
|
||||
temp_dir = None
|
||||
if os.path.exists(composite_path):
|
||||
# 创建一个临时目录来保存composite文件夹
|
||||
temp_dir = os.path.join(os.path.dirname(folder_path), f"temp_composite_{os.getpid()}")
|
||||
shutil.move(composite_path, temp_dir)
|
||||
|
||||
# 删除整个文件夹
|
||||
shutil.rmtree(folder_path)
|
||||
|
||||
# 如果之前有composite文件夹，将其移回原位
|
||||
if temp_dir and os.path.exists(temp_dir):
|
||||
# rmtree 删除了父目录本身，需要先重建
|
||||
os.makedirs(folder_path, exist_ok=True)
|
||||
shutil.move(temp_dir, composite_path)
|
||||
logging.info(f"已删除本地文件夹内容，但保留了composite子文件夹: {folder_path}")
|
||||
else:
|
||||
logging.info(f"已删除本地文件夹: {folder_path}")
|
||||
|
||||
return True
|
||||
except Exception as e:
|
||||
logging.error(f"删除文件夹失败: {e}")
|
||||
# 尝试恢复临时文件夹以防出错
|
||||
if temp_dir and os.path.exists(temp_dir):
|
||||
try:
|
||||
os.makedirs(folder_path, exist_ok=True)
|
||||
shutil.move(temp_dir, os.path.join(folder_path, 'composite'))
|
||||
except Exception as restore_error:
|
||||
logging.error(f"恢复composite文件夹失败: {restore_error}")
|
||||
return False
|
||||
|
||||
|
||||
def download_from_cos(cos_target, cos_source, coscli_path, cos_config_path, local_dest_path):
|
||||
"""
|
||||
从COS下载文件到 local_dest_path。
|
||||
:param cos_target: 目录
|
||||
:param cos_source: 完整 COS 路径
|
||||
:param coscli_path: cos exe文件地址
|
||||
:param cos_config_path: cos 配置文件
|
||||
:param local_dest_path: 下载到的本地目标路径 —— 必须是调用方(process_single_video)按
|
||||
derive_video_name() 算出的 collision-safe 路径, 不能在这里自己按 os.path.basename(cos_source)
|
||||
重新拼一个 './video/<文件名>' —— 之前就是这么写的, 跟调用方检查是否存在用的 local_video_path
|
||||
(带路径MD5后缀)对不上, 下载明明成功了、下一步却判"文件不存在", 100%必现。
|
||||
:return:
|
||||
"""
|
||||
try:
|
||||
# cos_source 是完整 COS 路径时直接使用，否则拼接目录
|
||||
full_cos_path = cos_source if cos_source.startswith("cos://") else f"{cos_target}{cos_source}"
|
||||
cmd = [f"{coscli_path}", "-c", f"{cos_config_path}", "cp",
|
||||
full_cos_path, local_dest_path]
|
||||
# 执行命令（带超时：网络卡住时不会永久阻塞 worker，超时即判下载失败）
|
||||
result = subprocess.run(cmd, capture_output=True, text=True, timeout=1800)
|
||||
|
||||
if result.returncode == 0:
|
||||
logging.info(f"成功下载到: {local_dest_path}")
|
||||
return cos_source
|
||||
else:
|
||||
logging.error(f"下载失败: {result.stderr}")
|
||||
return False
|
||||
|
||||
except subprocess.TimeoutExpired:
|
||||
logging.error(f"下载超时(>1800s)，跳过: {cos_source}")
|
||||
return False
|
||||
except Exception as e:
|
||||
logging.error(f"下载过程中出错: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def process_single_video(cos_source, cos_target):
|
||||
"""
|
||||
单个视频处理流程
|
||||
:param cos_source: 单个视频名称
|
||||
:param cos_target: 上传位置
|
||||
:return: True 处理成功, False 处理失败
|
||||
"""
|
||||
# 1. 读取参数
|
||||
split_frame_config = config.get_splitframe_config()
|
||||
download_flag = split_frame_config.get('download')
|
||||
coscli_path = split_frame_config.get('coscli_path')
|
||||
cos_config_path = split_frame_config.get('cosconfig_path')
|
||||
filename = os.path.basename(cos_source)
|
||||
root_path = "./video"
|
||||
video_name = derive_video_name(cos_source)
|
||||
local_video_path = os.path.join(root_path, video_name + os.path.splitext(filename)[1])
|
||||
|
||||
# 使用线程安全的日志器，避免多线程下重置单例导致状态混乱
|
||||
logger = UnifiedLogger()
|
||||
logger.reset()
|
||||
logger.init_logging(f"log/{video_name}.json")
|
||||
logger.start_program("main")
|
||||
|
||||
try:
|
||||
return _process_video_pipeline(
|
||||
cos_source, cos_target, video_name, local_video_path,
|
||||
split_frame_config, download_flag, coscli_path, cos_config_path, logger
|
||||
)
|
||||
except OcrIncompleteError as e:
|
||||
# 不丢帧: OCR没全成功 → 不产出带洞SRT, 把视频记到"OCR失败待重跑"列表, 判失败(以后用另一jsonl补跑)
|
||||
try:
|
||||
fail_list = config.get_subtitle_config().get('ocr_fail_list', './OCR失败待重跑.txt')
|
||||
_append_long_video(cos_source, fail_list)
|
||||
except Exception:
|
||||
pass
|
||||
logger.log(f"视频{video_name} OCR未完成(不丢帧→不产出带洞SRT, 已记入待重跑): {e}",
|
||||
level="error", program="main")
|
||||
_save_error_metadata(
|
||||
cos_source, video_name, error_reason=f"OcrIncomplete: {e}",
|
||||
download_flag=download_flag, interval_seconds=split_frame_config.get('interval_seconds', 0.33)
|
||||
)
|
||||
logger.end_program("main")
|
||||
logger.save_logs()
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.log(f"视频{video_name}处理流程异常: {type(e).__name__}: {e}", level="error", program="main")
|
||||
_save_error_metadata(
|
||||
cos_source, video_name, error_reason=f"{type(e).__name__}: {e}",
|
||||
download_flag=download_flag, interval_seconds=split_frame_config.get('interval_seconds', 0.33)
|
||||
)
|
||||
logger.end_program("main")
|
||||
logger.save_logs()
|
||||
return False
|
||||
|
||||
|
||||
def _process_video_pipeline(cos_source, cos_target, video_name, local_video_path,
|
||||
split_frame_config, download_flag, coscli_path, cos_config_path, logger):
|
||||
"""
|
||||
视频处理内部流水线，将各阶段错误统一记录并抛出
|
||||
"""
|
||||
# 2. 进行数据下载
|
||||
logger.step(video_name, "开始")
|
||||
if download_flag:
|
||||
logger.log("从COS下载视频...", program="main")
|
||||
if not os.path.exists(local_video_path):
|
||||
if not download_from_cos(cos_target, cos_source, coscli_path, cos_config_path, local_video_path):
|
||||
raise RuntimeError(f"视频{video_name}下载失败")
|
||||
logger.step(video_name, "下载完成")
|
||||
|
||||
# 检查文件是否存在
|
||||
if not os.path.exists(local_video_path):
|
||||
raise FileNotFoundError(f"视频{video_name}不存在: {local_video_path}")
|
||||
|
||||
# 判断字幕文件是否存在
|
||||
name = os.path.basename(local_video_path).split(".")[0]
|
||||
if os.path.exists(f"./split/{name}/composite/{name}.srt"):
|
||||
return True
|
||||
|
||||
# 视频时长过滤（抽帧前 ffprobe 取时长，省掉抽帧/拼接/OCR；也避免超长视频撑爆内存/卡死worker）
|
||||
_sub_cfg = config.get_subtitle_config()
|
||||
min_dur = _sub_cfg.get('min_duration_seconds', 300)
|
||||
max_dur = _sub_cfg.get('max_duration_seconds', 5400) # 默认1.5小时上限
|
||||
vdur = _probe_duration(local_video_path)
|
||||
# 在【删视频前】把时长存进日志stat: ① ffprobe能处理AV1(cv2不行) ② delete_video删源视频后main.py就读不到了
|
||||
if vdur is not None and vdur > 0:
|
||||
logger.set_stat("video_duration", round(vdur, 2))
|
||||
if vdur is not None and 0 < vdur < min_dur:
|
||||
logger.step(video_name, "已过滤", f"视频过短 {vdur:.0f}s < {min_dur}s")
|
||||
raise ValueError(f"视频{video_name}被过滤: 视频过短({vdur:.0f}秒, 阈值{min_dur}秒)")
|
||||
if vdur is not None and max_dur and vdur > max_dur:
|
||||
logger.step(video_name, "已过滤", f"视频过长 {vdur:.0f}s > {max_dur}s")
|
||||
# 把超长视频的 COS 路径单独留一份，供以后在 GPU 机器/放大上限后补跑
|
||||
# (当前 run 仍记 数据.jsonl 跳过它，不占 worker)
|
||||
long_list = _sub_cfg.get('long_video_list', './超长待重跑.txt')
|
||||
_append_long_video(cos_source, long_list)
|
||||
raise ValueError(f"视频{video_name}被过滤: 视频过长({vdur:.0f}秒, 阈值{max_dur}秒)")
|
||||
|
||||
# 3. 进行视频抽帧
|
||||
interval_seconds = split_frame_config.get('interval_seconds', 0.33)
|
||||
try:
|
||||
_t = time.time()
|
||||
duration, planned_frames, actual_frames, backend = split_frame(local_video_path, interval_seconds)
|
||||
logger.set_stat("frame_count", actual_frames)
|
||||
logger.log("抽帧成功", program="main")
|
||||
logger.step(video_name, "抽帧完成",
|
||||
f"{actual_frames}帧, {backend}, {time.time()-_t:.0f}s, 视频{duration:.0f}s")
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"视频{video_name}抽帧处理失败: {e}") from e
|
||||
|
||||
# 抽帧成功但一帧没出: 视频没被 min/max_duration_seconds 过滤掉、真去抽了、但结果是空——
|
||||
# 不是"没字幕"的正常情况, 是抽帧本身出了问题(比如时长探测失败导致期望帧数算成0, 见
|
||||
# split_frame_ffmpeg 的时长兜底), 应该报错让人查, 而不是让空帧列表流进 process_images()
|
||||
# 炸出一个看不懂的 "list index out of range"。
|
||||
if actual_frames == 0:
|
||||
raise RuntimeError(f"视频{video_name}抽帧结果为0帧(时长{duration:.0f}s), 抽帧本身有问题, 非正常无字幕")
|
||||
|
||||
try:
|
||||
_t = time.time()
|
||||
result = process_images(
|
||||
f"./split/{video_name}", f"./split/{video_name}", logger,
|
||||
interval_seconds=interval_seconds, video_path=local_video_path)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"视频{video_name}图片拼接失败: {e}") from e
|
||||
|
||||
duration, subtitle_region, tile_width, tile_height, compositenum, audio_text, sample_indices, watermark, candidate_regions = result
|
||||
if subtitle_region is None:
|
||||
logger.step(video_name, "已过滤", "未检测到字幕区域")
|
||||
raise ValueError(f"视频{video_name}被过滤: 未检测到字幕区域")
|
||||
logger.step(video_name, "拼接完成", f"{compositenum}张合成图, {time.time()-_t:.0f}s")
|
||||
|
||||
# 4. 删除本地抽帧文件（保留 composite 子文件夹）
|
||||
# delete_folder = os.path.join("./split", video_name)
|
||||
# if not delete_local_folder(delete_folder):
|
||||
# logger.log("本地文件夹删除失败", level="warning", program="main")
|
||||
|
||||
delete_flag = split_frame_config.get('delete_video')
|
||||
if delete_flag:
|
||||
try:
|
||||
if os.path.exists(local_video_path):
|
||||
os.remove(local_video_path)
|
||||
logger.log(f"已删除本地视频文件: {local_video_path}", program="main")
|
||||
except OSError as e:
|
||||
logger.log(f"删除视频文件失败: {e}", level="warning", program="main")
|
||||
|
||||
# 5. 转字幕
|
||||
composite_path = os.path.join("./split", video_name)
|
||||
os.makedirs(composite_path, exist_ok=True)
|
||||
if not os.path.exists(composite_path):
|
||||
raise FileNotFoundError(f"视频{video_name}的 composite 文件夹不存在: {composite_path}")
|
||||
|
||||
try:
|
||||
_t = time.time()
|
||||
score, conf, srt_path, contents, best_grids, audio_frame, h = processsrt(
|
||||
composite_path, "ocr_result.jsonl", f"{video_name}.ocr.srt", interval_seconds, logger)
|
||||
_getsrt_elapsed = time.time() - _t
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"视频{video_name}字幕生成失败: {e}") from e
|
||||
|
||||
# 7. 记录元数据
|
||||
if score is None:
|
||||
filter_reason = srt_path if isinstance(srt_path, str) else "未知原因"
|
||||
logger.step(video_name, "已过滤", str(filter_reason))
|
||||
logger.log(f"视频{video_name}被过滤，原因：{filter_reason}", level="warning", program="main")
|
||||
_save_filter_metadata(
|
||||
cos_source, video_name, filter_reason=filter_reason,
|
||||
download_flag=download_flag, interval_seconds=interval_seconds,
|
||||
duration=duration, subtitle_region=subtitle_region,
|
||||
tile_width=tile_width, tile_height=tile_height, compositenum=compositenum,
|
||||
audio_text=audio_text, sample_indices=sample_indices,
|
||||
watermark=watermark, candidate_regions=candidate_regions,
|
||||
contents=contents, best_grids=best_grids, audio_frame=audio_frame,
|
||||
h=h, score=score, conf=conf
|
||||
)
|
||||
logger.end_program("main")
|
||||
logger.save_logs()
|
||||
return False
|
||||
|
||||
_save_success_metadata(
|
||||
cos_source, video_name, srt_path=srt_path,
|
||||
download_flag=download_flag, interval_seconds=interval_seconds,
|
||||
duration=duration, subtitle_region=subtitle_region,
|
||||
tile_width=tile_width, tile_height=tile_height, compositenum=compositenum,
|
||||
audio_text=audio_text, sample_indices=sample_indices,
|
||||
watermark=watermark, candidate_regions=candidate_regions,
|
||||
contents=contents, best_grids=best_grids, audio_frame=audio_frame,
|
||||
h=h, score=score, conf=conf
|
||||
)
|
||||
logger.step(video_name, "字幕完成",
|
||||
f"质量分{score:.0f}, {_getsrt_elapsed:.0f}s, 等待上传")
|
||||
logger.log(f"视频{video_name}处理完成", program="main")
|
||||
logger.end_program("main")
|
||||
logger.save_logs()
|
||||
return True
|
||||
|
||||
|
||||
def _build_metadata_common(download_flag, interval_seconds, duration, subtitle_region,
|
||||
tile_width, tile_height, compositenum, audio_text,
|
||||
sample_indices, watermark, candidate_regions,
|
||||
contents, best_grids, audio_frame, h):
|
||||
"""构建元数据的公共部分，避免代码重复"""
|
||||
best_grid_data = best_grids.get("composite") if best_grids else None
|
||||
return {
|
||||
"download_phase": {
|
||||
"download_flag": download_flag
|
||||
},
|
||||
"frame_extract": {
|
||||
"interval": interval_seconds,
|
||||
"duration": duration,
|
||||
"success": True
|
||||
},
|
||||
"subtitle_localize": {
|
||||
"sample_indices": sample_indices,
|
||||
"watermark": watermark,
|
||||
"candidate_regions": candidate_regions,
|
||||
"subtitle_position": {
|
||||
"dilation": "65",
|
||||
"final_region": subtitle_region
|
||||
},
|
||||
"audio_extract": audio_text
|
||||
},
|
||||
"composite_build": {
|
||||
"BORDER_SIZE": 40,
|
||||
"tile_width": tile_width,
|
||||
"tile_height": tile_height,
|
||||
"composite_num": compositenum,
|
||||
"success": True
|
||||
},
|
||||
"ocr_process": {
|
||||
"ocr_result": contents,
|
||||
"best_grid": {
|
||||
"composite": {
|
||||
"grid_position": best_grid_data[0] if best_grid_data else None,
|
||||
"grid_width": best_grid_data[1] if best_grid_data else None,
|
||||
"grid_height": best_grid_data[2] if best_grid_data else None,
|
||||
"grid_top_bound": best_grid_data[3] if best_grid_data else None,
|
||||
"grid_left_bound": best_grid_data[4] if best_grid_data else None
|
||||
}
|
||||
}
|
||||
},
|
||||
"frame_processing": {
|
||||
"audio_frame": audio_frame,
|
||||
"average_height": h,
|
||||
"merge_params": {
|
||||
"confidence_bottom_line": 75,
|
||||
"merge_srt_similarity_bottom_line": 0.75,
|
||||
"audio_similarity_bottom_line": 0.8
|
||||
}
|
||||
},
|
||||
"custom_score": None,
|
||||
"ocr_confidence": None
|
||||
}
|
||||
|
||||
|
||||
def _save_success_metadata(cos_source, video_name, srt_path, **kwargs):
|
||||
"""保存成功的元数据到 info 目录"""
|
||||
os.makedirs("./info", exist_ok=True)
|
||||
score = kwargs.pop("score")
|
||||
conf = kwargs.pop("conf")
|
||||
common = _build_metadata_common(**kwargs)
|
||||
common["custom_score"] = score
|
||||
common["ocr_confidence"] = conf
|
||||
common["srt_output"] = srt_path
|
||||
|
||||
metadata = {"main_process": {"input_video": cos_source}}
|
||||
metadata["main_process"].update(common)
|
||||
|
||||
metadata_js = os.path.join("./info", f"{video_name}.json")
|
||||
with open(metadata_js, 'w', encoding='utf-8') as f:
|
||||
json.dump(metadata, f, ensure_ascii=False, indent=2)
|
||||
|
||||
|
||||
def _save_filter_metadata(cos_source, video_name, filter_reason, **kwargs):
|
||||
"""保存被过滤视频的元数据到 info_error 目录"""
|
||||
os.makedirs("./info_error", exist_ok=True)
|
||||
score = kwargs.pop("score")
|
||||
conf = kwargs.pop("conf")
|
||||
common = _build_metadata_common(**kwargs)
|
||||
common["custom_score"] = score
|
||||
common["ocr_confidence"] = conf
|
||||
common["error"] = filter_reason
|
||||
|
||||
metadata = {"main_process": {"input_video": cos_source}}
|
||||
metadata["main_process"].update(common)
|
||||
|
||||
metadata_js = os.path.join("./info_error", f"{video_name}.json")
|
||||
with open(metadata_js, 'w', encoding='utf-8') as f:
|
||||
json.dump(metadata, f, ensure_ascii=False, indent=2)
|
||||
|
||||
|
||||
def _save_error_metadata(cos_source, video_name, error_reason, download_flag, interval_seconds):
|
||||
"""保存处理异常视频的元数据到 info_error 目录"""
|
||||
os.makedirs("./info_error", exist_ok=True)
|
||||
metadata = {
|
||||
"main_process": {
|
||||
"input_video": cos_source,
|
||||
"download_phase": {"download_flag": download_flag},
|
||||
"frame_extract": {"interval": interval_seconds},
|
||||
"error": error_reason
|
||||
}
|
||||
}
|
||||
metadata_js = os.path.join("./info_error", f"{video_name}.json")
|
||||
with open(metadata_js, 'w', encoding='utf-8') as f:
|
||||
json.dump(metadata, f, ensure_ascii=False, indent=2)
|
||||
227
tool/whisper_manager.py
Normal file
227
tool/whisper_manager.py
Normal file
@@ -0,0 +1,227 @@
|
||||
import os
|
||||
import hashlib
|
||||
import subprocess
|
||||
import tempfile
|
||||
import threading
|
||||
from collections import OrderedDict
|
||||
from typing import Optional
|
||||
|
||||
from pydub import AudioSegment
|
||||
|
||||
from config.config_utils import get_ffmpeg, get_ffprobe
|
||||
|
||||
# 让 pydub 也用项目自带的 ffmpeg/ffprobe(显式指定，不依赖 PATH)
|
||||
AudioSegment.converter = get_ffmpeg()
|
||||
AudioSegment.ffmpeg = get_ffmpeg()
|
||||
AudioSegment.ffprobe = get_ffprobe()
|
||||
|
||||
_MAX_CACHE_SIZE = 512
|
||||
|
||||
|
||||
class WhisperManager:
|
||||
_instance = None
|
||||
_lock = threading.Lock()
|
||||
|
||||
def __new__(cls):
|
||||
if cls._instance is None:
|
||||
with cls._lock:
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
cls._instance._initialized = False
|
||||
return cls._instance
|
||||
|
||||
def __init__(self):
|
||||
if not self._initialized:
|
||||
self.model = None
|
||||
self.model_name = None
|
||||
self._load_lock = threading.Lock() # 仅保护模型加载，推理不加锁
|
||||
self._cache = OrderedDict()
|
||||
self._cache_lock = threading.Lock()
|
||||
self._initialized = True
|
||||
|
||||
@staticmethod
|
||||
def _cuda_available() -> bool:
|
||||
"""nvidia-smi 可用即认为有可用 GPU。"""
|
||||
try:
|
||||
r = subprocess.run(["nvidia-smi"], capture_output=True, timeout=10)
|
||||
return r.returncode == 0
|
||||
except (FileNotFoundError, OSError, subprocess.TimeoutExpired):
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _gpu_count() -> int:
|
||||
"""数 GPU 张数（nvidia-smi -L 行数）。"""
|
||||
try:
|
||||
r = subprocess.run(["nvidia-smi", "-L"], capture_output=True, text=True, timeout=10)
|
||||
if r.returncode == 0:
|
||||
return sum(1 for ln in r.stdout.splitlines() if ln.strip().startswith("GPU "))
|
||||
except (FileNotFoundError, OSError, subprocess.TimeoutExpired):
|
||||
pass
|
||||
return 0
|
||||
|
||||
def _resolve_device(self):
|
||||
"""决定 whisper 跑在 cpu 还是 gpu。读 config.audio.device：auto/cuda/gpu/cpu。
|
||||
device_index=auto/last 时用【最后一张卡】，与抽帧(用前 N-1 张)错开。
|
||||
返回 (device, compute_type, device_index)。"""
|
||||
try:
|
||||
from config.config_utils import config
|
||||
acfg = config.get_audio_config() or {}
|
||||
except Exception:
|
||||
acfg = {}
|
||||
dev = str(acfg.get("device", "auto")).lower()
|
||||
idx_cfg = acfg.get("device_index", "auto")
|
||||
ctype = str(acfg.get("compute_type", "float16"))
|
||||
if dev in ("cuda", "gpu") or (dev == "auto" and self._cuda_available()):
|
||||
if str(idx_cfg).lower() in ("auto", "last"):
|
||||
idx = max(0, self._gpu_count() - 1) # 用最后一张
|
||||
else:
|
||||
try:
|
||||
idx = int(idx_cfg)
|
||||
except (TypeError, ValueError):
|
||||
idx = 0
|
||||
return "cuda", ctype, idx
|
||||
return "cpu", "int8", 0
|
||||
|
||||
def load_model(self, model_name: str = "base"):
|
||||
"""加载 faster-whisper 模型：有 GPU 用 cuda(float16)，否则 CPU(int8)；GPU 失败自动回退 CPU。"""
|
||||
with self._load_lock:
|
||||
if self.model is None or self.model_name != model_name:
|
||||
from faster_whisper import WhisperModel
|
||||
device, ctype, idx = self._resolve_device()
|
||||
# CPU 推理限制内部线程数：防止进程池下每个 worker 的 whisper 各抓一堆核 → 线程爆炸。
|
||||
# whisper 只识别首尾各5秒,2 线程足够;可用 config.audio.cpu_threads 调整。
|
||||
try:
|
||||
from config.config_utils import config as _cfg
|
||||
cpu_threads = int((_cfg.get_audio_config() or {}).get("cpu_threads", 2) or 2)
|
||||
except Exception:
|
||||
cpu_threads = 2
|
||||
cpu_threads = max(1, cpu_threads)
|
||||
try:
|
||||
print(f"[WhisperManager] 加载 faster-whisper '{model_name}' (device={device}, {ctype})...")
|
||||
if device == "cuda":
|
||||
self.model = WhisperModel(model_name, device="cuda",
|
||||
device_index=idx, compute_type=ctype)
|
||||
else:
|
||||
self.model = WhisperModel(model_name, device="cpu",
|
||||
compute_type="int8", cpu_threads=cpu_threads)
|
||||
except Exception as e:
|
||||
print(f"[WhisperManager] GPU 加载失败({type(e).__name__}: {e})，回退 CPU(int8, cpu_threads={cpu_threads})")
|
||||
self.model = WhisperModel(model_name, device="cpu",
|
||||
compute_type="int8", cpu_threads=cpu_threads)
|
||||
self.model_name = model_name
|
||||
print(f"[WhisperManager] 模型 '{model_name}' 加载完成。")
|
||||
return self.model
|
||||
|
||||
def extract_audio_from_video(
|
||||
self,
|
||||
video_path: str,
|
||||
output_audio_path: str,
|
||||
logger: Optional[object] = None
|
||||
) -> Optional[float]:
|
||||
if not os.path.exists(video_path):
|
||||
if logger:
|
||||
logger.log(f"视频文件不存在: {video_path}", level="error")
|
||||
return None
|
||||
|
||||
cmd = [
|
||||
get_ffmpeg(), "-loglevel", "error",
|
||||
"-i", video_path,
|
||||
"-vn", "-acodec", "pcm_s16le",
|
||||
"-ar", "16000", "-ac", "1",
|
||||
"-y", output_audio_path
|
||||
]
|
||||
try:
|
||||
subprocess.run(cmd, check=True, capture_output=True, text=True)
|
||||
except subprocess.CalledProcessError as e:
|
||||
if logger:
|
||||
logger.log(f"ffmpeg 提取音频失败: {e.stderr}", level="error")
|
||||
return None
|
||||
|
||||
try:
|
||||
audio = AudioSegment.from_file(output_audio_path)
|
||||
duration = len(audio) / 1000.0
|
||||
if logger:
|
||||
logger.log(f"音频提取成功，时长: {duration:.2f} 秒")
|
||||
return duration
|
||||
except Exception as e:
|
||||
if logger:
|
||||
logger.log(f"获取音频时长失败: {e}", level="error")
|
||||
return None
|
||||
|
||||
def extract_audio_text_with_whisper(
|
||||
self,
|
||||
audio_path: str,
|
||||
start_time: float,
|
||||
end_time: float,
|
||||
logger: Optional[object] = None,
|
||||
use_cache: bool = True
|
||||
) -> str:
|
||||
if self.model is None:
|
||||
raise RuntimeError("Whisper 模型未加载，请先调用 load_model()")
|
||||
|
||||
if not os.path.exists(audio_path):
|
||||
if logger:
|
||||
logger.log(f"音频文件不存在: {audio_path}", level="error")
|
||||
return ""
|
||||
|
||||
cache_key = None
|
||||
if use_cache:
|
||||
mtime = os.path.getmtime(audio_path)
|
||||
cache_key = hashlib.md5(
|
||||
f"{audio_path}_{mtime}_{start_time}_{end_time}".encode()
|
||||
).hexdigest()
|
||||
with self._cache_lock:
|
||||
if cache_key in self._cache:
|
||||
self._cache.move_to_end(cache_key)
|
||||
if logger:
|
||||
logger.log("使用缓存识别结果")
|
||||
return self._cache[cache_key]
|
||||
|
||||
try:
|
||||
full_audio = AudioSegment.from_file(audio_path)
|
||||
start_ms = int(start_time * 1000)
|
||||
end_ms = int(end_time * 1000) + 500
|
||||
segment = full_audio[start_ms:end_ms]
|
||||
|
||||
if len(segment) == 0:
|
||||
if logger:
|
||||
logger.log("音频片段为空，跳过识别")
|
||||
return ""
|
||||
|
||||
segment = segment.set_frame_rate(16000).set_channels(1).set_sample_width(2)
|
||||
|
||||
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
|
||||
temp_path = tmp.name
|
||||
segment.export(temp_path, format="wav")
|
||||
|
||||
try:
|
||||
# faster-whisper CPU 模式线程安全，无需加锁，多线程并发推理
|
||||
segments, _ = self.model.transcribe(
|
||||
temp_path, language="zh", beam_size=5
|
||||
)
|
||||
text = "".join(seg.text for seg in segments).strip()
|
||||
finally:
|
||||
try:
|
||||
os.unlink(temp_path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
if use_cache and cache_key:
|
||||
with self._cache_lock:
|
||||
self._cache[cache_key] = text
|
||||
self._cache.move_to_end(cache_key)
|
||||
while len(self._cache) > _MAX_CACHE_SIZE:
|
||||
self._cache.popitem(last=False)
|
||||
|
||||
if logger:
|
||||
logger.log(f"音频片段识别结果: {text}")
|
||||
|
||||
return text
|
||||
|
||||
except Exception as e:
|
||||
if logger:
|
||||
logger.log(f"音频识别失败: {e}", level="error")
|
||||
return ""
|
||||
|
||||
|
||||
whisper_manager = WhisperManager()
|
||||
Reference in New Issue
Block a user