""" 滑块验证码识别核心模块 ======================== 基于 YOLOv8-seg（实例分割）的 ONNX 模型，识别滑块验证码的滑块和缺口位置。 整体流程： 原始图片 → 预处理(letterbox+归一化) → ONNX模型推理 → 后处理(NMS+掩膜) → 滑块/缺口坐标 主要对外方法： identify(source) → 返回 (gap_box, confidence) 旧版接口，只返回缺口 identify_both(source) → 返回 dict，包含滑块+缺口+偏移量 新版接口 identify_offset(source) → 返回 (offset, confidence) 只返回滑块x坐标 关键参数（可调，改这里调整识别灵敏度）： CONF_THRESHOLD = 0.5 置信度阈值，低于此值的目标会被丢弃 IOU_THRESHOLD = 0.8 NMS 用的 IoU 阈值 Y_IOU_THRESHOLD = 0.85 Y轴方向 IoU 阈值，用于区分滑块和缺口 """ import base64 import os import random import time from pathlib import Path from typing import List, Tuple, Union import cv2 # OpenCV：图片读写、缩放、绘制 import numpy as np # 数值计算 import onnxruntime as ort # ONNX 模型推理引擎 from shapely.geometry import Polygon # 多边形几何计算（计算IoU用） # ============================================================ # 全局阈值参数（改这里可以调整识别的灵敏度） # ============================================================ CONF_THRESHOLD = 0.5 # 置信度阈值：模型认为"这可能是缺口"的最低分数 # 调低 → 识别更多目标（可能误报增多） # 调高 → 只保留高置信度结果（可能漏报） IOU_THRESHOLD = 0.8 # NMS（非极大值抑制）的 IoU 阈值 # 两个框重叠度超过此值，只保留分数更高的那个 Y_IOU_THRESHOLD = 0.85 # Y轴方向 IoU 阈值，用于 pick_out_mask # 判断两个目标是否在"同一水平线"上（滑块和缺口通常y位置接近） class Slider: def __init__(self): """ 初始化：加载 ONNX 模型文件 模型路径: captcha_recognizer/models/slider.onnx """ root_dir = os.path.dirname(os.path.dirname(__file__)) slider_model_path = os.path.join(root_dir, 'captcha_recognizer', 'models', 'slider.onnx') # 根据是否有 GPU 选择推理设备 # HF Spaces 免费层没有 GPU，所以通常走 CPUExecutionProvider self.session = ort.InferenceSession( slider_model_path, providers=["CUDAExecutionProvider", "CPUExecutionProvider"] if ort.get_device() == 'GPU' else [ "CPUExecutionProvider"], ) # 模型类别：只有一类 's'（slider/缺口） # 如果以后换成多类别模型，这里要加 self.classes = {0: 's'} # ============================================================ # 核心推理流程 # ============================================================ def predict(self, img: np.ndarray, conf: float = 0.25, iou: float = 0.7, imgsz: Union[int, Tuple[int, int]] = 640) -> List: """ 完整推理流程：预处理 → 模型推理 → 后处理 参数： img: 原始图片（BGR格式的numpy数组） conf: 置信度阈值（传给NMS） iou: IoU阈值（传给NMS） imgsz: 模型输入尺寸，默认640x640 ↑ 更大=更准但更慢，更小=更快但可能不准 ↑ 常见值: 320, 640, 1280 """ imgsz = (imgsz, imgsz) if isinstance(imgsz, int) else imgsz prep_img = self.preprocess(img, imgsz) outs = self.session.run(None, {self.session.get_inputs()[0].name: prep_img}) return self.postprocess(img, prep_img, outs, conf=conf, iou=iou) @staticmethod def letterbox(img: np.ndarray, new_shape: Tuple[int, int] = (640, 640)) -> np.ndarray: """ Letterbox 缩放：保持宽高比缩放图片，不足部分用灰色(114,114,114)填充 为什么不直接 resize？因为直接拉伸会变形，影响识别准确率。 Letterbox 相当于"等比缩放 + 补边"，是 YOLO 系列模型的标准做法。 示例：原图 300x200，目标 640x640 → 等比缩放到 640x427 → 上下各补 106 像素灰色边，最终 640x640 """ shape = img.shape[:2] # 当前尺寸 [height, width] # 计算缩放比例（取宽高比中较小的，保证整个图片都能放下） r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) # 确保尺寸不越界 new_unpad = (max(1, min(new_unpad[0], new_shape[1])), max(1, min(new_unpad[1], new_shape[0]))) # 缩放 if shape[::-1] != new_unpad: img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR) # 计算需要填充的像素数 dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] dw, dh = float(dw), float(dh) # 上下、左右各填一半 top, bottom = int(round(dh / 2)), int(round(dh / 2)) left, right = int(round(dw / 2)), int(round(dw / 2)) # 填充灰色边 img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114)) # 最终确保精确尺寸（防止四舍五入导致差1像素） if img.shape[0] != new_shape[0] or img.shape[1] != new_shape[1]: img = cv2.resize(img, new_shape, interpolation=cv2.INTER_LINEAR) return img def preprocess(self, img: np.ndarray, new_shape: Tuple[int, int]) -> np.ndarray: """ 图片预处理：letterbox缩放 → BGR转RGB → 转置 → 归一化 处理后格式: (1, 3, 640, 640) float32，值域 [0, 1] 这是 YOLO 模型的标准输入格式 """ img = self.letterbox(img, new_shape) # 等比缩放+填充 img = img[..., ::-1].transpose([2, 0, 1])[None] # BGR→RGB, HWC→CHW, 加batch维度 img = np.ascontiguousarray(img) # 确保内存连续 img = img.astype(np.float32) / 255 # 归一化到 0~1 return img def postprocess(self, img: np.ndarray, prep_img: np.ndarray, outs: List, conf: float = 0.25, iou: float = 0.7) -> List: """ 后处理：模型输出 → 有意义的检测框和掩膜 模型输出两个部分： preds: 检测框 + 类别 + 置信度 protos: 掩膜原型（用于生成分割掩膜） """ preds, protos = outs preds = self.non_max_suppression(preds, conf, iou, nc=len(self.classes)) results = [] for i, pred in enumerate(preds): if len(pred) == 0: results.append([pred, None]) continue # 把检测框从模型输入坐标映射回原图坐标 pred[:, :4] = self.scale_boxes(prep_img.shape[2:], pred[:, :4], img.shape) # 用掩膜原型 + 检测框系数 → 生成分割掩膜 masks = self.process_mask(protos[i], pred[:, 6:], pred[:, :4], img.shape[:2]) results.append([pred[:, :6], masks]) return results # ============================================================ # 掩膜处理相关 # ============================================================ def process_mask(self, protos: np.ndarray, masks_in: np.ndarray, bboxes: np.ndarray, shape: Tuple[int, int]) -> np.ndarray: """ 从掩膜原型生成最终的二值掩膜 原理：掩膜系数 × 掩膜原型 = 每个目标的分割掩膜 """ c, mh, mw = protos.shape masks = (masks_in @ protos.reshape(c, -1)).reshape(-1, mh, mw) masks = self.scale_masks(masks, shape) masks = self.crop_mask(masks, bboxes) return masks > 0.0 @staticmethod def masks_to_segments(masks: Union[np.ndarray,], strategy: str = "largest") -> List[np.ndarray]: """ 将二值掩膜转换为多边形轮廓点 为什么需要这个？因为后面要计算两个多边形的 IoU 来判断哪个是滑块、哪个是缺口 strategy: 'largest' - 只保留最大轮廓（默认，适合大多数情况） 'all' - 合并所有轮廓 'none' - 保留所有轮廓不合并 """ masks_np = masks.astype("uint8") if masks_np.ndim == 2: masks_np = masks_np[np.newaxis, ...] segments = [] for mask in masks_np: contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) if not contours: segments.append(np.zeros((0, 2), dtype=np.float32)) continue if strategy == "all" and len(contours) > 1: contour = np.concatenate([x.reshape(-1, 2) for x in contours]) elif strategy == "largest": contour = max(contours, key=lambda x: cv2.arcLength(x, closed=True)) contour = contour.reshape(-1, 2) else: contour = contours[0].reshape(-1, 2) segments.append(contour.astype(np.float32)) return segments[0] if masks_np.shape[0] == 1 else segments @staticmethod def draw_segments(image, boxes, masks, mask_alpha=0.5, box_thickness=2, draw_labels=True): """ 在图片上绘制检测框和掩膜（调试用，API服务中不会调用） """ output = image.copy() if boxes is None and masks is None: return output if masks is not None: color_mask = np.zeros_like(image) for i, mask in enumerate(masks): color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) mask = mask.astype(bool) color_mask[mask] = color output = cv2.addWeighted(output, 1, color_mask, mask_alpha, 0) if boxes is not None: for box in boxes: x1, y1, x2, y2, score, class_id = box[:6] color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) cv2.rectangle(output, (int(x1), int(y1)), (int(x2), int(y2)), color, box_thickness) if draw_labels: label = f"{int(class_id)}: {score:.2f}" (label_width, label_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1) cv2.rectangle(output, (int(x1), int(y1) - label_height - 5), (int(x1) + label_width, int(y1)), color, -1) cv2.putText(output, label, (int(x1), int(y1) - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1, cv2.LINE_AA) return output # ============================================================ # 图片输入处理 # ============================================================ @staticmethod def image_to_array(source: Union[str, Path, bytes, np.ndarray] = None): """ 把各种格式的输入统一转成 OpenCV 图片（numpy数组） 支持的输入： - base64 字符串（data:image/png;base64,...） - 文件路径 - 字节流 - 已经是 numpy 数组的（直接返回） """ if isinstance(source, str) and source.startswith('data:image'): header, encoded = source.split(',', 1) data = base64.b64decode(encoded) np_arr = np.frombuffer(data, np.uint8) return cv2.imdecode(np_arr, cv2.IMREAD_COLOR) elif isinstance(source, (str, Path)): return cv2.imread(str(source)) elif isinstance(source, bytes): np_arr = np.frombuffer(source, np.uint8) return cv2.imdecode(np_arr, cv2.IMREAD_COLOR) elif isinstance(source, np.ndarray): return source else: raise TypeError("Unsupported source type. Only str, Path, bytes, or numpy.ndarray are supported.") # ============================================================ # IoU 计算相关（用于区分滑块和缺口） # ============================================================ @staticmethod def normalize_points(points): """将多边形点集归一化到以质心为原点（用于形状比较，消除位置影响）""" centroid = np.mean(points, axis=0) normalized_points = points - centroid return normalized_points @staticmethod def y_iou(segment1, segment2): """ 计算 Y 轴方向的一维 IoU 用途：判断两个目标是否在同一水平线上 滑块和缺口通常 y 坐标接近（在同一高度），而背景干扰物可能在不同高度 """ start = max(segment1[0], segment2[0]) end = min(segment1[1], segment2[1]) intersection = max(0, end - start) len1 = segment1[1] - segment1[0] len2 = segment2[1] - segment2[0] union = len1 + len2 - intersection iou = intersection / union if union != 0 else 0 return iou def polygon_iou(self, poly1, poly2): """ 计算两个多边形的形状 IoU（先归一化位置，只比较形状相似度） 用途：滑块的形状和缺口的形状通常相似（都是拼图块形状） 通过比较形状 IoU 来判断哪个目标是缺口 """ p1 = self.normalize_points(poly1) p2 = self.normalize_points(poly2) poly1 = Polygon(p1).buffer(0) # buffer(0) 修复自相交等无效多边形 poly2 = Polygon(p2).buffer(0) intersect = poly1.intersection(poly2).area union = poly1.union(poly2).area iou = intersect / union if union > 0 else 0.0 return iou # ============================================================ # 核心：从多个检测目标中区分滑块和缺口 # ============================================================ def pick_out_mask(self, boxes: list, segments): """ 从多个检测目标中区分滑块和缺口 返回: (slider_box, gap_box) slider_box: 滑块（x最小的检测框，通常在图片左侧） gap_box: 缺口（与滑块形状最相似的目标，通常在图片右侧） 策略： 1. 找 x 坐标最小的目标 → 这通常是滑块（在图片左侧） 2. 在剩余目标中，找 y 位置与滑块接近的（同一水平线） 3. 在同水平线的目标中，找形状与滑块最相似的 → 这就是缺口 为什么这样判断？ - 滑块验证码的布局：滑块在左，缺口在右 - 滑块和缺口形状相同（都是拼图块），但位置不同 - 背景干扰物形状不同，且可能不在同一水平线 """ # 第一步：找 x 最小的 = 滑块 box_slider = min(boxes, key=lambda x: x[0]) box_slider_index = boxes.index(box_slider) segment_slider = segments[box_slider_index] # 剩余目标（排除滑块） box_sample = boxes[:box_slider_index] + boxes[box_slider_index + 1:] segment_sample = segments[:box_slider_index] + segments[box_slider_index + 1:] # 第二步：Y轴方向 IoU 过滤（只保留同一水平线的目标） box_filtered = [] segment_filtered = [] for index, box in enumerate(box_sample): if self.y_iou([box_slider[1], box_slider[3]], [box[1], box[3]]) > Y_IOU_THRESHOLD: box_filtered.append(box) segment_filtered.append(segment_sample[index]) # 如果Y轴过滤没有保留任何目标，退回到全部候选 if not box_filtered: box_filtered = box_sample segment_filtered = segment_sample if len(box_filtered) == 1: return box_slider, box_filtered[0] # 第三步：找形状与滑块最相似的 = 缺口 iou_flag = 0 iou_index = 0 for index, segment in enumerate(segment_filtered): segment_iou = self.polygon_iou(segment_slider, segment) if segment_iou > iou_flag: iou_flag = segment_iou iou_index = index return box_slider, box_filtered[iou_index] # ============================================================ # 对外接口：识别缺口 # ============================================================ def identify(self, source: Union[str, Path, bytes, np.ndarray], conf=CONF_THRESHOLD, iou=IOU_THRESHOLD, show=False): """ 识别滑块验证码缺口位置（兼容旧接口） 返回： (gap_box, confidence) gap_box = [x1, y1, x2, y2] 缺口的左上角和右下角坐标 confidence = 0~1 置信度 如果没检测到缺口：返回 ([], 0.0) """ result = self.identify_both(source, conf=conf, iou=iou, show=show) return result['gap'], result['gap_confidence'] def identify_both(self, source: Union[str, Path, bytes, np.ndarray], conf=CONF_THRESHOLD, iou=IOU_THRESHOLD, show=False): """ 同时识别滑块和缺口位置（新版接口） 参数： source: 图片输入（路径/base64/字节/numpy数组） conf: 置信度阈值（默认0.5） iou: NMS IoU阈值（默认0.8） show: 是否显示识别结果（调试用，服务器上别开） 返回 dict： slider: [x1, y1, x2, y2] 滑块坐标（空列表=未检测到） slider_confidence: float 滑块置信度 gap: [x1, y1, x2, y2] 缺口坐标（空列表=未检测到） gap_confidence: float 缺口置信度 offset: int 滑动距离 = 缺口x1 - 滑块x1（0表示无法计算） """ slider_box_list = [] gap_box_list = [] original_image: np.ndarray = self.image_to_array(source) results = self.predict(original_image, conf=conf, iou=iou, imgsz=640) if results: boxes, masks = results[0] if len(boxes) == 0: pass # 没检测到任何目标 elif len(boxes) == 1: # 只检测到一个目标，无法区分滑块/缺口，当作缺口 gap_box_list = boxes[0].tolist() else: # 多目标：区分滑块和缺口 segments = self.masks_to_segments(masks) slider_box_list, gap_box_list = self.pick_out_mask(boxes.tolist(), segments) # 调试用：显示识别结果 if show: draw_boxes = [] draw_masks = [] boxes_np, masks_np = results[0] if results else (np.zeros((0, 6)), None) if gap_box_list and masks_np is not None: gap_idx = boxes_np.tolist().index(gap_box_list) if gap_box_list in boxes_np.tolist() else -1 if gap_idx >= 0: draw_boxes.append(gap_box_list) draw_masks.append(masks_np[gap_idx]) if slider_box_list and masks_np is not None: slider_idx = boxes_np.tolist().index(slider_box_list) if slider_box_list in boxes_np.tolist() else -1 if slider_idx >= 0: draw_boxes.append(slider_box_list) draw_masks.append(masks_np[slider_idx]) if draw_boxes: sample = self.draw_segments(original_image, draw_boxes, draw_masks) cv2.imshow('result', sample) cv2.waitKey(0) cv2.destroyAllWindows() # 提取坐标和置信度 slider = [int(x) for x in slider_box_list[:4]] if slider_box_list else [] slider_conf = float(slider_box_list[4]) if slider_box_list else 0.0 gap = [int(x) for x in gap_box_list[:4]] if gap_box_list else [] gap_conf = float(gap_box_list[4]) if gap_box_list else 0.0 # 计算滑动距离：缺口x1 - 滑块x1 # 这就是滑块需要从当前位置滑动到缺口的距离 if slider and gap: offset = gap[0] - slider[0] else: offset = 0 return { 'slider': slider, 'slider_confidence': slider_conf, 'gap': gap, 'gap_confidence': gap_conf, 'offset': offset, } def identify_offset(self, source: Union[str, Path, bytes, np.ndarray], conf=CONF_THRESHOLD, iou=IOU_THRESHOLD, show=False): """ 识别缺口并直接返回偏移量（滑块初始x坐标） 注意：这里返回的 offset 是滑块自身的 x1 坐标， 不是滑动距离。要算滑动距离请用 identify_both() 的 offset 字段。 用途：某些验证码的滑块有固定偏移量，可用此方法获取 """ box_list = [] mask_ndarray = None original_image: np.ndarray = self.image_to_array(source) results = self.predict(original_image, conf=conf, iou=iou, imgsz=640) if results: boxes, masks = results[0] if len(boxes) == 0: pass elif len(boxes) == 1: box_list = boxes[0].tolist() mask_ndarray = masks[0] else: # 多目标时选 x 最小的（最左边的 = 滑块位置） box_left = min(boxes, key=lambda x: x[0]) box_list = box_left.tolist() mask_ndarray = masks[boxes.tolist().index(box_list)] if show and box_list and mask_ndarray is not None: sample = self.draw_segments(original_image, [box_list, ], [mask_ndarray, ]) cv2.imshow('result', sample) cv2.waitKey(0) cv2.destroyAllWindows() if box_list: box = box_list[:4] box_conf = float(box_list[4]) offset = box[0] # 缺口/滑块的 x1 坐标 else: offset = 0 box_conf = 0.0 return offset, box_conf # ============================================================ # 以下是 YOLO 后处理的标准工具方法 # 一般不需要修改，除非换模型 # ============================================================ def scale_boxes(self, img1_shape, boxes, img0_shape, ratio_pad=None, padding=True, xywh=False): """将检测框从模型输入坐标映射回原图坐标（逆letterbox）""" if ratio_pad is None: gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) pad = ( round((img1_shape[1] - img0_shape[1] * gain) / 2), round((img1_shape[0] - img0_shape[0] * gain) / 2), ) else: gain = ratio_pad[0][0] pad = ratio_pad[1] if padding: boxes[..., 0] -= pad[0] boxes[..., 1] -= pad[1] if not xywh: boxes[..., 2] -= pad[0] boxes[..., 3] -= pad[1] boxes[..., :4] /= gain return self.clip_boxes(boxes, img0_shape) @staticmethod def get_covariance_matrix(boxes: np.ndarray): """从旋转边界框生成协方差矩阵（用于旋转框NMS）""" gbbs = np.concatenate((np.power(boxes[:, 2:4], 2) / 12, boxes[:, 4:]), axis=-1) a, b, c = np.split(gbbs, [1, 2], axis=-1) cos = np.cos(c) sin = np.sin(c) cos2 = np.power(cos, 2) sin2 = np.power(sin, 2) return a * cos2 + b * sin2, a * sin2 + b * cos2, (a - b) * cos * sin def batch_probiou(self, obb1, obb2, eps=1e-7): """计算旋转边界框的概率IoU""" x1, y1 = np.split(obb1[..., :2], 2, axis=-1) x2, y2 = (np.expand_dims(x.squeeze(-1), 0) for x in np.split(obb2[..., :2], 2, axis=-1)) a1, b1, c1 = self.get_covariance_matrix(obb1) a2, b2, c2 = (np.expand_dims(x.squeeze(-1), 0) for x in self.get_covariance_matrix(obb2)) t1 = ( ((a1 + a2) * np.power(y1 - y2, 2) + (b1 + b2) * np.power(x1 - x2, 2)) / ( (a1 + a2) * (b1 + b2) - np.power(c1 + c2, 2) + eps) ) * 0.25 t2 = (((c1 + c2) * (x2 - x1) * (y1 - y2)) / ((a1 + a2) * (b1 + b2) - np.power(c1 + c2, 2) + eps)) * 0.5 term1_log = (a1 * b1 - np.power(c1, 2)).clip(0) term2_log = (a2 * b2 - np.power(c2, 2)).clip(0) denominator = 4 * np.sqrt(term1_log * term2_log) + eps t3_numerator = (a1 + a2) * (b1 + b2) - np.power(c1 + c2, 2) t3_arg = np.clip(t3_numerator / denominator + eps, eps, None) t3 = np.log(t3_arg) * 0.5 bd = (t1 + t2 + t3).clip(eps, 100.0) hd = np.sqrt(1.0 - np.exp(-bd) + eps) return 1 - hd def nms_rotated(self, boxes, scores, threshold=0.45): """旋转边界框的NMS""" sorted_idx = np.argsort(scores)[::-1] boxes = boxes[sorted_idx] ious = self.batch_probiou(boxes, boxes) n = boxes.shape[0] ious[np.tril_indices(n)] = 0 pick = np.where((ious >= threshold).sum(axis=0) <= 0)[0] return sorted_idx[pick] def clip_boxes(self, boxes, shape): """将检测框裁剪到图片范围内（防止越界）""" boxes[..., [0, 2]] = np.clip(boxes[..., [0, 2]], 0, shape[1]) boxes[..., [1, 3]] = np.clip(boxes[..., [1, 3]], 0, shape[0]) return boxes @staticmethod def xywh2xyxy(x): """坐标格式转换：中心点+宽高 → 左上角+右下角""" assert x.shape[-1] == 4, f"input shape last dimension expected 4 but input shape is {x.shape}" y = np.empty_like(x, dtype=np.float32) xy = x[..., :2] wh = x[..., 2:] / 2 y[..., :2] = xy - wh y[..., 2:] = xy + wh return y @staticmethod def crop_mask(masks, boxes): """将掩膜裁剪到检测框范围内（框外的掩膜置零）""" _, h, w = masks.shape boxes = boxes[:, :, None] if boxes.ndim == 2 else boxes x1, y1, x2, y2 = np.split(boxes, 4, axis=1) r = np.arange(w, dtype=x1.dtype)[None, None, :] c = np.arange(h, dtype=x1.dtype)[None, :, None] return masks * ((r >= x1) * (r < x2) * (c >= y1) * (c < y2)) def process_mask_np(self, protos, masks_in, bboxes, shape, upsample=False): """另一版本的掩膜处理（带下采样坐标映射）""" c, mh, mw = protos.shape ih, iw = shape masks = (masks_in @ protos.reshape(c, -1)).reshape(-1, mh, mw) width_ratio = mw / iw height_ratio = mh / ih downsampled_bboxes = bboxes.copy() downsampled_bboxes[:, 0] *= width_ratio downsampled_bboxes[:, 2] *= width_ratio downsampled_bboxes[:, 3] *= height_ratio downsampled_bboxes[:, 1] *= height_ratio masks = self.crop_mask(masks, downsampled_bboxes) if upsample: masks = cv2.resize(masks.transpose((1, 2, 0)), (shape[1], shape[0]), interpolation=cv2.INTER_LINEAR).transpose((2, 0, 1)) return masks > 0.0 @staticmethod def scale_masks(masks, shape, padding=True): """将掩膜从模型输出尺寸缩放到原图尺寸""" mh, mw = masks.shape[1:] gain = min(mh / shape[0], mw / shape[1]) pad = [mw - shape[1] * gain, mh - shape[0] * gain] if padding: pad[0] /= 2 pad[1] /= 2 top, left = (int(round(pad[1])), int(round(pad[0]))) if padding else (0, 0) bottom, right = (mh - int(round(pad[1])), mw - int(round(pad[0]))) masks_cropped = masks[:, top:bottom, left:right] resized_masks = np.zeros((masks_cropped.shape[0], shape[0], shape[1]), dtype=masks_cropped.dtype) for i, mask in enumerate(masks_cropped): resized_masks[i] = cv2.resize(mask, (shape[1], shape[0]), interpolation=cv2.INTER_LINEAR) return resized_masks def non_max_suppression(self, prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False, labels=(), max_det=300, nc=0, max_time_img=0.05, max_nms=30000, max_wh=7680, in_place=True, rotated=False, end2end=False, return_idxs=False): """ 非极大值抑制（NMS） 作用：模型可能对同一个目标检测出多个重叠的框， NMS 会保留分数最高的，去掉与它重叠太多的其他框。 一般不需要改这里的参数，调 CONF_THRESHOLD 和 IOU_THRESHOLD 就够了 """ assert 0 <= conf_thres <= 1, f"Invalid Confidence threshold {conf_thres}" assert 0 <= iou_thres <= 1, f"Invalid IoU {iou_thres}" if isinstance(prediction, (list, tuple)): prediction = prediction[0] if classes is not None: classes = np.array(classes) if prediction.shape[-1] == 6 or end2end: output = [pred[pred[:, 4] > conf_thres][:max_det] for pred in prediction] if classes is not None: output = [pred[np.any(pred[:, 5:6] == classes, axis=1)] for pred in output] return output bs = prediction.shape[0] nc = nc or (prediction.shape[1] - 4) extra = prediction.shape[1] - nc - 4 mi = 4 + nc xc = np.amax(prediction[:, 4:mi], axis=1) > conf_thres xinds = np.stack([np.arange(len(i)) for i in xc])[..., None] time_limit = 2.0 + max_time_img * bs multi_label &= nc > 1 prediction = np.transpose(prediction, (0, 2, 1)) if not rotated: if in_place: prediction[..., :4] = self.xywh2xyxy(prediction[..., :4]) else: prediction = np.concatenate((self.xywh2xyxy(prediction[..., :4]), prediction[..., 4:]), axis=-1) t = time.time() output = [np.zeros((0, 6 + extra), dtype=np.float32)] * bs keepi = [np.zeros((0, 1), dtype=np.int64)] * bs for xi, (x, xk) in enumerate(zip(prediction, xinds)): filt = xc[xi] x, xk = x[filt], xk[filt] if labels and len(labels) > xi and len(labels[xi]) and not rotated: lb = np.array(labels[xi]) if lb.size > 0: v = np.zeros((len(lb), nc + extra + 4), dtype=np.float32) v[:, :4] = self.xywh2xyxy(lb[:, 1:5]) v[range(len(lb)), lb[:, 0].astype(np.int64) + 4] = 1.0 x = np.concatenate((x, v), axis=0) if not x.shape[0]: continue box, cls, mask = np.split(x, [4, 4 + nc], axis=1) if multi_label: i, j = np.where(cls > conf_thres) x = np.concatenate((box[i], x[i, 4 + j, None], j[:, None].astype(np.float32), mask[i]), axis=1) xk = xk[i] else: conf = np.amax(cls, axis=1, keepdims=True) j = np.argmax(cls, axis=1, keepdims=True) filt = conf.squeeze(-1) > conf_thres x = np.concatenate((box, conf, j.astype(np.float32), mask), axis=1)[filt] xk = xk[filt] if classes is not None: filt = np.any(x[:, 5:6] == classes, axis=1) x, xk = x[filt], xk[filt] n = x.shape[0] if not n: continue if n > max_nms: filt = np.argsort(x[:, 4])[::-1][:max_nms] x, xk = x[filt], xk[filt] c = x[:, 5:6] * (0 if agnostic else max_wh) scores = x[:, 4] if rotated: boxes = np.concatenate((x[:, :2] + c, x[:, 2:4], x[:, -1:]), axis=-1) i = self.nms_rotated(boxes, scores, iou_thres) else: boxes = x[:, :4] + c i = [] if boxes.shape[0] > 0: y1, x1, y2, x2 = boxes[:, 1], boxes[:, 0], boxes[:, 3], boxes[:, 2] area = (x2 - x1) * (y2 - y1) order = scores.argsort()[::-1] while order.size > 0: idx = order[0] i.append(idx) xx1 = np.maximum(x1[idx], x1[order[1:]]) yy1 = np.maximum(y1[idx], y1[order[1:]]) xx2 = np.minimum(x2[idx], x2[order[1:]]) yy2 = np.minimum(y2[idx], y2[order[1:]]) w = np.maximum(0.0, xx2 - xx1) h = np.maximum(0.0, yy2 - yy1) inter = w * h iou = inter / (area[idx] + area[order[1:]] - inter) order = order[np.where(iou <= iou_thres)[0] + 1] i = np.array(i) i = i[:max_det] output[xi], keepi[xi] = x[i], xk[i].reshape(-1) if (time.time() - t) > time_limit: break return (output, keepi) if return_idxs else output if __name__ == "__main__": """本地测试：直接运行识别单张图片""" model = Slider() # 测试新版接口（同时返回滑块和缺口） res = model.identify_both(source='img_example.png', show=True) print(f'滑块: {res["slider"]}, 置信度: {res["slider_confidence"]}') print(f'缺口: {res["gap"]}, 置信度: {res["gap_confidence"]}') print(f'滑动距离: {res["offset"]}')