yml/live-hub-py
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

type: 收窄虎牙 OCR 与登录类型

Browse Source
This commit is contained in:
yml2213
2026-08-30 20:24:25 +08:00
parent cbd314f4e5
commit 92dc461e52
4 changed files with 324 additions and 214 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/huya/verification/ocr.py
+55 -21
View File
@@ -5,6 +5,7 @@ from __future__ import annotations
from functools import lru_cache
from io import BytesIO
from pathlib import Path
from typing import cast
import cv2
import numpy as np
@@ -19,7 +20,12 @@ 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):
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)
@@ -68,7 +74,9 @@ class YoloOnnx:
return np.expand_dims(input_tensor, axis=0)
@staticmethod
def nms(boxes: np.ndarray, scores: np.ndarray, iou_threshold: float = 0.3) -> list[int]:
def nms(
boxes: np.ndarray, scores: np.ndarray, iou_threshold: float = 0.3
) -> list[int]:
order = scores.argsort()[::-1]
keep = []
@@ -84,7 +92,9 @@ class YoloOnnx:
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_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]
)
@@ -134,20 +144,25 @@ class YoloOnnx:
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.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})
outputs = cast(
list[np.ndarray],
self.session.run([self.output_name], {self.input_name: input_tensor}),
)
return self.postprocess(outputs)
@@ -158,7 +173,7 @@ class SimilarityOnnx:
providers = providers or ["CPUExecutionProvider"]
ort.set_default_logger_severity(3)
self.session = ort.InferenceSession(str(model_path), providers=providers)
self.input_shape = [64, 64]
self.input_shape: tuple[int, int] = (64, 64)
@staticmethod
def sigmoid(value: np.ndarray) -> np.ndarray:
@@ -175,24 +190,38 @@ class SimilarityOnnx:
return Image.open(value)
def _tensor(self, value) -> np.ndarray:
image = self._to_image(value).convert("RGB").resize(tuple(reversed(self.input_shape)), 1)
image = (
self._to_image(value)
.convert("RGB")
.resize((self.input_shape[1], self.input_shape[0]), 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]
out = self.session.run(
None, {"x1": self._tensor(image_1), "x2": self._tensor(image_2)}
)
similarity = self.sigmoid(cast(np.ndarray, 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):
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)
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:
@@ -209,11 +238,16 @@ class HuyaCaptchaOcr:
"""识别点选坐标，返回按目标顺序排列的点击点。"""
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]
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]
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])