OpenClaw Agent
e0a2b323d9
- obsidian-kb: Add OCR processors (easyocr, tesseract) - obsidian-kb: Fix subtitle detection with test files - obsidian-kb: Add cookie verification script - Add find-skills skill for discovering capabilities - Add video processing scripts - Update clawhub lock file
297 lines
9.5 KiB
Python
297 lines
9.5 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
字幕检测器测试脚本
|
||
|
||
验证：
|
||
1. 预设区域定义
|
||
2. 自动字幕区域检测
|
||
3. 缓存机制
|
||
"""
|
||
|
||
import sys
|
||
import os
|
||
import cv2
|
||
import numpy as np
|
||
|
||
# 添加项目路径
|
||
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||
sys.path.insert(0, project_root)
|
||
|
||
from src.extractors.subtitle_detector import SubtitleDetector
|
||
|
||
|
||
def create_test_frame_with_text(width=1280, height=720):
|
||
"""
|
||
创建包含模拟字幕的测试帧
|
||
|
||
Args:
|
||
width: 帧宽度
|
||
height: 帧高度
|
||
|
||
Returns:
|
||
numpy.ndarray: 测试帧
|
||
"""
|
||
# 创建黑色背景
|
||
frame = np.zeros((height, width, 3), dtype=np.uint8)
|
||
|
||
# 添加一些噪声背景
|
||
noise = np.random.randint(0, 50, (height, width, 3), dtype=np.uint8)
|
||
frame = cv2.add(frame, noise)
|
||
|
||
# 在底部添加白色文字（模拟B站字幕）
|
||
font = cv2.FONT_HERSHEY_SIMPLEX
|
||
font_scale = 1.0
|
||
thickness = 2
|
||
|
||
# 添加多行字幕
|
||
texts = [
|
||
"这是测试字幕内容",
|
||
"B站硬字幕提取测试",
|
||
"00:01:23 第一行字幕",
|
||
"00:01:25 第二行字幕"
|
||
]
|
||
|
||
y_start = int(height * 0.85)
|
||
line_height = 35
|
||
|
||
for i, text in enumerate(texts):
|
||
y = y_start + i * line_height
|
||
x = int(width * 0.15)
|
||
|
||
# 绘制文字阴影（黑色）
|
||
cv2.putText(frame, text, (x + 2, y + 2), font, font_scale, (0, 0, 0), thickness + 2)
|
||
# 绘制白色文字
|
||
cv2.putText(frame, text, (x, y), font, font_scale, (255, 255, 255), thickness)
|
||
|
||
return frame
|
||
|
||
|
||
def test_preset_regions():
|
||
"""测试1: 预设区域定义"""
|
||
print("\n" + "="*60)
|
||
print("测试1: 预设区域定义")
|
||
print("="*60)
|
||
|
||
detector = SubtitleDetector()
|
||
frame_width, frame_height = 1280, 720
|
||
|
||
# 测试所有预设区域
|
||
for region_name in [detector.REGION_BOTTOM, detector.REGION_CENTER, detector.REGION_TOP]:
|
||
try:
|
||
x, y, width, height = detector.get_preset_region(region_name, frame_height, frame_width)
|
||
print(f"\n✓ 区域 '{region_name}':")
|
||
print(f" 坐标: x={x}, y={y}")
|
||
print(f" 尺寸: width={width}, height={height}")
|
||
print(f" 相对位置: top={y/frame_height:.2f}, bottom={(y+height)/frame_height:.2f}")
|
||
except Exception as e:
|
||
print(f"✗ 区域 '{region_name}' 失败: {e}")
|
||
|
||
# 测试无效区域名称
|
||
try:
|
||
detector.get_preset_region("invalid", frame_height, frame_width)
|
||
print("\n✗ 错误: 应该抛出异常但没有")
|
||
except ValueError as e:
|
||
print(f"\n✓ 无效区域名称正确抛出异常: {e}")
|
||
|
||
print("\n测试1完成!")
|
||
|
||
|
||
def test_auto_detection():
|
||
"""测试2: 自动字幕区域检测"""
|
||
print("\n" + "="*60)
|
||
print("测试2: 自动字幕区域检测")
|
||
print("="*60)
|
||
|
||
detector = SubtitleDetector()
|
||
|
||
# 创建测试帧（带字幕）
|
||
print("\n创建测试帧...")
|
||
test_frame = create_test_frame_with_text()
|
||
height, width = test_frame.shape[:2]
|
||
print(f"帧尺寸: {width}x{height}")
|
||
|
||
# 保存测试帧用于可视化
|
||
test_frame_path = "/tmp/test_subtitle_frame.png"
|
||
cv2.imwrite(test_frame_path, test_frame)
|
||
print(f"✓ 测试帧已保存到: {test_frame_path}")
|
||
|
||
# 执行自动检测
|
||
print("\n执行自动检测...")
|
||
try:
|
||
x, y, w, h = detector.auto_detect_subtitle_region(test_frame, video_id=None, use_cache=False)
|
||
print(f"✓ 检测到字幕区域:")
|
||
print(f" 坐标: x={x}, y={y}")
|
||
print(f" 尺寸: width={w}, height={h}")
|
||
print(f" 相对位置: top={y/height:.2f}, bottom={(y+h)/height:.2f}")
|
||
|
||
# 在测试帧上绘制检测框
|
||
frame_with_box = test_frame.copy()
|
||
cv2.rectangle(frame_with_box, (x, y), (x+w, y+h), (0, 255, 0), 2)
|
||
|
||
# 保存结果
|
||
result_path = "/tmp/test_subtitle_detected.png"
|
||
cv2.imwrite(result_path, frame_with_box)
|
||
print(f"✓ 检测结果已保存到: {result_path}")
|
||
|
||
# 验证检测区域是否合理（应该在底部）
|
||
region_bottom = (y + h) / height
|
||
if region_bottom > 0.7: # 底部30%
|
||
print(f"✓ 检测区域位置合理（底部）")
|
||
else:
|
||
print(f"! 警告: 检测区域可能不在底部 (bottom={region_bottom:.2f})")
|
||
|
||
except Exception as e:
|
||
print(f"✗ 自动检测失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
|
||
print("\n测试2完成!")
|
||
|
||
|
||
def test_cache_mechanism():
|
||
"""测试3: 缓存机制"""
|
||
print("\n" + "="*60)
|
||
print("测试3: 缓存机制")
|
||
print("="*60)
|
||
|
||
detector = SubtitleDetector(cache_size=3) # 小缓存便于测试
|
||
|
||
# 创建测试帧
|
||
test_frame = create_test_frame_with_text()
|
||
|
||
# 测试1: 首次检测并缓存
|
||
print("\n1. 首次检测并缓存...")
|
||
video_id1 = "BV1xx411c7mD"
|
||
region1 = detector.auto_detect_subtitle_region(test_frame, video_id=video_id1, use_cache=True)
|
||
print(f"✓ 检测完成，区域: {region1}")
|
||
print(f"✓ 缓存大小: {detector.get_cache_size()}")
|
||
|
||
# 测试2: 再次检测同一视频（应使用缓存）
|
||
print("\n2. 再次检测同一视频（应使用缓存）...")
|
||
region2 = detector.auto_detect_subtitle_region(test_frame, video_id=video_id1, use_cache=True)
|
||
print(f"✓ 缓存结果: {region2}")
|
||
if region1 == region2:
|
||
print(f"✓ 缓存命中，结果一致")
|
||
else:
|
||
print(f"✗ 警告: 缓存结果不一致")
|
||
|
||
# 测试3: 检测不同视频
|
||
print("\n3. 检测不同视频...")
|
||
video_id2 = "BV1yy411c7mE"
|
||
region3 = detector.auto_detect_subtitle_region(test_frame, video_id=video_id2, use_cache=True)
|
||
print(f"✓ 新视频检测完成，区域: {region3}")
|
||
print(f"✓ 缓存大小: {detector.get_cache_size()}")
|
||
|
||
# 测试4: 查看所有缓存
|
||
print("\n4. 查看所有缓存...")
|
||
cached = detector.get_cached_regions()
|
||
print(f"✓ 缓存内容:")
|
||
for vid, reg in cached.items():
|
||
print(f" {vid}: {reg}")
|
||
|
||
# 测试5: 清空指定视频缓存
|
||
print("\n5. 清空指定视频缓存...")
|
||
detector.clear_cache(video_id=video_id1)
|
||
print(f"✓ 已删除 {video_id1}")
|
||
print(f"✓ 缓存大小: {detector.get_cache_size()}")
|
||
|
||
# 测试6: 清空所有缓存
|
||
print("\n6. 清空所有缓存...")
|
||
detector.clear_cache()
|
||
print(f"✓ 缓存已清空")
|
||
print(f"✓ 缓存大小: {detector.get_cache_size()}")
|
||
|
||
# 测试7: 缓存容量限制（LRU）
|
||
print("\n7. 测试缓存容量限制（LRU）...")
|
||
for i in range(5):
|
||
vid = f"BV{i:012d}"
|
||
detector.auto_detect_subtitle_region(test_frame, video_id=vid, use_cache=True)
|
||
print(f" 添加视频 {vid}，缓存大小: {detector.get_cache_size()}")
|
||
|
||
cached = detector.get_cached_regions()
|
||
print(f"✓ 缓存中有 {len(cached)} 个视频（应为最大容量3）")
|
||
print(f"✓ 保留的视频: {list(cached.keys())}")
|
||
|
||
print("\n测试3完成!")
|
||
|
||
|
||
def test_integration():
|
||
"""测试4: 集成测试"""
|
||
print("\n" + "="*60)
|
||
print("测试4: 集成测试")
|
||
print("="*60)
|
||
|
||
detector = SubtitleDetector()
|
||
test_frame = create_test_frame_with_text()
|
||
height, width = test_frame.shape[:2]
|
||
|
||
print("\n模拟B站视频字幕检测流程...")
|
||
|
||
# 步骤1: 尝试自动检测
|
||
print("\n步骤1: 自动检测字幕区域")
|
||
video_id = "BV1xx411c7mD"
|
||
region = detector.auto_detect_subtitle_region(test_frame, video_id=video_id)
|
||
print(f"✓ 自动检测结果: {region}")
|
||
|
||
# 步骤2: 后续帧使用缓存
|
||
print("\n步骤2: 后续帧使用缓存")
|
||
for i in range(3):
|
||
cached_region = detector.auto_detect_subtitle_region(test_frame, video_id=video_id)
|
||
if cached_region == region:
|
||
print(f"✓ 帧 {i+1}: 使用缓存，结果一致")
|
||
else:
|
||
print(f"✗ 帧 {i+1}: 缓存结果异常")
|
||
|
||
# 步骤3: 对比预设区域
|
||
print("\n步骤3: 对比预设区域")
|
||
bottom_region = detector.get_preset_region(detector.REGION_BOTTOM, height, width)
|
||
print(f" 检测区域: y={region[1]}, height={region[3]}")
|
||
print(f" 底部预设: y={bottom_region[1]}, height={bottom_region[3]}")
|
||
|
||
# 步骤4: 提取字幕区域
|
||
print("\n步骤4: 提取字幕区域用于OCR")
|
||
x, y, w, h = region
|
||
subtitle_roi = test_frame[y:y+h, x:x+w]
|
||
subtitle_path = "/tmp/subtitle_roi.png"
|
||
cv2.imwrite(subtitle_path, subtitle_roi)
|
||
print(f"✓ 字幕ROI已保存到: {subtitle_path}")
|
||
print(f" ROI尺寸: {w}x{h}")
|
||
|
||
print("\n测试4完成!")
|
||
|
||
|
||
def main():
|
||
"""主测试函数"""
|
||
print("\n" + "="*60)
|
||
print("字幕检测器完整测试")
|
||
print("="*60)
|
||
|
||
try:
|
||
# 运行所有测试
|
||
test_preset_regions()
|
||
test_auto_detection()
|
||
test_cache_mechanism()
|
||
test_integration()
|
||
|
||
print("\n" + "="*60)
|
||
print("所有测试完成！")
|
||
print("="*60)
|
||
|
||
print("\n生成的测试文件:")
|
||
print(" - /tmp/test_subtitle_frame.png (原始测试帧)")
|
||
print(" - /tmp/test_subtitle_detected.png (带检测框的结果)")
|
||
print(" - /tmp/subtitle_roi.png (提取的字幕区域)")
|
||
|
||
except Exception as e:
|
||
print(f"\n✗ 测试失败: {e}")
|
||
import traceback
|
||
traceback.print_exc()
|
||
return 1
|
||
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|