model_watermark_verify/watermark_verify/inference/ssd.py

import numpy as np
import onnxruntime
from PIL import Image

from watermark_verify.inference.yolox import compute_ciou
from watermark_verify.tools import parse_qrcode_label_file
from watermark_verify.utils.anchors import get_anchors
from watermark_verify.utils.utils_bbox import BBoxUtility


# ---------------------------------------------------------#
#   将图像转换成RGB图像，防止灰度图在预测时报错。
#   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
# ---------------------------------------------------------#
def cvtColor(image):
    if len(np.shape(image)) == 3 and np.shape(image)[2] == 3:
        return image
    else:
        image = image.convert('RGB')
        return image


# ---------------------------------------------------#
#   对输入图像进行resize
# ---------------------------------------------------#
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


# ---------------------------------------------------#
#   获得学习率
# ---------------------------------------------------#
def preprocess_input(inputs):
    MEANS = (104, 117, 123)
    return inputs - MEANS


# ---------------------------------------------------#
#   处理输入图像
# ---------------------------------------------------#
def deal_img(img_path, resized_size):
    image = Image.open(img_path)
    image_shape = np.array(np.shape(image)[0:2])
    # ---------------------------------------------------------#
    #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
    #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
    # ---------------------------------------------------------#
    image = cvtColor(image)
    image_data = resize_image(image, resized_size, False)
    image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0)
    image_data = image_data.astype('float32')
    return image_data, image_shape


# ---------------------------------------------------#
#   检测图像水印
# ---------------------------------------------------#
def detect_watermark(results, watermark_box, threshold=0.5):
    # 解析输出结果
    if len(results[0]) == 0:
        return False
    top_label = np.array(results[0][:, 4], dtype='int32')
    top_conf = results[0][:, 5]
    top_boxes = results[0][:, :4]
    for box, score, cls in zip(top_boxes, top_conf, top_label):
        wm_box_coords = watermark_box[:4]
        wm_cls = watermark_box[4]
        if cls == wm_cls:
            ciou = compute_ciou(box, wm_box_coords)
            if ciou > threshold:
                return True
    return False


def predict_and_detect(image_path, model_file, watermark_txt, input_shape) -> bool:
    """
    使用指定onnx文件进行预测并进行黑盒水印检测
    :param image_path: 输入图像路径
    :param model_file: 模型文件路径
    :param watermark_txt: 水印标签文件路径
    :param input_shape: 模型输入图像大小，tuple
    :return:
    """
    image_data, image_shape = deal_img(image_path, input_shape)
    # 解析标签嵌入位置
    parse_label = parse_qrcode_label_file.load_watermark_info(watermark_txt, image_path)
    if len(parse_label) < 5:
        return False
    x_center, y_center, w, h, cls = parse_label

    # 计算绝对坐标
    height, width = image_shape
    x1 = (x_center - w / 2) * width
    y1 = (y_center - h / 2) * height
    x2 = (x_center + w / 2) * width
    y2 = (y_center + h / 2) * height
    watermark_box = [x1, y1, x2, y2, cls]
    if len(watermark_box) == 0:
        return False
    # 使用onnx进行推理
    session = onnxruntime.InferenceSession(model_file)
    ort_inputs = {session.get_inputs()[0].name: image_data}
    output = session.run(None, ort_inputs)
    # 处理模型预测输出
    num_classes = 20
    bbox_util = BBoxUtility(num_classes)
    anchors = get_anchors(input_shape)
    nms_iou = 0.45
    confidence = 0.5
    results = bbox_util.decode_box(output, anchors, image_shape, input_shape, False, nms_iou=nms_iou,
                                   confidence=confidence)

    if results is not None:
        detect_result = detect_watermark(results, watermark_box)
        return detect_result
    else:
        return False