fix: 复制 obsidian-kb 的 src 目录以解决依赖缺失
This commit is contained in:
OpenClaw Agent
@@ -0,0 +1,309 @@
|
||||
"""
|
||||
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
|
||||
}
|
||||
Reference in New Issue
Block a user