import time import numpy as np # 修复np.int已弃用的问题 if not hasattr(np, 'int'): np.int = int from paddleocr import PaddleOCR from service.sensitive_service import get_all_sensitive_words # AC自动机节点定义 class AhoNode: def __init__(self): self.children = {} # 子节点映射（字符->节点） self.fail = None # 失败指针（类似KMP的next数组） self.is_end = False # 标记是否为某个模式串的结尾 self.word = None # 存储当前结尾对应的完整违禁词 # AC自动机实现（多模式字符串匹配） class AhoCorasick: def __init__(self): self.root = AhoNode() # 根节点 def add_word(self, word): """添加违禁词到Trie树""" if not isinstance(word, str) or not word.strip(): return # 过滤无效词 node = self.root for char in word: if char not in node.children: node.children[char] = AhoNode() node = node.children[char] node.is_end = True node.word = word # 记录完整词 def build_fail(self): """构建失败指针（BFS遍历）""" queue = [] # 根节点的子节点失败指针指向根节点 for child in self.root.children.values(): child.fail = self.root queue.append(child) # BFS处理其他节点 while queue: current_node = queue.pop(0) # 遍历当前节点的所有子节点 for char, child in current_node.children.items(): # 寻找失败指针目标节点 fail_node = current_node.fail while fail_node is not None and char not in fail_node.children: fail_node = fail_node.fail # 确定失败指针指向 child.fail = fail_node.children[char] if (fail_node and char in fail_node.children) else self.root queue.append(child) def match(self, text): """匹配文本中所有出现的违禁词（去重）""" result = set() node = self.root for char in text: # 沿失败链查找可用节点 while node is not None and char not in node.children: node = node.fail # 重置到根节点（如果没找到） node = node.children[char] if (node and char in node.children) else self.root # 收集所有匹配的违禁词（包括失败链上的） temp = node while temp != self.root: if temp.is_end: result.add(temp.word) temp = temp.fail return list(result) # 全局变量 _ocr_engine = None _ac_automaton = None # 替换原有的_forbidden_words集合 _conf_threshold = 0.5 def set_forbidden_words(new_words): """更新违禁词（使用AC自动机存储）""" global _ac_automaton if not isinstance(new_words, (set, list, tuple)): raise TypeError("新违禁词必须是集合、列表或元组类型") # 初始化AC自动机并添加有效词 _ac_automaton = AhoCorasick() valid_words = [word for word in new_words if isinstance(word, str) and word.strip()] for word in valid_words: _ac_automaton.add_word(word.strip()) # 构建失败指针（关键步骤） _ac_automaton.build_fail() print(f"已通过函数更新违禁词，当前数量: {len(valid_words)}") def load_forbidden_words(): """从敏感词服务加载违禁词并初始化AC自动机""" global _ac_automaton try: sensitive_words = get_all_sensitive_words() # 保持原接口不变（返回list[str]） _ac_automaton = AhoCorasick() # 添加所有有效敏感词 valid_words = [word for word in sensitive_words if isinstance(word, str) and word.strip()] for word in valid_words: _ac_automaton.add_word(word.strip()) # 构建失败指针 _ac_automaton.build_fail() print(f"加载的违禁词数量: {len(valid_words)}") return True except Exception as e: print(f"Forbidden words load error: {e}") return False def init_ocr_engine(): """初始化OCR引擎和违禁词自动机""" global _ocr_engine try: _ocr_engine = PaddleOCR( use_angle_cls=True, lang="ch", show_log=False, use_gpu=True, max_text_length=1024 ) load_result = load_forbidden_words() if not load_result: print("警告：违禁词加载失败，可能影响检测功能") print("OCR引擎初始化完成") return True except Exception as e: print(f"OCR引擎初始化错误: {e}") _ocr_engine = None return False def detect(frame, conf_threshold=0.8): """检测帧中的文本是否包含违禁词（拆分OCR和匹配时间）""" print("开始进行OCR检测...") total_start = time.time() # 总耗时开始 ocr_time = 0.0 # OCR及结果解析耗时 match_time = 0.0 # 违禁词匹配耗时 try: if not _ocr_engine or not _ac_automaton: return (False, "OCR引擎或违禁词库未初始化") # 1. OCR识别及结果解析阶段 ocr_start = time.time() ocr_res = _ocr_engine.ocr(frame, cls=True) if not ocr_res or not isinstance(ocr_res, list): return (False, "无OCR结果") texts = [] confs = [] # 解析OCR结果 for line in ocr_res: if line is None: continue items_to_process = line if isinstance(line, list) else [line] for item in items_to_process: # 过滤坐标类数据 if isinstance(item, list) and len(item) == 4: is_coordinate = all(isinstance(p, list) and len(p) == 2 and all(isinstance(c, (int, float)) for c in p) for p in item) if is_coordinate: continue # 提取文本和置信度 if isinstance(item, tuple) and len(item) == 2: text, conf = item if isinstance(text, str) and isinstance(conf, (int, float)): texts.append(text.strip()) confs.append(float(conf)) continue if isinstance(item, list) and len(item) >= 2: text_data = item[1] if isinstance(text_data, tuple) and len(text_data) == 2: text, conf = text_data if isinstance(text, str) and isinstance(conf, (int, float)): texts.append(text.strip()) confs.append(float(conf)) continue elif isinstance(text_data, str): texts.append(text_data.strip()) confs.append(1.0) continue ocr_end = time.time() ocr_time = ocr_end - ocr_start # 计算OCR阶段耗时 if len(texts) != len(confs): return (False, "OCR结果格式异常") # 2. 违禁词匹配阶段 match_start = time.time() vio_words = [] for txt, conf in zip(texts, confs): if conf < _conf_threshold: continue # 用AC自动机匹配当前文本中的所有违禁词 matched_words = _ac_automaton.match(txt) # 全局去重并保持顺序 for word in matched_words: if word not in vio_words: vio_words.append(word) match_end = time.time() match_time = match_end - match_start # 计算匹配阶段耗时 has_text = len(texts) > 0 has_violation = len(vio_words) > 0 if not has_text: return (False, "未识别到文本") elif has_violation: return (True, ", ".join(vio_words)) else: return (False, "未检测到违禁词") except Exception as e: print(f"OCR detect error: {e}") return (False, f"检测错误: {str(e)}") finally: # 打印各阶段耗时 total_time = time.time() - total_start print(f"当前帧耗时明细：") print(f" OCR识别及解析：{ocr_time:.8f}秒") print(f" 违禁词匹配：{match_time:.8f}秒") print(f" 总耗时：{total_time:.8f}秒")