OpenClaw Agent
310 lines
9.9 KiB
Python
310 lines
9.9 KiB
Python
"""
|
||
OCR 文字识别处理器
|
||
|
||
使用 PaddleOCR 实现视频帧文字识别功能，主要用于提取硬字幕
|
||
"""
|
||
|
||
import os
|
||
import re
|
||
import numpy as np
|
||
from typing import List, Dict, Any, Tuple
|
||
from paddleocr import PaddleOCR
|
||
import logging
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class OCRProcessor:
|
||
"""
|
||
OCR 文字识别处理器
|
||
|
||
使用 PaddleOCR 从视频帧中识别文字，支持批量识别和结果过滤
|
||
"""
|
||
|
||
# 相似度阈值：用于判断两个文本是否相似（基于编辑距离）
|
||
SIMILARITY_THRESHOLD = 0.85
|
||
|
||
# 置信度阈值
|
||
CONFIDENCE_THRESHOLD = 0.7
|
||
|
||
# 最小字幕长度（字符数）
|
||
MIN_SUBTITLE_LENGTH = 2
|
||
|
||
# 注意：Python re 模块不支持 \p{P} 等属性类，在 _is_noise 方法中使用自定义逻辑
|
||
|
||
def __init__(self, use_textline_orientation: bool = True, lang: str = 'ch',
|
||
use_gpu: bool = False):
|
||
"""
|
||
初始化 OCR 处理器
|
||
|
||
Args:
|
||
use_textline_orientation: 是否使用文本行方向分类器（PaddleOCR 3.x 参数名）
|
||
lang: 语言支持，'ch' 为中英文混合
|
||
use_gpu: 是否使用 GPU 加速（通过环境变量设置，如需启用请先设置 CUDA_VISIBLE_DEVICES）
|
||
"""
|
||
self.use_textline_orientation = use_textline_orientation
|
||
self.lang = lang
|
||
self.use_gpu = use_gpu
|
||
|
||
# 初始化 PaddleOCR 实例
|
||
self._ocr = None
|
||
self._initialize_ocr()
|
||
|
||
def _initialize_ocr(self):
|
||
"""初始化 PaddleOCR 实例"""
|
||
try:
|
||
# 强制禁用 oneDNN 和 MKL
|
||
os.environ['FLAGS_use_mkldnn'] = '0'
|
||
os.environ['MKL_SERVICE_FORCE_INTEL'] = '1'
|
||
os.environ['OMP_NUM_THREADS'] = '1'
|
||
os.environ['PADDLE_ON_INFER'] = '1'
|
||
|
||
# PaddleOCR 3.x 版本参数
|
||
self._ocr = PaddleOCR(
|
||
use_textline_orientation=self.use_textline_orientation,
|
||
lang=self.lang
|
||
)
|
||
logger.info("PaddleOCR 初始化成功（oneDNN 已禁用）")
|
||
except Exception as e:
|
||
logger.error(f"PaddleOCR 初始化失败: {e}")
|
||
raise RuntimeError(f"无法初始化 PaddleOCR: {e}")
|
||
|
||
def recognize_frames(self, frames: List[Tuple[np.ndarray, float]]) -> List[Dict[str, Any]]:
|
||
"""
|
||
批量识别帧中的文字
|
||
|
||
Args:
|
||
frames: 帧列表，每个元素为 (图像数组, 时间戳) 元组
|
||
- 图像数组：numpy.ndarray，形状为 (H, W, 3) 的 RGB 图像
|
||
- 时间戳：float，单位为秒
|
||
|
||
Returns:
|
||
List[Dict[str, Any]]: 识别结果列表，每个结果包含：
|
||
- time: float, 时间戳（秒）
|
||
- text: str, 识别的文本
|
||
- confidence: float, 置信度 (0-1)
|
||
"""
|
||
if not frames:
|
||
logger.warning("输入帧列表为空")
|
||
return []
|
||
|
||
results = []
|
||
|
||
for image, timestamp in frames:
|
||
try:
|
||
# 调用 PaddleOCR 进行识别
|
||
# 注意：PaddleOCR 3.x 版本的 ocr() 方法参数已简化
|
||
ocr_result = self._ocr.ocr(image)
|
||
|
||
# PaddleOCR 返回格式: [[[box], (text, confidence)], ...]
|
||
if ocr_result and ocr_result[0]:
|
||
for line in ocr_result[0]:
|
||
box, (text, confidence) = line
|
||
if text and confidence:
|
||
results.append({
|
||
'time': timestamp,
|
||
'text': text,
|
||
'confidence': float(confidence)
|
||
})
|
||
except Exception as e:
|
||
logger.warning(f"帧识别失败 (时间戳 {timestamp}s): {e}")
|
||
continue
|
||
|
||
logger.info(f"原始识别结果数量: {len(results)}")
|
||
|
||
# 过滤和清洗结果
|
||
filtered_results = self._filter_and_clean(results)
|
||
logger.info(f"过滤后结果数量: {len(filtered_results)}")
|
||
|
||
return filtered_results
|
||
|
||
def _filter_and_clean(self, results: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||
"""
|
||
过滤和清洗识别结果
|
||
|
||
Args:
|
||
results: 原始识别结果列表
|
||
|
||
Returns:
|
||
List[Dict[str, Any]]: 过滤后的结果列表
|
||
"""
|
||
if not results:
|
||
return []
|
||
|
||
# 步骤 1: 按置信度过滤
|
||
filtered = self._filter_by_confidence(results)
|
||
logger.debug(f"置信度过滤后: {len(filtered)}")
|
||
|
||
# 步骤 2: 过滤噪声（纯符号、单字符等）
|
||
filtered = self._filter_noise(filtered)
|
||
logger.debug(f"噪声过滤后: {len(filtered)}")
|
||
|
||
# 步骤 3: 去除重复或相似的文本
|
||
filtered = self._remove_duplicates(filtered)
|
||
logger.debug(f"去重后: {len(filtered)}")
|
||
|
||
# 步骤 4: 按时间戳排序
|
||
filtered.sort(key=lambda x: x['time'])
|
||
|
||
return filtered
|
||
|
||
def _filter_by_confidence(self, results: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||
"""
|
||
按置信度过滤结果
|
||
|
||
Args:
|
||
results: 识别结果列表
|
||
|
||
Returns:
|
||
List[Dict[str, Any]]: 高置信度结果
|
||
"""
|
||
return [r for r in results if r['confidence'] >= self.CONFIDENCE_THRESHOLD]
|
||
|
||
def _filter_noise(self, results: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||
"""
|
||
过滤噪声结果
|
||
|
||
过滤条件：
|
||
- 文本长度小于最小长度
|
||
- 纯符号或空白字符
|
||
|
||
Args:
|
||
results: 识别结果列表
|
||
|
||
Returns:
|
||
List[Dict[str, Any]]: 过滤后的结果
|
||
"""
|
||
filtered = []
|
||
|
||
for result in results:
|
||
text = result['text'].strip()
|
||
|
||
# 过滤过短的文本
|
||
if len(text) < self.MIN_SUBTITLE_LENGTH:
|
||
continue
|
||
|
||
# 过滤纯符号或空白字符
|
||
if self._is_noise(text):
|
||
continue
|
||
|
||
filtered.append(result)
|
||
|
||
return filtered
|
||
|
||
def _is_noise(self, text: str) -> bool:
|
||
"""
|
||
判断文本是否为噪声
|
||
|
||
Args:
|
||
text: 要判断的文本
|
||
|
||
Returns:
|
||
bool: 如果是噪声返回 True
|
||
"""
|
||
# 使用正则判断是否只包含符号和空白
|
||
# 注意：这里用简单的正则，更复杂的判断可以后续优化
|
||
try:
|
||
# 移除标点符号后检查是否为空
|
||
# 先替换常见的中文标点
|
||
cleaned = re.sub(r'[^\w\u4e00-\u9fff]', '', text)
|
||
return len(cleaned) == 0
|
||
except:
|
||
# 如果正则出错，保守处理，不算噪声
|
||
return False
|
||
|
||
def _remove_duplicates(self, results: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||
"""
|
||
去除重复或相似的文本
|
||
|
||
基于编辑距离（Levenshtein distance）判断文本相似度
|
||
|
||
Args:
|
||
results: 识别结果列表（假设已按时间戳排序）
|
||
|
||
Returns:
|
||
List[Dict[str, Any]]: 去重后的结果
|
||
"""
|
||
if not results:
|
||
return []
|
||
|
||
deduped = [results[0]]
|
||
|
||
for current in results[1:]:
|
||
last = deduped[-1]
|
||
|
||
# 计算相似度
|
||
similarity = self._calculate_similarity(current['text'], last['text'])
|
||
|
||
# 如果相似度低于阈值，则保留
|
||
if similarity < self.SIMILARITY_THRESHOLD:
|
||
deduped.append(current)
|
||
else:
|
||
# 相似文本：保留置信度更高的那个
|
||
if current['confidence'] > last['confidence']:
|
||
deduped[-1] = current
|
||
|
||
return deduped
|
||
|
||
def _calculate_similarity(self, text1: str, text2: str) -> float:
|
||
"""
|
||
计算两个文本的相似度（基于归一化编辑距离）
|
||
|
||
Args:
|
||
text1: 文本1
|
||
text2: 文本2
|
||
|
||
Returns:
|
||
float: 相似度 (0-1)，1 表示完全相同
|
||
"""
|
||
if text1 == text2:
|
||
return 1.0
|
||
|
||
len1, len2 = len(text1), len(text2)
|
||
|
||
if len1 == 0 or len2 == 0:
|
||
return 0.0
|
||
|
||
# 动态规划计算编辑距离
|
||
dp = np.zeros((len1 + 1, len2 + 1))
|
||
|
||
for i in range(len1 + 1):
|
||
dp[i, 0] = i
|
||
|
||
for j in range(len2 + 1):
|
||
dp[0, j] = j
|
||
|
||
for i in range(1, len1 + 1):
|
||
for j in range(1, len2 + 1):
|
||
if text1[i-1] == text2[j-1]:
|
||
dp[i, j] = dp[i-1, j-1]
|
||
else:
|
||
dp[i, j] = min(
|
||
dp[i-1, j] + 1, # 删除
|
||
dp[i, j-1] + 1, # 插入
|
||
dp[i-1, j-1] + 1 # 替换
|
||
)
|
||
|
||
edit_distance = dp[len1, len2]
|
||
max_len = max(len1, len2)
|
||
|
||
# 归一化相似度：1 - (编辑距离 / 最大长度)
|
||
similarity = 1.0 - (edit_distance / max_len)
|
||
|
||
return similarity
|
||
|
||
def get_ocr_info(self) -> Dict[str, Any]:
|
||
"""
|
||
获取 OCR 处理器信息
|
||
|
||
Returns:
|
||
Dict[str, Any]: OCR 配置信息
|
||
"""
|
||
return {
|
||
'use_textline_orientation': self.use_textline_orientation,
|
||
'lang': self.lang,
|
||
'use_gpu': self.use_gpu,
|
||
'confidence_threshold': self.CONFIDENCE_THRESHOLD,
|
||
'similarity_threshold': self.SIMILARITY_THRESHOLD,
|
||
'min_subtitle_length': self.MIN_SUBTITLE_LENGTH
|
||
}
|