Rename onnx_inference.py to infer_onnx.py (#4)
Browse files- Rename onnx_inference.py to infer_onnx.py (c7d7a6f43f39e8f1b95ded6600786e774de9c4a1)
Co-authored-by: Xiaodong Wang <XiaodongWang@users.noreply.huggingface.co>
onnx_inference.py → infer_onnx.py
RENAMED
|
@@ -53,9 +53,9 @@ def make_parser():
|
|
| 53 |
parser = argparse.ArgumentParser("onnxruntime inference sample")
|
| 54 |
parser.add_argument(
|
| 55 |
"-m",
|
| 56 |
-
"--
|
| 57 |
type=str,
|
| 58 |
-
default="./
|
| 59 |
help="input your onnx model.",
|
| 60 |
)
|
| 61 |
parser.add_argument(
|
|
@@ -99,13 +99,13 @@ names = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', '
|
|
| 99 |
|
| 100 |
if __name__ == '__main__':
|
| 101 |
args = make_parser().parse_args()
|
| 102 |
-
onnx_path = args.
|
| 103 |
if args.ipu:
|
| 104 |
providers = ["VitisAIExecutionProvider"]
|
| 105 |
provider_options = [{"config_file": args.provider_config}]
|
| 106 |
-
|
| 107 |
else:
|
| 108 |
-
|
| 109 |
grid = np.load("./grid.npy", allow_pickle=True)
|
| 110 |
anchor_grid = np.load("./anchor_grid.npy", allow_pickle=True)
|
| 111 |
path = args.image_path
|
|
@@ -114,8 +114,8 @@ if __name__ == '__main__':
|
|
| 114 |
|
| 115 |
img0 = cv2.imread(path)
|
| 116 |
img = pre_process(img0)
|
| 117 |
-
onnx_input = {
|
| 118 |
-
onnx_output =
|
| 119 |
onnx_output = [torch.tensor(item).permute(0, 3, 1, 2) for item in onnx_output]
|
| 120 |
onnx_output = post_process(onnx_output)
|
| 121 |
pred = non_max_suppression(
|
|
@@ -137,3 +137,4 @@ if __name__ == '__main__':
|
|
| 137 |
# Stream results
|
| 138 |
im0 = annotator.result()
|
| 139 |
cv2.imwrite(new_path, im0)
|
|
|
|
|
|
| 53 |
parser = argparse.ArgumentParser("onnxruntime inference sample")
|
| 54 |
parser.add_argument(
|
| 55 |
"-m",
|
| 56 |
+
"--onnx_model",
|
| 57 |
type=str,
|
| 58 |
+
default="./yolov5s.onnx",
|
| 59 |
help="input your onnx model.",
|
| 60 |
)
|
| 61 |
parser.add_argument(
|
|
|
|
| 99 |
|
| 100 |
if __name__ == '__main__':
|
| 101 |
args = make_parser().parse_args()
|
| 102 |
+
onnx_path = args.onnx_model
|
| 103 |
if args.ipu:
|
| 104 |
providers = ["VitisAIExecutionProvider"]
|
| 105 |
provider_options = [{"config_file": args.provider_config}]
|
| 106 |
+
onnx_weight = onnxruntime.InferenceSession(onnx_path, providers=providers, provider_options=provider_options)
|
| 107 |
else:
|
| 108 |
+
onnx_weight = onnxruntime.InferenceSession(onnx_path)
|
| 109 |
grid = np.load("./grid.npy", allow_pickle=True)
|
| 110 |
anchor_grid = np.load("./anchor_grid.npy", allow_pickle=True)
|
| 111 |
path = args.image_path
|
|
|
|
| 114 |
|
| 115 |
img0 = cv2.imread(path)
|
| 116 |
img = pre_process(img0)
|
| 117 |
+
onnx_input = {onnx_weight.get_inputs()[0].name: img.transpose(0, 2, 3, 1)}
|
| 118 |
+
onnx_output = onnx_weight.run(None, onnx_input)
|
| 119 |
onnx_output = [torch.tensor(item).permute(0, 3, 1, 2) for item in onnx_output]
|
| 120 |
onnx_output = post_process(onnx_output)
|
| 121 |
pred = non_max_suppression(
|
|
|
|
| 137 |
# Stream results
|
| 138 |
im0 = annotator.result()
|
| 139 |
cv2.imwrite(new_path, im0)
|
| 140 |
+
|