新增对华为昇腾平台的支持

README.md

@@ -5,6 +5,23 @@
 - `master`分支使用onnx权重文件进行水印提取流程
 - `om`分支使用昇腾平台专用om权重文件格式进行水印提取流程
 
+## 转换为om格式脚本
+需要使用ATC工具进行模型转换，需要安装MXvision工具
+- GoogleNet转换
+```shell
+# 修改--model参数，指定onnx文件位置，--output参数，指定输出的om文件名
+atc --model="model.onnx" --framework=5 --output=googlenet --input_shape="input:-1,3,224,224" --dynamic_batch_size="1,10" --soc_version=Ascend310B4 --log=info --host_env_cpu=aarch64
+```
+- SSD转换
+```shell
+atc --model="model.onnx" --framework=5 --output=ssd --input_shape="images:1,3,300,300" --soc_version=Ascend310B4 --log=info --host_env_cpu=aarch64
+```
+- YOLOX转换
+```shell
+atc --model="yolox.onnx" --framework=5 --output=yolox --input_shape="images:1,3,640,640" --soc_version=Ascend310B4 --log=info --host_env_cpu=aarch64
+```
+- Faster-RCNN转换
+
 ## 水印提取流程说明
 - 黑盒水印提取流程
 1. 初始化参数和参数检查（模型绝对路径、触发集位置、触发集水印标记文件qrcode_positions.txt、公钥文件public.key）

tests/inference_test.py

@@ -1,12 +1,15 @@
 from watermark_verify.inference.rcnn_inference import FasterRCNNInference
 from watermark_verify.inference.ssd_inference import SSDInference
 from watermark_verify.inference.yolox_inference import YOLOXInference
+from watermark_verify.inference.classification_inference import ClassificationInference
 
 if __name__ == '__main__':
-    # model_filename = "/mnt/e/DevTools/workspace/model_watermark_detect/tests/models/origin/yolox/yolox.onnx"
+    model_filename = "models/origin/googlenet/googlenet.om"
+    dets = ClassificationInference(model_filename).predict('ILSVRC2012_val_00002138.JPEG')
+    # model_filename = "models/origin/yolox/yolox.om"
     # dets = YOLOXInference(model_filename).predict('coco/val2017/000000005992.jpg')
-    # model_filename = "/mnt/e/DevTools/workspace/model_watermark_detect/tests/models/origin/ssd/models.onnx"
+    # model_filename = "models/origin/ssd/ssd.om"
     # dets = SSDInference(model_filename).predict('VOC2007/JPEGImages/000001.jpg')
-    model_filename = "/mnt/e/DevTools/workspace/model_watermark_detect/tests/models/origin/faster-rcnn/model.onnx"
-    dets = FasterRCNNInference(model_filename).predict('VOC2007/JPEGImages/000001.jpg')
+    # model_filename = "models/origin/faster-rcnn/rcnn_linux_aarch64.om"
+    # dets = FasterRCNNInference(model_filename).predict('VOC2007/JPEGImages/000001.jpg')
     print(f"dets: {dets}")

watermark_verify/inference/classification_inference.py

@@ -4,7 +4,8 @@
 
 import numpy as np
 from PIL import Image
-import onnxruntime as ort
+from mindx.sdk import Tensor # mxVision 中的 Tensor 数据结构
+from mindx.sdk import base # mxVision 推理接口
 
 
 class ClassificationInference:
@@ -48,11 +49,13 @@ class ClassificationInference:
         :param image_path: 图片路径
         :return: 推理结果
         """
-        session = ort.InferenceSession(self.model_path)  # 加载 ONNX 模型
-        input_name = session.get_inputs()[0].name
         image = self.input_processing(image_path)
-        # 执行预测
-        outputs = session.run(None, {input_name: np.expand_dims(image, axis=0)})
+        # 执行推理
+        base.mx_init()  # 初始化 mxVision 资源
+        img = Tensor(image[None, :, :, :])
+        model = base.model(modelPath=self.model_path)  # 初始化 base.model 类
+        outputs = model.infer([img])[0]  # 执行推理。输入数据类型：List[base.Tensor]， 返回模型推理输出的List[base.Tensor]
+        outputs = np.array(outputs.to_host())  # 将outputs移动到内存中，并且将base.Tensor类转为numpy array
         return outputs
 
     def predict_batch(self, image_paths):
@@ -61,8 +64,6 @@ class ClassificationInference:
         :param image_paths: 待推理的图片路径列表
         :return: 批量推理结果
         """
-        session = ort.InferenceSession(self.model_path)  # 加载 ONNX 模型
-        input_name = session.get_inputs()[0].name
         batch_images = []
 
         for image_path in image_paths:
