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

472 lines
21 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

from 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)