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

watermark_verify/inference/rcnn_inference.py

@@ -0,0 +1,93 @@
+"""
+定义Faster-RCNN推理流程
+"""
+import numpy as np
+import onnxruntime as ort
+from PIL import Image
+from watermark_verify.utils.utils_bbox import DecodeBox
+
+
+class FasterRCNNInference:
+    def __init__(self, model_path, input_size=(600, 600), num_classes=20, num_iou=0.3, confidence=0.5, swap=(2, 0, 1)):
+        """
+        初始化Faster-RCNN模型推理流程
+        :param model_path: 图像分类模型onnx文件路径
+        :param input_size: 模型输入大小
+        :param num_classes: 模型目标检测分类数
+        :param num_iou: iou阈值
+        :param confidence: 置信度阈值
+        :param swap: 变换方式，pytorch需要进行轴变换(默认参数)，tensorflow无需进行轴变换
+        """
+        self.model_path = model_path
+        self.input_size = input_size
+        self.swap = swap
+        self.num_classes = num_classes
+        self.nms_iou = num_iou
+        self.confidence = confidence
+
+    def input_processing(self, image_path):
+        """
+        对输入图片进行预处理
+        :param image_path: 图片路径
+        :return: 图片经过处理完成的ndarray
+        """
+        image = Image.open(image_path)
+        image_shape = np.array(np.shape(image)[0:2])
+        # ---------------------------------------------------------#
+        #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
+        #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
+        # ---------------------------------------------------------#
+        if not (len(np.shape(image)) == 3 and np.shape(image)[2] == 3):
+            image = image.convert('RGB')
+        image_data = resize_image(image, self.input_size, False)
+        MEANS = (104, 117, 123)
+        image_data = np.array(image_data, dtype='float32')
+        image_data = image_data - MEANS
+        image_data = np.expand_dims(np.transpose(image_data, self.swap).copy(), 0)
+        image_data = image_data.astype('float32')
+        return image_data, image_shape
+
+    def predict(self, image_path):
+        """
+        对单张图片进行推理
+        :param image_path: 图片路径
+        :return: 推理结果
+        """
+        image_data, image_shape = self.input_processing(image_path)
+        # 使用onnx文件进行推理
+        session = ort.InferenceSession(self.model_path)
+        ort_inputs = {session.get_inputs()[0].name: image_data,
+                      session.get_inputs()[1].name: np.array(1.0).astype('float64')}
+        output = session.run(None, ort_inputs)
+        output = self.output_processing(output, image_shape)
+        return output
+
+    def output_processing(self, outputs, image_shape):
+        """
+        对模型输出进行后处理工作
+        :param outputs: 模型原始输出
+        :param image_shape: 原始图像大小
+        :return: 经过处理完成的模型输出
+        """
+        # 处理模型预测输出
+        roi_cls_locs, roi_scores, rois, _ = outputs
+        bbox_util = DecodeBox(self.num_classes)
+        results = bbox_util.forward(roi_cls_locs, roi_scores, rois, image_shape, self.input_size,
+                                    nms_iou=self.nms_iou, confidence=self.confidence)
+        return results
+
+
+def resize_image(image, size, letterbox_image):
+    iw, ih = image.size
+    w, h = size
+    if letterbox_image:
+        scale = min(w / iw, h / ih)
+        nw = int(iw * scale)
+        nh = int(ih * scale)
+
+        image = image.resize((nw, nh), Image.BICUBIC)
+        new_image = Image.new('RGB', size, (128, 128, 128))
+        new_image.paste(image, ((w - nw) // 2, (h - nh) // 2))
+    else:
+        new_image = image.resize((w, h), Image.BICUBIC)
+    return new_image

watermark_verify/process/faster-rcnn_pytorch_blackbox_process.py

@@ -3,16 +3,12 @@ faster-rcnn基于pytorch框架的黑盒水印处理验证流程
 """
 
 import os
-
 import numpy as np
-import onnxruntime
 from PIL import Image
-
+from watermark_verify.inference.rcnn_inference import FasterRCNNInference
 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
-from watermark_verify.utils.utils_bbox import DecodeBox
 
 
 class ModelWatermarkProcessor(BlackBoxWatermarkProcessDefine):
@@ -42,21 +38,6 @@ class ModelWatermarkProcessor(BlackBoxWatermarkProcessDefine):
         verify_result = self.verify_label()  # 模型标签检测通过，进行标签验证
         return verify_result
 
-    def preprocess_image(self, image_path, input_size, swap=(2, 0, 1)):
-        image = Image.open(image_path)
-        image_shape = np.array(np.shape(image)[0:2])
-        # ---------------------------------------------------------#
-        #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
-        #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
-        # ---------------------------------------------------------#
-        if not (len(np.shape(image)) == 3 and np.shape(image)[2] == 3):
-            image = image.convert('RGB')
-        image_data = resize_image(image, input_size, False)
-        image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), swap).copy(),
-                                    0)
-        image_data = image_data.astype('float32')
-        return image_data, image_shape
-
     def detect_secret_label(self, image_path, model_file, watermark_txt, input_shape) -> bool:
         """
         使用指定onnx文件进行预测并进行黑盒水印检测
@@ -66,7 +47,8 @@ class ModelWatermarkProcessor(BlackBoxWatermarkProcessDefine):
         :param input_shape: 模型输入图像大小，tuple
         :return:
         """
-        image_data, image_shape = self.preprocess_image(image_path, input_shape)
+        image = Image.open(image_path)
+        image_shape = np.array(np.shape(image)[0:2])
         # 解析标签嵌入位置
         parse_label = parse_qrcode_label_file.load_watermark_info(watermark_txt, image_path)
         if len(parse_label) < 5:
@@ -83,18 +65,9 @@ class ModelWatermarkProcessor(BlackBoxWatermarkProcessDefine):
         if len(watermark_box) == 0:
             return False
         # 使用onnx进行推理
-        session = onnxruntime.InferenceSession(model_file)
-        ort_inputs = {session.get_inputs()[0].name: image_data,
-                      session.get_inputs()[1].name: np.array(1.0).astype('float64')}
-        output = session.run(None, ort_inputs)
-        roi_cls_locs, roi_scores, rois, _ = output
-        # 处理模型预测输出
-        num_classes = 20
-        bbox_util = DecodeBox(num_classes)
-        nms_iou = 0.3
-        confidence = 0.5
-        results = bbox_util.forward(roi_cls_locs, roi_scores, rois, image_shape, input_shape,
-                                    nms_iou=nms_iou, confidence=confidence)
+        results = FasterRCNNInference(self.model_filename).predict(image_path)
+
+        # 检测模型是否存在黑盒水印
         if results is not None:
             detect_result = detect_watermark(results, watermark_box)
             return detect_result