"""虎牙滑块/点选验证码 OCR。""" from __future__ import annotations from functools import lru_cache from io import BytesIO from pathlib import Path import cv2 import numpy as np import onnxruntime as ort from PIL import Image from scipy.optimize import linear_sum_assignment MODEL_DIR = Path(__file__).resolve().parent / "models" class YoloOnnx: """轻量 YOLO ONNX 推理封装。""" def __init__(self, model_path: str | Path, classes: list[str], providers: list[str] | None = None): providers = providers or ["CPUExecutionProvider"] ort.set_default_logger_severity(3) self.session = ort.InferenceSession(str(model_path), providers=providers) self.input_name = self.session.get_inputs()[0].name self.output_name = self.session.get_outputs()[0].name self.input_shape = self.session.get_inputs()[0].shape[2:] self.names = classes self.scale = 1.0 self.pad_top = 0 self.pad_left = 0 @staticmethod def xywh2xyxy(value: np.ndarray) -> np.ndarray: result = np.copy(value) result[..., 0] = value[..., 0] - value[..., 2] / 2 result[..., 1] = value[..., 1] - value[..., 3] / 2 result[..., 2] = value[..., 0] + value[..., 2] / 2 result[..., 3] = value[..., 1] + value[..., 3] / 2 return result def preprocess(self, image: Image.Image) -> np.ndarray: img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) height, width = img.shape[:2] self.scale = min(self.input_shape[0] / height, self.input_shape[1] / width) new_size = (int(width * self.scale), int(height * self.scale)) resized = cv2.resize(img, new_size, interpolation=cv2.INTER_LINEAR) dh = self.input_shape[0] - new_size[1] dw = self.input_shape[1] - new_size[0] self.pad_top = dh // 2 pad_bottom = dh - self.pad_top self.pad_left = dw // 2 pad_right = dw - self.pad_left padded = cv2.copyMakeBorder( resized, self.pad_top, pad_bottom, self.pad_left, pad_right, cv2.BORDER_CONSTANT, value=(114, 114, 114), ) input_tensor = padded.transpose(2, 0, 1).astype(np.float32) / 255.0 return np.expand_dims(input_tensor, axis=0) @staticmethod def nms(boxes: np.ndarray, scores: np.ndarray, iou_threshold: float = 0.3) -> list[int]: order = scores.argsort()[::-1] keep = [] while order.size > 0: index = order[0] keep.append(index) xx1 = np.maximum(boxes[index, 0], boxes[order[1:], 0]) yy1 = np.maximum(boxes[index, 1], boxes[order[1:], 1]) xx2 = np.minimum(boxes[index, 2], boxes[order[1:], 2]) yy2 = np.minimum(boxes[index, 3], boxes[order[1:], 3]) width = np.maximum(0.0, xx2 - xx1) height = np.maximum(0.0, yy2 - yy1) intersection = width * height area_i = (boxes[index, 2] - boxes[index, 0]) * (boxes[index, 3] - boxes[index, 1]) area_j = (boxes[order[1:], 2] - boxes[order[1:], 0]) * ( boxes[order[1:], 3] - boxes[order[1:], 1] ) iou = intersection / (area_i + area_j - intersection + 1e-7) inds = np.where(iou <= iou_threshold)[0] order = order[inds + 1] return keep def postprocess( self, outputs: list[np.ndarray], conf_threshold: float = 0.1, iou_threshold: float = 0.3, ) -> list[dict]: predictions = np.squeeze(outputs[0]).T class_scores = predictions[:, 4:] scores = np.max(class_scores, axis=1) class_ids = np.argmax(class_scores, axis=1) valid_indices = scores > conf_threshold predictions = predictions[valid_indices] scores = scores[valid_indices] class_ids = class_ids[valid_indices] if predictions.size == 0: return [] boxes = predictions[:, :4].copy() boxes[:, 0] = (boxes[:, 0] - self.pad_left) / self.scale boxes[:, 1] = (boxes[:, 1] - self.pad_top) / self.scale boxes[:, 2] /= self.scale boxes[:, 3] /= self.scale xyxy_boxes = self.xywh2xyxy(boxes) keep_indices = [] for class_id in np.unique(class_ids): class_indices = np.where(class_ids == class_id)[0] class_keep = self.nms( xyxy_boxes[class_indices], scores[class_indices], iou_threshold, ) keep_indices.extend(class_indices[class_keep]) results = [] for index in keep_indices: class_id = int(class_ids[index]) x1, y1, x2, y2 = map(int, xyxy_boxes[index].tolist()) results.append({ "label_id": class_id, "label_name": self.names[class_id], "confidence": float(scores[index]), "box_mid_xy": [(x1 + x2) // 2, (y1 + y2) // 2], "xyxy": [x1, y1, x2, y2], }) results.sort(key=lambda item: item["confidence"], reverse=True) return results def detect(self, image: Image.Image) -> list[dict]: input_tensor = self.preprocess(image) outputs = self.session.run([self.output_name], {self.input_name: input_tensor}) return self.postprocess(outputs) class SimilarityOnnx: """点选验证码中目标图和候选字的相似度模型。""" def __init__(self, model_path: str | Path, providers: list[str] | None = None): providers = providers or ["CPUExecutionProvider"] ort.set_default_logger_severity(3) self.session = ort.InferenceSession(str(model_path), providers=providers) self.input_shape = [64, 64] @staticmethod def sigmoid(value: np.ndarray) -> np.ndarray: return 1 / (1 + np.exp(-value)) @staticmethod def _to_image(value) -> Image.Image: if isinstance(value, np.ndarray): return Image.fromarray(value) if isinstance(value, bytes): return Image.open(BytesIO(value)) if isinstance(value, Image.Image): return value return Image.open(value) def _tensor(self, value) -> np.ndarray: image = self._to_image(value).convert("RGB").resize(tuple(reversed(self.input_shape)), 1) array = np.array(image).astype(np.float32) / 255.0 return np.expand_dims(np.transpose(array, (2, 0, 1)), 0) def score(self, image_1, image_2) -> int: out = self.session.run(None, {"x1": self._tensor(image_1), "x2": self._tensor(image_2)}) similarity = self.sigmoid(out[0])[0][0] return int(round(similarity.item(), 2) * 100) class HuyaCaptchaOcr: """封装虎牙滑块和点选识别。""" def __init__(self, model_dir: str | Path = MODEL_DIR, providers: list[str] | None = None): model_dir = Path(model_dir) self.similarity = SimilarityOnnx(model_dir / "weights.onnx", providers=providers) self.click_model = YoloOnnx(model_dir / "best.onnx", classes=["target", "char"], providers=providers) self.slider_model = YoloOnnx(model_dir / "slider_2.onnx", classes=["slider"], providers=providers) @staticmethod def _open_image(image_bytes: bytes) -> Image.Image: return Image.open(BytesIO(image_bytes)).convert("RGB") def ocr_slider_box(self, image_bytes: bytes) -> list[int]: """识别滑块缺口框，返回 [x1, y1, x2, y2]。""" image = self._open_image(image_bytes) detections = self.slider_model.detect(image) boxes = [item["xyxy"] for item in detections] return boxes[0] if boxes else [100, 100, 100, 100] def ocr_click_points(self, image_bytes: bytes) -> list[list[int]]: """识别点选坐标，返回按目标顺序排列的点击点。""" image = self._open_image(image_bytes) detections = self.click_model.detect(image) results = [{**item, "cropped_image": image.crop(tuple(item["xyxy"]))} for item in detections] char_list = [item for item in results if item.get("label_name") == "char"] target_list = [item for item in results if item.get("label_name") == "target"] target_list = [item for item in target_list if item["xyxy"][2] - item["xyxy"][0] > 10] char_list.sort(key=lambda item: item["xyxy"][0]) target_list.sort(key=lambda item: item["xyxy"][0]) if not char_list or not target_list: return [] score_matrix = np.zeros((len(target_list), len(char_list))) for target_index, target in enumerate(target_list): for char_index, char in enumerate(char_list): try: score_matrix[target_index][char_index] = -self.similarity.score( target["cropped_image"], char["cropped_image"], ) except Exception: score_matrix[target_index][char_index] = 1e6 row_ind, col_ind = linear_sum_assignment(score_matrix) matched = [] for target_index, char_index in zip(row_ind, col_ind): if score_matrix[target_index][char_index] == 1e6: continue char = char_list[char_index] matched.append({**char, "index": target_index + 1}) return [item["box_mid_xy"] for item in matched] @lru_cache(maxsize=1) def default_ocr() -> HuyaCaptchaOcr: """懒加载默认 OCR，避免导入 core.huya 时立即加载 150MB 模型。""" return HuyaCaptchaOcr()