修改YOLOX模型黑盒水印检测流程，将黑盒检测流程与模型推理流程分离

watermark_verify/inference/yolox_inference.py

@@ -0,0 +1,178 @@
+"""
+定义yolox推理流程
+"""
+import cv2
+import numpy as np
+import onnxruntime as ort
+
+
+class YOLOXInference:
+    def __init__(self, model_path, input_size=(640, 640), swap=(2, 0, 1)):
+        """
+        初始化YOLOX模型推理流程
+        :param model_path: 图像分类模型onnx文件路径
+        :param input_size: 模型输入大小
+        :param swap: 变换方式，pytorch需要进行轴变换(默认参数)，tensorflow无需进行轴变换
+        """
+        self.model_path = model_path
+        self.input_size = input_size
+        self.swap = swap
+
+    def input_processing(self, image_path):
+        """
+        对输入图片进行预处理
+        :param image_path: 图片路径
+        :return: 图片经过处理完成的ndarray
+        """
+        img = cv2.imread(image_path)
+        if len(img.shape) == 3:
+            padded_img = np.ones((self.input_size[0], self.input_size[1], 3), dtype=np.uint8) * 114
+        else:
+            padded_img = np.ones(self.input_size, dtype=np.uint8) * 114
+
+        r = min(self.input_size[0] / img.shape[0], self.input_size[1] / img.shape[1])
+        resized_img = cv2.resize(
+            img,
+            (int(img.shape[1] * r), int(img.shape[0] * r)),
+            interpolation=cv2.INTER_LINEAR,
+        ).astype(np.uint8)
+        padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img
+
+        padded_img = padded_img.transpose(self.swap).copy()
+        padded_img = np.ascontiguousarray(padded_img, dtype=np.float32)
+        height, width, channels = img.shape
+        return padded_img, r, height, width, channels
+
+    def predict(self, image_path):
+        """
+        对单张图片进行推理
+        :param image_path: 图片路径
+        :return: 推理结果
+        """
+        img, ratio, height, width, channels = self.input_processing(image_path)
+
+        session = ort.InferenceSession(self.model_path)
+
+        ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]}
+        output = session.run(None, ort_inputs)
+        output = self.output_processing(output[0], ratio)
+        return output
+
+    def output_processing(self, outputs, ratio, p6=False):
+        """
+        对模型输出进行后处理工作
+        :param outputs: 模型原始输出
+        :return: 经过处理完成的模型输出
+        """
+        grids = []
+        expanded_strides = []
+        strides = [8, 16, 32] if not p6 else [8, 16, 32, 64]
+
+        hsizes = [self.input_size[0] // stride for stride in strides]
+        wsizes = [self.input_size[1] // stride for stride in strides]
+
+        for hsize, wsize, stride in zip(hsizes, wsizes, strides):
+            xv, yv = np.meshgrid(np.arange(wsize), np.arange(hsize))
+            grid = np.stack((xv, yv), 2).reshape(1, -1, 2)
+            grids.append(grid)
+            shape = grid.shape[:2]
+            expanded_strides.append(np.full((*shape, 1), stride))
+
+        grids = np.concatenate(grids, 1)
+        expanded_strides = np.concatenate(expanded_strides, 1)
+        outputs[..., :2] = (outputs[..., :2] + grids) * expanded_strides
+        outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * expanded_strides
+        outputs = outputs[0]  # 获取第一张图片的检测结果
+        boxes = outputs[:, :4]
+        scores = outputs[:, 4:5] * outputs[:, 5:]
+
+        boxes_xyxy = np.ones_like(boxes)
+        boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.
+        boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.
+        boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.
+        boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.
+        boxes_xyxy /= ratio
+        dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.1)
+        return dets
+
+
+def nms(boxes, scores, nms_thr):
+    """Single class NMS implemented in Numpy."""
+    x1 = boxes[:, 0]
+    y1 = boxes[:, 1]
+    x2 = boxes[:, 2]
+    y2 = boxes[:, 3]
+
+    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
+    order = scores.argsort()[::-1]
+
+    keep = []
+    while order.size > 0:
+        i = order[0]
+        keep.append(i)
+        xx1 = np.maximum(x1[i], x1[order[1:]])
+        yy1 = np.maximum(y1[i], y1[order[1:]])
+        xx2 = np.minimum(x2[i], x2[order[1:]])
+        yy2 = np.minimum(y2[i], y2[order[1:]])
+
+        w = np.maximum(0.0, xx2 - xx1 + 1)
+        h = np.maximum(0.0, yy2 - yy1 + 1)
+        inter = w * h
+        ovr = inter / (areas[i] + areas[order[1:]] - inter)
+
+        inds = np.where(ovr <= nms_thr)[0]
+        order = order[inds + 1]
+
+    return keep
+
+
+def multiclass_nms_class_agnostic(boxes, scores, nms_thr, score_thr):
+    """Multiclass NMS implemented in Numpy. Class-agnostic version."""
+    cls_inds = scores.argmax(1)
+    cls_scores = scores[np.arange(len(cls_inds)), cls_inds]
+
+    valid_score_mask = cls_scores > score_thr
+    if valid_score_mask.sum() == 0:
+        return None
+    valid_scores = cls_scores[valid_score_mask]
+    valid_boxes = boxes[valid_score_mask]
+    valid_cls_inds = cls_inds[valid_score_mask]
+    keep = nms(valid_boxes, valid_scores, nms_thr)
+    if keep:
+        dets = np.concatenate(
+            [valid_boxes[keep], valid_scores[keep, None], valid_cls_inds[keep, None]], 1
+        )
+    return dets
+
+
+def multiclass_nms_class_aware(boxes, scores, nms_thr, score_thr):
+    """Multiclass NMS implemented in Numpy. Class-aware version."""
+    final_dets = []
+    num_classes = scores.shape[1]
+    for cls_ind in range(num_classes):
+        cls_scores = scores[:, cls_ind]
+        valid_score_mask = cls_scores > score_thr
+        if valid_score_mask.sum() == 0:
+            continue
+        else:
+            valid_scores = cls_scores[valid_score_mask]
+            valid_boxes = boxes[valid_score_mask]
+            keep = nms(valid_boxes, valid_scores, nms_thr)
+            if len(keep) > 0:
+                cls_inds = np.ones((len(keep), 1)) * cls_ind
+                dets = np.concatenate(
+                    [valid_boxes[keep], valid_scores[keep, None], cls_inds], 1
+                )
+                final_dets.append(dets)
+    if len(final_dets) == 0:
+        return None
+    return np.concatenate(final_dets, 0)
+
+
+def multiclass_nms(boxes, scores, nms_thr, score_thr, class_agnostic=True):
+    """Multiclass NMS implemented in Numpy"""
+    if class_agnostic:
+        nms_method = multiclass_nms_class_agnostic
+    else:
+        nms_method = multiclass_nms_class_aware
+    return nms_method(boxes, scores, nms_thr, score_thr)

watermark_verify/process/yolox_pytorch_blackbox_process.py

@@ -2,12 +2,9 @@
 yolox基于pytorch框架的黑盒水印处理验证流程
 """
 import os
-
 import cv2
-import numpy as np
-import onnxruntime
+from watermark_verify.inference.yolox_inference import YOLOXInference
 from watermark_verify.process.general_process_define import BlackBoxWatermarkProcessDefine
-
 from watermark_verify.tools import parse_qrcode_label_file
 from watermark_verify.tools.evaluate_tool import calculate_ciou
 
@@ -27,7 +24,7 @@ class DetectionProcess(BlackBoxWatermarkProcessDefine):
         for cls, images in cls_image_mapping.items():
             for image in images:
                 image_path = os.path.join(self.trigger_dir, image)
-                detect_result = self.detect_secret_label(image_path, self.model_filename, self.qrcode_positions_file, (640, 640))
+                detect_result = self.detect_secret_label(image_path, self.qrcode_positions_file, (640, 640))
                 if detect_result:
                     accessed_cls.add(cls)
                     break
@@ -37,44 +34,16 @@ class DetectionProcess(BlackBoxWatermarkProcessDefine):
         verify_result = self.verify_label()  # 模型标签检测通过，进行标签验证
         return verify_result
 
-
-    def preprocess_image(self, image_path, input_size, swap=(2, 0, 1)):
-        """
-        对输入图片进行预处理
-        :param swap: 维度变换元组，默认按照(2，0，1)进行变换
-        :param image_path: 图片路径
-        :param input_size: 模型输入大小
-        :return: 图片经过处理完成的ndarray
-        """
-        img = cv2.imread(image_path)
-        if len(img.shape) == 3:
-            padded_img = np.ones((input_size[0], input_size[1], 3), dtype=np.uint8) * 114
-        else:
-            padded_img = np.ones(input_size, dtype=np.uint8) * 114
-
-        r = min(input_size[0] / img.shape[0], input_size[1] / img.shape[1])
-        resized_img = cv2.resize(
-            img,
-            (int(img.shape[1] * r), int(img.shape[0] * r)),
-            interpolation=cv2.INTER_LINEAR,
-        ).astype(np.uint8)
-        padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img
-
-        padded_img = padded_img.transpose(swap).copy()
-        padded_img = np.ascontiguousarray(padded_img, dtype=np.float32)
-        height, width, channels = img.shape
-        return padded_img, r, height, width, channels
-
-    def detect_secret_label(self, image_path, model_file, watermark_txt, input_shape) -> bool:
+    def detect_secret_label(self, image_path, watermark_txt, input_shape) -> bool:
         """
         对模型使用触发集进行检查，判断是否存在黑盒模型水印，如果对嵌入水印的图片样本正确率高于阈值，证明模型存在黑盒水印
         :param image_path: 输入图像路径
-        :param model_file: 模型文件路径
         :param watermark_txt: 水印标签文件路径
         :param input_shape: 模型输入图像大小，tuple
         :return: 检测结果
         """
-        img, ratio, height, width, channels = self.preprocess_image(image_path, input_shape)
+        img = cv2.imread(image_path)
+        height, width, channels = img.shape
         x_center, y_center, w, h, cls = parse_qrcode_label_file.load_watermark_info(watermark_txt, image_path)
         # 计算绝对坐标
         x1 = (x_center - w / 2) * width
@@ -85,11 +54,8 @@ class DetectionProcess(BlackBoxWatermarkProcessDefine):
         if len(watermark_box) == 0:
             return False
 
-        session = onnxruntime.InferenceSession(model_file)
+        dets = YOLOXInference(self.model_filename,input_size=input_shape).predict(image_path)
 
-        ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]}
-        output = session.run(None, ort_inputs)
-        dets = postprocess(output[0], input_shape, ratio)[0]
         if dets is not None:
             detect_result = detect_watermark(dets, watermark_box)
             return detect_result
@@ -97,123 +63,6 @@ class DetectionProcess(BlackBoxWatermarkProcessDefine):
             return False
 
 
-def postprocess(outputs, img_size, ratio, p6=False):
-    grids = []
-    expanded_strides = []
-    outputs = outputs[0]
-    strides = [8, 16, 32] if not p6 else [8, 16, 32, 64]
-
-    hsizes = [img_size[0] // stride for stride in strides]
-    wsizes = [img_size[1] // stride for stride in strides]
-
-    for hsize, wsize, stride in zip(hsizes, wsizes, strides):
-        xv, yv = np.meshgrid(np.arange(wsize), np.arange(hsize))
-        grid = np.stack((xv, yv), 2).reshape(1, -1, 2)
-        grids.append(grid)
-        shape = grid.shape[:2]
-        expanded_strides.append(np.full((*shape, 1), stride))
-
-    grids = np.concatenate(grids, 1)
-    expanded_strides = np.concatenate(expanded_strides, 1)
-    outputs[..., :2] = (outputs[..., :2] + grids) * expanded_strides
-    outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * expanded_strides
-
-    boxes = outputs[:, :4]
-    scores = outputs[:, 4:5] * outputs[:, 5:]
-
-    boxes_xyxy = np.ones_like(boxes)
-    boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.
-    boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.
-    boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.
-    boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.
-    boxes_xyxy /= ratio
-    dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.1)
-
-    return dets
-
-
-def nms(boxes, scores, nms_thr):
-    """Single class NMS implemented in Numpy."""
-    x1 = boxes[:, 0]
-    y1 = boxes[:, 1]
-    x2 = boxes[:, 2]
-    y2 = boxes[:, 3]
-
-    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
-    order = scores.argsort()[::-1]
-
-    keep = []
-    while order.size > 0:
-        i = order[0]
-        keep.append(i)
-        xx1 = np.maximum(x1[i], x1[order[1:]])
-        yy1 = np.maximum(y1[i], y1[order[1:]])
-        xx2 = np.minimum(x2[i], x2[order[1:]])
-        yy2 = np.minimum(y2[i], y2[order[1:]])
-
-        w = np.maximum(0.0, xx2 - xx1 + 1)
-        h = np.maximum(0.0, yy2 - yy1 + 1)
-        inter = w * h
-        ovr = inter / (areas[i] + areas[order[1:]] - inter)
-
-        inds = np.where(ovr <= nms_thr)[0]
-        order = order[inds + 1]
-
-    return keep
-
-
-def multiclass_nms_class_agnostic(boxes, scores, nms_thr, score_thr):
-    """Multiclass NMS implemented in Numpy. Class-agnostic version."""
-    cls_inds = scores.argmax(1)
-    cls_scores = scores[np.arange(len(cls_inds)), cls_inds]
-
-    valid_score_mask = cls_scores > score_thr
-    if valid_score_mask.sum() == 0:
-        return None
-    valid_scores = cls_scores[valid_score_mask]
-    valid_boxes = boxes[valid_score_mask]
-    valid_cls_inds = cls_inds[valid_score_mask]
-    keep = nms(valid_boxes, valid_scores, nms_thr)
-    if keep:
-        dets = np.concatenate(
-            [valid_boxes[keep], valid_scores[keep, None], valid_cls_inds[keep, None]], 1
-        )
-    return dets
-
-
-def multiclass_nms_class_aware(boxes, scores, nms_thr, score_thr):
-    """Multiclass NMS implemented in Numpy. Class-aware version."""
-    final_dets = []
-    num_classes = scores.shape[1]
-    for cls_ind in range(num_classes):
-        cls_scores = scores[:, cls_ind]
-        valid_score_mask = cls_scores > score_thr
-        if valid_score_mask.sum() == 0:
-            continue
-        else:
-            valid_scores = cls_scores[valid_score_mask]
-            valid_boxes = boxes[valid_score_mask]
-            keep = nms(valid_boxes, valid_scores, nms_thr)
-            if len(keep) > 0:
-                cls_inds = np.ones((len(keep), 1)) * cls_ind
-                dets = np.concatenate(
-                    [valid_boxes[keep], valid_scores[keep, None], cls_inds], 1
-                )
-                final_dets.append(dets)
-    if len(final_dets) == 0:
-        return None
-    return np.concatenate(final_dets, 0)
-
-
-def multiclass_nms(boxes, scores, nms_thr, score_thr, class_agnostic=True):
-    """Multiclass NMS implemented in Numpy"""
-    if class_agnostic:
-        nms_method = multiclass_nms_class_agnostic
-    else:
-        nms_method = multiclass_nms_class_aware
-    return nms_method(boxes, scores, nms_thr, score_thr)
-
-
 def detect_watermark(dets, watermark_box, threshold=0.5):
     if dets.size == 0:  # 检查是否为空
         return False
@@ -225,5 +74,3 @@ def detect_watermark(dets, watermark_box, threshold=0.5):
             if ciou > threshold:
                 return True
     return False
-
-