Files
FastDeploy/examples/vision/detection/yolov5/python/infer.py
Jason 0990ab9b50 Revert "[Backend] Add KunlunXin XPU deploy support" (#893)
Revert "[Backend] Add KunlunXin XPU deploy support (#747)"

This reverts commit 5be839b322.
2022-12-15 21:17:59 +08:00

62 lines
1.5 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

import fastdeploy as fd
import cv2
def parse_arguments():
import argparse
import ast
parser = argparse.ArgumentParser()
parser.add_argument(
"--model", default=None, help="Path of yolov5 onnx model.")
parser.add_argument(
"--image", default=None, help="Path of test image file.")
parser.add_argument(
"--device",
type=str,
default='cpu',
help="Type of inference device, support 'cpu' or 'gpu'.")
parser.add_argument(
"--use_trt",
type=ast.literal_eval,
default=False,
help="Wether to use tensorrt.")
return parser.parse_args()
def build_option(args):
option = fd.RuntimeOption()
if args.device.lower() == "gpu":
option.use_gpu()
if args.use_trt:
option.use_trt_backend()
option.set_trt_input_shape("images", [1, 3, 640, 640])
return option
args = parse_arguments()
if args.model is None:
model = fd.download_model(name='YOLOv5s')
else:
model = args.model
# 配置runtime加载模型
runtime_option = build_option(args)
model = fd.vision.detection.YOLOv5(model, runtime_option=runtime_option)
# 预测图片检测结果
if args.image is None:
image = fd.utils.get_detection_test_image()
else:
image = args.image
im = cv2.imread(image)
result = model.predict(im)
print(result)
# 预测结果可视化
vis_im = fd.vision.vis_detection(im, result)
cv2.imwrite("visualized_result.jpg", vis_im)
print("Visualized result save in ./visualized_result.jpg")