@@ -71,7 +72,10 @@ class ClassificationInference:
 
         # 将批次图片堆叠成 (batch_size, 3, 224, 224) 维度
         batch_images = np.stack(batch_images)
-
-        # 执行预测
-        outputs = session.run(None, {input_name: batch_images})
+        batch_images = Tensor(batch_images)
+        # 执行推理
+        base.mx_init()  # 初始化 mxVision 资源
+        model = base.model(modelPath=self.model_path)  # 初始化 base.model 类
+        outputs = model.infer([batch_images])[0]  # 执行推理。输入数据类型：List[base.Tensor]， 返回模型推理输出的List[base.Tensor]
+        outputs = np.array(outputs.to_host())  # 将outputs移动到内存中，并且将base.Tensor类转为numpy array
         return outputs

watermark_verify/inference/rcnn_inference.py

@@ -2,7 +2,8 @@
 定义Faster-RCNN推理流程
 """
 import numpy as np
-import onnxruntime as ort
+from mindx.sdk import Tensor  # mxVision 中的 Tensor 数据结构
+from mindx.sdk import base  # mxVision 推理接口
 from PIL import Image
 from watermark_verify.utils.utils_bbox import DecodeBox
 
@@ -11,7 +12,7 @@ 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 model_path: 图像分类模型om文件路径
         :param input_size: 模型输入大小
         :param num_classes: 模型目标检测分类数
         :param num_iou: iou阈值
@@ -53,11 +54,22 @@ class FasterRCNNInference:
         :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('float32')}
-        output = session.run(None, ort_inputs)
+        # 使用mindx框架进行om权重文件推理
+        base.mx_init()
+        model = base.model(modelPath=self.model_path)  # 初始化 base.model 类
+        if model is None:
+            raise Exception("模型导入失败！请检查model_path和device_id.")
+        # 确保img_tensor是正确的输入格式
+        input_tensors = Tensor(image_data)  # 将numpy转为转为Tensor类
+        scale = Tensor(1.0)
+        outputs = model.infer([input_tensors, scale])  # 执行推理
+        output = []
+        for item in outputs:
+            item.to_host()  # 将Tensor数据转移到内存
+            item = np.array(item)
+            if item.size == 0:
+                return False
+            output.append(item)
         output = self.output_processing(output, image_shape)
         return output
 

watermark_verify/inference/ssd_inference.py

@@ -2,7 +2,8 @@
 定义SSD推理流程
 """
 import numpy as np
-import onnxruntime as ort
+from mindx.sdk import Tensor  # mxVision 中的 Tensor 数据结构
+from mindx.sdk import base  # mxVision 推理接口
 from PIL import Image
 from watermark_verify.utils.anchors import get_anchors
 from watermark_verify.utils.utils_bbox import BBoxUtility
@@ -55,10 +56,21 @@ class SSDInference:
         :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}
-        output = session.run(None, ort_inputs)
+        # 使用mindx框架进行om权重文件推理
+        base.mx_init()
+        model = base.model(modelPath=self.model_path)  # 初始化 base.model 类
+        if model is None:
+            raise Exception("模型导入失败！请检查model_path和device_id.")
+        # 确保img_tensor是正确的输入格式
+        input_tensors = Tensor(image_data)  # 将numpy转为转为Tensor类
+        outputs = model.infer([input_tensors])  # 执行推理
+        output = []
+        for item in outputs:
+            item.to_host()  # 将Tensor数据转移到内存
+            item = np.array(item)
+            if item.size == 0:
+                return False
+            output.append(item)
         output = self.output_processing(output, image_shape)
         return output
 

watermark_verify/inference/yolox_inference.py

@@ -3,7 +3,8 @@
 """
 import cv2
 import numpy as np
-import onnxruntime as ort
+from mindx.sdk import Tensor  # mxVision 中的 Tensor 数据结构
+from mindx.sdk import base  # mxVision 推理接口
 
 
 class YOLOXInference:
@@ -50,12 +51,18 @@ class YOLOXInference:
         :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)
+        # 使用mindx框架进行om权重文件推理
+        base.mx_init()
+        model = base.model(modelPath=self.model_path)  # 初始化 base.model 类
+        if model is None:
+            raise Exception("模型导入失败！请检查model_path和device_id.")
+        # 确保img_tensor是正确的输入格式
+        input_tensors = img[None, :, :, :]
+        input_tensors = Tensor(input_tensors)  # 将numpy转为转为Tensor类
+        outputs = model.infer([input_tensors])[0]  # 执行推理
+        outputs.to_host()
+        outputs = np.array(outputs)  # 将Tensor数据转移到内存
+        output = self.output_processing(outputs, ratio)
         return output
 
     def output_processing(self, outputs, ratio, p6=False):