249 lines
9.4 KiB
Python
249 lines
9.4 KiB
Python
"""虎牙滑块/点选验证码 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()
|