Ver código fonte

修改黑盒推理流程

liyan 7 meses atrás
pai
commit
daa2270d4f
1 arquivos alterados com 1 adições e 1 exclusões
  1. 1 1
      watermark_verify/inference/rcnn.py

+ 1 - 1
watermark_verify/inference/rcnn.py

@@ -149,7 +149,7 @@ def predict_and_detect(image_path, model_file, watermark_txt, input_shape) -> bo
         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')}
+    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)
     roi_cls_locs, roi_scores, rois, _ = output
     # 处理模型预测输出