""" 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 }