YogeLiu/TexTeller
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update infer_det.py

Browse Source 
增加使用gpu进行onnx模型推理的功能
This commit is contained in:
TonyLee1256
2024-05-09 00:19:39 +08:00
committed by GitHub
parent e495640690
commit 48043d11e3
1 changed files with 7 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
src/infer_det.py
View File

@@ -3,7 +3,7 @@ import argparse
import glob import glob
import subprocess import subprocess
from onnxruntime import InferenceSession import onnxruntime
from pathlib import Path from pathlib import Path
from models.det_model.inference import PredictConfig, predict_image from models.det_model.inference import PredictConfig, predict_image
@@ -17,6 +17,7 @@ parser.add_argument('--onnx_file', type=str, help="onnx model file path",
parser.add_argument("--image_dir", type=str, default='./testImgs') parser.add_argument("--image_dir", type=str, default='./testImgs')
parser.add_argument("--image_file", type=str) parser.add_argument("--image_file", type=str)
parser.add_argument("--imgsave_dir", type=str, default="./detect_results") parser.add_argument("--imgsave_dir", type=str, default="./detect_results")
parser.add_argument('--use_gpu', action='store_true', help='Whether to use GPU for inference', default=True)
def get_test_images(infer_dir, infer_img): def get_test_images(infer_dir, infer_img):
@@ -71,8 +72,11 @@ if __name__ == '__main__':
# load image list # load image list
img_list = get_test_images(FLAGS.image_dir, FLAGS.image_file) img_list = get_test_images(FLAGS.image_dir, FLAGS.image_file)
# load predictor
predictor = InferenceSession(FLAGS.onnx_file) if FLAGS.use_gpu:
predictor = onnxruntime.InferenceSession(FLAGS.onnx_file, providers=['CUDAExecutionProvider'])
else:
predictor = onnxruntime.InferenceSession(FLAGS.onnx_file, providers=['CPUExecutionProvider'])
# load infer config # load infer config
infer_config = PredictConfig(FLAGS.infer_cfg) infer_config = PredictConfig(FLAGS.infer_cfg)