508 lines
17 KiB
Python
508 lines
17 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Subtitle Safe Area QC · 字幕安全区域质检器
|
||
=============================================
|
||
自动检查字幕有没有遮脸、遮身体表演、贴边、看不清。
|
||
|
||
依赖：
|
||
pip install opencv-python pillow numpy
|
||
# 可选：pip install face-recognition dlib # 人脸识别
|
||
# 可选：pip install ultralytics # YOLO 人体检测
|
||
|
||
用法：
|
||
# 检查单张帧
|
||
python subtitle-safe-area-qc.py --frame frame.png --subtitle-box x,y,w,h --output qc-report.json
|
||
|
||
# 批量检查视频帧
|
||
python subtitle-safe-area-qc.py --video input.mp4 --subtitle-ass subtitles.ass --output qc-report.json
|
||
|
||
# 交互式标注模式（手动框选人脸/身体区域）
|
||
python subtitle-safe-area-qc.py --frame frame.png --interactive
|
||
|
||
# 生成质检报告（包含问题帧截图）
|
||
python subtitle-safe-area-qc.py --video input.mp4 --subtitle-ass subtitles.ass --generate-report --output-dir ./qc-report/
|
||
|
||
输出 JSON 格式：
|
||
{
|
||
"frame": "frame_00123.png",
|
||
"timestamp": 12.34,
|
||
"issues": [
|
||
{"type": "face_occlusion", "severity": "high", "bbox": [x,y,w,h]},
|
||
{"type": "body_occlusion", "severity": "medium", "bbox": [x,y,w,h]},
|
||
{"type": "edge_too_close", "severity": "low", "distance": 5},
|
||
{"type": "low_contrast", "severity": "medium", "contrast_ratio": 2.1}
|
||
],
|
||
"safe": false
|
||
}
|
||
|
||
路径：
|
||
video-ai-system/engines/subtitle-pipeline/qc-tools/subtitle-safe-area-qc.py
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import os
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
try:
|
||
import cv2
|
||
import numpy as np
|
||
from PIL import Image
|
||
except ImportError as e:
|
||
print(f"[ERROR] 缺少依赖：{e}")
|
||
print("请先安装：pip install opencv-python pillow numpy")
|
||
sys.exit(1)
|
||
|
||
|
||
# 尝试导入可选依赖
|
||
try:
|
||
import face_recognition
|
||
FACE_RECOGNITION_AVAILABLE = True
|
||
except ImportError:
|
||
FACE_RECOGNITION_AVAILABLE = False
|
||
print("[WARN] 未安装 face-recognition，将使用简化人脸检测")
|
||
|
||
try:
|
||
from ultralytics import YOLO
|
||
YOLO_AVAILABLE = True
|
||
except ImportError:
|
||
YOLO_AVAILABLE = False
|
||
print("[WARN] 未安装 ultralytics，将使用简化身体检测")
|
||
|
||
|
||
def detect_faces(image: np.ndarray) -> list:
|
||
"""
|
||
检测人脸区域
|
||
|
||
:param image: 图片数组
|
||
:return: 人脸边界框列表 [[x,y,w,h], ...]
|
||
"""
|
||
if FACE_RECOGNITION_AVAILABLE:
|
||
# 使用 face_recognition 库（基于 dlib）
|
||
rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||
face_locations = face_recognition.face_locations(rgb_image)
|
||
|
||
# 转换为 [x,y,w,h] 格式
|
||
boxes = []
|
||
for top, right, bottom, left in face_locations:
|
||
x = left
|
||
y = top
|
||
w = right - left
|
||
h = bottom - top
|
||
boxes.append([x, y, w, h])
|
||
|
||
return boxes
|
||
else:
|
||
# 简化方案：使用 OpenCV Haar Cascade
|
||
cascade_path = cv2.data.haarcascades + "haarcascade_frontalface_default.xml"
|
||
face_cascade = cv2.CascadeClassifier(cascade_path)
|
||
|
||
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
||
faces = face_cascade.detectMultiScale(gray, 1.1, 4)
|
||
|
||
return faces.tolist() if len(faces) > 0 else []
|
||
|
||
|
||
def detect_body(image: np.ndarray) -> list:
|
||
"""
|
||
检测身体区域
|
||
|
||
:param image: 图片数组
|
||
:return: 身体边界框列表 [[x,y,w,h], ...]
|
||
"""
|
||
if YOLO_AVAILABLE:
|
||
# 使用 YOLO 检测人体
|
||
model = YOLO("yolov8n.pt") # 自动下载
|
||
results = model(image)
|
||
|
||
boxes = []
|
||
for result in results:
|
||
for box in result.boxes:
|
||
cls = int(box.cls[0])
|
||
# COCO 数据集中，person 的类别 ID 是 0
|
||
if cls == 0:
|
||
x1, y1, x2, y2 = box.xyxy[0].tolist()
|
||
boxes.append([int(x1), int(y1), int(x2-x1), int(y2-y1)])
|
||
|
||
return boxes
|
||
else:
|
||
# 简化方案：使用背景减除或轮廓检测
|
||
print("[WARN] 未安装 YOLO，使用简化身体检测（可能不准确）")
|
||
|
||
# 策略：检测图像中的大轮廓（假设身体是大区域）
|
||
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
||
_, binary = cv2.threshold(gray, 128, 255, cv2.THRESH_BINARY_INV)
|
||
|
||
contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||
|
||
boxes = []
|
||
for cnt in contours:
|
||
area = cv2.contourArea(cnt)
|
||
if area > image.shape[0] * image.shape[1] * 0.1: # 面积 > 10%
|
||
x, y, w, h = cv2.boundingRect(cnt)
|
||
boxes.append([x, y, w, h])
|
||
|
||
return boxes
|
||
|
||
|
||
def detect_subtitle_box(image: np.ndarray) -> list:
|
||
"""
|
||
检测字幕框位置
|
||
|
||
:param image: 图片数组
|
||
:return: 字幕边界框 [x,y,w,h]
|
||
"""
|
||
# 策略：检测底部区域的白色文本
|
||
height, width = image.shape[:2]
|
||
|
||
# 裁剪底部区域
|
||
bottom_region = image[int(height * 0.85):, :]
|
||
|
||
# 转为灰度图
|
||
gray = cv2.cvtColor(bottom_region, cv2.COLOR_BGR2GRAY)
|
||
|
||
# 二值化（检测白色文本）
|
||
_, binary = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY)
|
||
|
||
# 查找轮廓
|
||
contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||
|
||
if not contours:
|
||
return [0, int(height * 0.85), width, int(height * 0.12)] # 默认底部区域
|
||
|
||
# 收集所有轮廓的边界框
|
||
all_x = []
|
||
all_y = []
|
||
all_w = []
|
||
all_h = []
|
||
|
||
for cnt in contours:
|
||
x, y, w, h = cv2.boundingRect(cnt)
|
||
all_x.append(x)
|
||
all_y.append(y)
|
||
all_w.append(w)
|
||
all_h.append(h)
|
||
|
||
# 合并所有轮廓的边界框
|
||
# 先在 bottom_region 坐标系里算完，最后再加偏移
|
||
x_bottom = min(all_x)
|
||
y_bottom = min(all_y)
|
||
max_x_bottom = max([all_x[i] + all_w[i] for i in range(len(all_x))])
|
||
max_y_bottom = max([all_y[i] + all_h[i] for i in range(len(all_y))])
|
||
w = max_x_bottom - x_bottom
|
||
h = max_y_bottom - y_bottom
|
||
|
||
# 最后统一加偏移
|
||
x = x_bottom
|
||
y = y_bottom + int(height * 0.85)
|
||
# w, h 不需要加偏移（它们是在bottom_region里的尺寸）
|
||
|
||
return [x, y, w, h]
|
||
|
||
|
||
def calculate_iou(box1: list, box2: list) -> float:
|
||
"""
|
||
计算两个边界框的 IOU（交并比）
|
||
|
||
:param box1: 边界框1 [x,y,w,h]
|
||
:param box2: 边界框2 [x,y,w,h]
|
||
:return: IOU 值（0-1）
|
||
"""
|
||
x1 = max(box1[0], box2[0])
|
||
y1 = max(box1[1], box2[1])
|
||
x2 = min(box1[0] + box1[2], box2[0] + box2[2])
|
||
y2 = min(box1[1] + box1[3], box2[1] + box2[3])
|
||
|
||
if x2 < x1 or y2 < y1:
|
||
return 0.0
|
||
|
||
intersection = (x2 - x1) * (y2 - y1)
|
||
area1 = box1[2] * box1[3]
|
||
area2 = box2[2] * box2[3]
|
||
union = area1 + area2 - intersection
|
||
|
||
return intersection / union if union > 0 else 0.0
|
||
|
||
|
||
def check_subtitle_safe_area(image_path: str, subtitle_box: list = None, interactive: bool = False) -> dict:
|
||
"""
|
||
检查字幕安全区域
|
||
|
||
:param image_path: 图片路径
|
||
:param subtitle_box: 字幕框 [x,y,w,h]（如果为 None，自动检测）
|
||
:param interactive: 是否交互式标注
|
||
:return: 质检报告字典
|
||
"""
|
||
# 读取图片
|
||
img = cv2.imread(image_path)
|
||
if img is None:
|
||
print(f"[ERROR] 无法读取图片：{image_path}")
|
||
return {}
|
||
|
||
height, width = img.shape[:2]
|
||
print(f"[INFO] 检查图片：{image_path} ({width}x{height})")
|
||
|
||
# 检测字幕框
|
||
if subtitle_box is None:
|
||
if interactive:
|
||
print("[INFO] 交互式标注模式：请手动框选字幕区域")
|
||
roi = cv2.selectROI("Select Subtitle Region", img, showCrosshair=True)
|
||
cv2.destroyAllWindows()
|
||
subtitle_box = [int(roi[0]), int(roi[1]), int(roi[2]), int(roi[3])]
|
||
else:
|
||
subtitle_box = detect_subtitle_box(img)
|
||
|
||
print(f"[INFO] 字幕框：x={subtitle_box[0]}, y={subtitle_box[1]}, w={subtitle_box[2]}, h={subtitle_box[3]}")
|
||
|
||
# 检测人脸
|
||
print("[INFO] 检测人脸...")
|
||
face_boxes = detect_faces(img)
|
||
print(f"[INFO] 找到 {len(face_boxes)} 个人脸")
|
||
|
||
# 检测身体
|
||
print("[INFO] 检测身体...")
|
||
body_boxes = detect_body(img)
|
||
print(f"[INFO] 找到 {len(body_boxes)} 个身体区域")
|
||
|
||
# 检查问题
|
||
issues = []
|
||
|
||
# 1. 检查是否遮挡脸部
|
||
for i, face_box in enumerate(face_boxes):
|
||
iou = calculate_iou(subtitle_box, face_box)
|
||
if iou > 0.1: # IOU > 10% 认为有遮挡
|
||
severity = "high" if iou > 0.5 else "medium"
|
||
issues.append({
|
||
"type": "face_occlusion",
|
||
"severity": severity,
|
||
"bbox": face_box,
|
||
"iou": round(iou, 2),
|
||
"message": f"字幕遮挡脸部（IOU={iou:.2f}）"
|
||
})
|
||
print(f"[ISSUE] 字幕遮挡脸部（IOU={iou:.2f}）")
|
||
|
||
# 2. 检查是否遮挡身体
|
||
for i, body_box in enumerate(body_boxes):
|
||
iou = calculate_iou(subtitle_box, body_box)
|
||
if iou > 0.2: # IOU > 20% 认为有遮挡
|
||
severity = "high" if iou > 0.6 else "medium"
|
||
issues.append({
|
||
"type": "body_occlusion",
|
||
"severity": severity,
|
||
"bbox": body_box,
|
||
"iou": round(iou, 2),
|
||
"message": f"字幕遮挡身体（IOU={iou:.2f}）"
|
||
})
|
||
print(f"[ISSUE] 字幕遮挡身体（IOU={iou:.2f}）")
|
||
|
||
# 3. 检查是否贴边
|
||
edge_threshold = 20 # 像素
|
||
if subtitle_box[0] < edge_threshold:
|
||
issues.append({
|
||
"type": "edge_too_close",
|
||
"severity": "low",
|
||
"distance": subtitle_box[0],
|
||
"message": f"字幕距离左边缘太近（{subtitle_box[0]}px）"
|
||
})
|
||
print(f"[ISSUE] 字幕距离左边缘太近（{subtitle_box[0]}px）")
|
||
|
||
if (subtitle_box[0] + subtitle_box[2]) > (width - edge_threshold):
|
||
distance = (subtitle_box[0] + subtitle_box[2]) - width
|
||
issues.append({
|
||
"type": "edge_too_close",
|
||
"severity": "low",
|
||
"distance": abs(distance),
|
||
"message": f"字幕距离右边缘太近（{abs(distance)}px）"
|
||
})
|
||
print(f"[ISSUE] 字幕距离右边缘太近（{abs(distance)}px）")
|
||
|
||
# 4. 检查对比度（字幕是否看不清）
|
||
# 简化方案：计算字幕区域和背景的平均亮度差
|
||
subtitle_region = img[
|
||
subtitle_box[1]:subtitle_box[1]+subtitle_box[3],
|
||
subtitle_box[0]:subtitle_box[0]+subtitle_box[2]
|
||
]
|
||
|
||
if subtitle_region.size > 0:
|
||
# 计算字幕区域的平均亮度
|
||
subtitle_gray = cv2.cvtColor(subtitle_region, cv2.COLOR_BGR2GRAY)
|
||
subtitle_brightness = np.mean(subtitle_gray)
|
||
|
||
# 计算背景区域的平均亮度（字幕框上方的区域）
|
||
background_region = img[
|
||
max(0, subtitle_box[1]-subtitle_box[3]):subtitle_box[1],
|
||
subtitle_box[0]:subtitle_box[0]+subtitle_box[2]
|
||
]
|
||
|
||
if background_region.size > 0:
|
||
background_gray = cv2.cvtColor(background_region, cv2.COLOR_BGR2GRAY)
|
||
background_brightness = np.mean(background_gray)
|
||
|
||
# 计算对比度（亮度差）
|
||
contrast = abs(subtitle_brightness - background_brightness)
|
||
contrast_ratio = contrast / max(subtitle_brightness, background_brightness)
|
||
|
||
if contrast_ratio < 0.3: # 对比度 < 30% 认为看不清
|
||
severity = "high" if contrast_ratio < 0.1 else "medium"
|
||
issues.append({
|
||
"type": "low_contrast",
|
||
"severity": severity,
|
||
"contrast_ratio": round(contrast_ratio, 2),
|
||
"message": f"字幕对比度太低（{contrast_ratio:.2f}），可能看不清"
|
||
})
|
||
print(f"[ISSUE] 字幕对比度太低（{contrast_ratio:.2f}），可能看不清")
|
||
|
||
# 生成报告
|
||
report = {
|
||
"image_path": image_path,
|
||
"video_width": width,
|
||
"video_height": height,
|
||
"subtitle_box": subtitle_box,
|
||
"face_boxes": face_boxes,
|
||
"body_boxes": body_boxes,
|
||
"issues": issues,
|
||
"safe": len(issues) == 0,
|
||
"issue_count": len(issues)
|
||
}
|
||
|
||
if report["safe"]:
|
||
print(f"[OK] 字幕安全区域检查通过")
|
||
else:
|
||
print(f"[WARN] 发现 {len(issues)} 个问题")
|
||
|
||
return report
|
||
|
||
|
||
def batch_check_video(video_path: str, ass_path: str, output_dir: str, sample_interval: float = 1.0):
|
||
"""
|
||
批量检查视频帧
|
||
|
||
:param video_path: 视频文件路径
|
||
:param ass_path: ASS 字幕文件路径
|
||
:param output_dir: 输出目录
|
||
:param sample_interval: 采样间隔（秒）
|
||
"""
|
||
# 打开视频
|
||
cap = cv2.VideoCapture(video_path)
|
||
if not cap.isOpened():
|
||
print(f"[ERROR] 无法打开视频：{video_path}")
|
||
return
|
||
|
||
fps = cap.get(cv2.CAP_PROP_FPS)
|
||
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||
|
||
print(f"[INFO] 视频信息：{total_frames} 帧，{fps} FPS")
|
||
|
||
# 创建输出目录
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
frames_dir = os.path.join(output_dir, "problem-frames/")
|
||
os.makedirs(frames_dir, exist_ok=True)
|
||
|
||
# 读取 ASS 字幕文件（获取字幕时间点）
|
||
# 简化方案：假设字幕在底部，直接检测
|
||
# TODO: 解析 ASS 文件，获取准确的字幕时间点
|
||
|
||
reports = []
|
||
frame_count = 0
|
||
|
||
while True:
|
||
ret, frame = cap.read()
|
||
if not ret:
|
||
break
|
||
|
||
timestamp = frame_count / fps
|
||
|
||
# 按采样间隔检查
|
||
if timestamp % sample_interval < (1.0 / fps):
|
||
# 保存帧为临时文件
|
||
temp_frame_path = os.path.join(output_dir, f"temp_frame_{frame_count:06d}.png")
|
||
cv2.imwrite(temp_frame_path, frame)
|
||
|
||
# 检查字幕安全区域
|
||
report = check_subtitle_safe_area(temp_frame_path)
|
||
|
||
if report and not report["safe"]:
|
||
# 保存问题帧
|
||
problem_frame_path = os.path.join(frames_dir, f"problem_{frame_count:06d}.png")
|
||
cv2.imwrite(problem_frame_path, frame)
|
||
|
||
report["timestamp"] = timestamp
|
||
report["problem_frame"] = problem_frame_path
|
||
reports.append(report)
|
||
|
||
print(f"[WARN] 发现问题的帧：{timestamp:.2f}s")
|
||
|
||
# 删除临时文件
|
||
os.remove(temp_frame_path)
|
||
|
||
frame_count += 1
|
||
|
||
cap.release()
|
||
|
||
# 保存报告
|
||
report_path = os.path.join(output_dir, "qc-report.json")
|
||
with open(report_path, "w", encoding="utf-8") as f:
|
||
json.dump(reports, f, ensure_ascii=False, indent=2)
|
||
|
||
print(f"\n[OK] 批量检查完成")
|
||
print(f"[INFO] 共检查 {frame_count} 帧，发现 {len(reports)} 个有问题帧")
|
||
print(f"[INFO] 报告已保存：{report_path}")
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(
|
||
description="Subtitle Safe Area QC · 字幕安全区域质检器",
|
||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||
epilog="""
|
||
示例：
|
||
python subtitle-safe-area-qc.py --frame frame.png --subtitle-box 100,1750,880,120 --output qc-report.json
|
||
python subtitle-safe-area-qc.py --video input.mp4 --subtitle-ass subtitles.ass --output qc-report.json
|
||
python subtitle-safe-area-qc.py --frame frame.png --interactive
|
||
"""
|
||
)
|
||
parser.add_argument("--frame", help="单张帧图片路径")
|
||
parser.add_argument("--video", help="视频文件路径（批量检查）")
|
||
parser.add_argument("--subtitle-box", help="字幕框坐标 x,y,w,h（逗号分隔）")
|
||
parser.add_argument("--subtitle-ass", help="ASS 字幕文件路径（用于批量检查）")
|
||
parser.add_argument("--output", help="输出报告文件路径")
|
||
parser.add_argument("--interactive", action="store_true", help="交互式标注模式")
|
||
parser.add_argument("--generate-report", action="store_true", help="生成质检报告（包含问题帧截图）")
|
||
parser.add_argument("--output-dir", help="输出目录（用于批量检查）")
|
||
parser.add_argument("--sample-interval", type=float, default=1.0, help="采样间隔（秒，用于批量检查）")
|
||
|
||
args = parser.parse_args()
|
||
|
||
if args.frame:
|
||
# 单张帧检查
|
||
subtitle_box = None
|
||
if args.subtitle_box:
|
||
subtitle_box = [int(x) for x in args.subtitle_box.split(",")]
|
||
|
||
report = check_subtitle_safe_area(args.frame, subtitle_box, args.interactive)
|
||
|
||
if not report:
|
||
sys.exit(1)
|
||
|
||
if args.output:
|
||
with open(args.output, "w", encoding="utf-8") as f:
|
||
json.dump(report, f, ensure_ascii=False, indent=2)
|
||
print(f"[OK] 质检报告已保存：{args.output}")
|
||
|
||
elif args.video:
|
||
# 批量检查
|
||
if not args.subtitle_ass:
|
||
print("[ERROR] 批量检查需要指定 --subtitle-ass")
|
||
sys.exit(1)
|
||
|
||
output_dir = args.output_dir or "./qc-report/"
|
||
batch_check_video(args.video, args.subtitle_ass, output_dir, args.sample_interval)
|
||
|
||
else:
|
||
print("[ERROR] 请指定 --frame 或 --video")
|
||
sys.exit(1)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|