重命名模型水印处理流程类

watermark_verify/process/faster-rcnn_pytorch_blackbox_process.py

@@ -15,9 +15,9 @@ from watermark_verify.tools.evaluate_tool import calculate_ciou
 from watermark_verify.utils.utils_bbox import DecodeBox
 
 
-class ClassificationProcess(BlackBoxWatermarkProcessDefine):
+class ModelWatermarkProcessor(BlackBoxWatermarkProcessDefine):
     def __init__(self, model_filename):
-        super(ClassificationProcess, self).__init__(model_filename)
+        super(ModelWatermarkProcessor, self).__init__(model_filename)
 
     def process(self) -> bool:
         # 获取权重文件，使用触发集进行模型推理, 将推理结果与触发集预先二维码保存位置进行比对，在误差范围内则进行下一步，否则返回False

watermark_verify/process/faster-rcnn_pytorch_whitebox_process.py

@@ -4,9 +4,9 @@ faster-rcnn基于pytorch框架的白盒水印处理验证流程
 from watermark_verify.process.general_process_define import WhiteBoxWatermarkProcessDefine
 
 
-class ClassificationProcess(WhiteBoxWatermarkProcessDefine):
+class ModelWatermarkProcessor(WhiteBoxWatermarkProcessDefine):
     def __init__(self, model_filename):
-        super(ClassificationProcess, self).__init__(model_filename)
+        super(ModelWatermarkProcessor, self).__init__(model_filename)
 
     def process(self) -> bool:
         """

watermark_verify/process/ssd_pytorch_blackbox_process.py

@@ -16,9 +16,9 @@ from watermark_verify.utils.anchors import get_anchors
 from watermark_verify.utils.utils_bbox import BBoxUtility
 
 
-class ClassificationProcess(BlackBoxWatermarkProcessDefine):
+class DetectionProcess(BlackBoxWatermarkProcessDefine):
     def __init__(self, model_filename):
-        super(ClassificationProcess, self).__init__(model_filename)
+        super(DetectionProcess, self).__init__(model_filename)
 
     def process(self) -> bool:
         # 获取权重文件，使用触发集进行模型推理, 将推理结果与触发集预先二维码保存位置进行比对，在误差范围内则进行下一步，否则返回False

watermark_verify/process/ssd_pytorch_whitebox_process.py

@@ -4,9 +4,9 @@ ssd基于pytorch框架的白盒水印处理验证流程
 from watermark_verify.process.general_process_define import WhiteBoxWatermarkProcessDefine
 
 
-class ClassificationProcess(WhiteBoxWatermarkProcessDefine):
+class DetectionProcess(WhiteBoxWatermarkProcessDefine):
     def __init__(self, model_filename):
-        super(ClassificationProcess, self).__init__(model_filename)
+        super(DetectionProcess, self).__init__(model_filename)
 
     def process(self) -> bool:
         """

watermark_verify/process/yolox_pytorch_blackbox_process.py

@@ -12,9 +12,9 @@ from watermark_verify.tools import parse_qrcode_label_file
 from watermark_verify.tools.evaluate_tool import calculate_ciou
 
 
-class ClassificationProcess(BlackBoxWatermarkProcessDefine):
+class DetectionProcess(BlackBoxWatermarkProcessDefine):
     def __init__(self, model_filename):
-        super(ClassificationProcess, self).__init__(model_filename)
+        super(DetectionProcess, self).__init__(model_filename)
 
     def process(self) -> bool:
         """
@@ -89,18 +89,7 @@ class ClassificationProcess(BlackBoxWatermarkProcessDefine):
 
         ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]}
         output = session.run(None, ort_inputs)
-        predictions = postprocess(output[0], input_shape)[0]
-
-        boxes = predictions[:, :4]
-        scores = predictions[:, 4:5] * predictions[:, 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)
+        dets = postprocess(output[0], input_shape, ratio)[0]
         if dets is not None:
             detect_result = detect_watermark(dets, watermark_box)
             return detect_result
@@ -108,9 +97,10 @@ class ClassificationProcess(BlackBoxWatermarkProcessDefine):
             return False
 
 
-def postprocess(outputs, img_size, p6=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]
@@ -128,7 +118,18 @@ def postprocess(outputs, img_size, p6=False):
     outputs[..., :2] = (outputs[..., :2] + grids) * expanded_strides
     outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * expanded_strides
 
-    return outputs
+    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):

watermark_verify/process/yolox_pytorch_whitebox_process.py

@@ -4,9 +4,9 @@ yolox基于pytorch框架的白盒水印处理验证流程
 from watermark_verify.process.general_process_define import WhiteBoxWatermarkProcessDefine
 
 
-class ClassificationProcess(WhiteBoxWatermarkProcessDefine):
+class DetectionProcess(WhiteBoxWatermarkProcessDefine):
     def __init__(self, model_filename):
-        super(ClassificationProcess, self).__init__(model_filename)
+        super(DetectionProcess, self).__init__(model_filename)
 
     def process(self) -> bool:
         """

watermark_verify/verify_tool.py

@@ -8,7 +8,7 @@ def label_verification(model_filename: str) -> bool:
     :return: 模型标签验证结果
     """
     # 初始化模型水印检测器
-    model_detector = classification_pytorch_blackbox_process.ClassificationProcess(model_filename)
+    model_detector = classification_pytorch_blackbox_process.ModelWatermarkProcessor(model_filename)
     result = model_detector.process()  # 获取模型水印检测结果
     return result