1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
|
- import sys
- import onnx
- import os
- import argparse
- import numpy as np
- import cv2
- import onnxruntime
- from tool.utils import *
- from tool.darknet2onnx import *
- def main(cfg_file, namesfile, weight_file, image_path, batch_size):
- if batch_size <= 0:
- onnx_path_demo = transform_to_onnx(cfg_file, weight_file, batch_size)
- else:
- # Transform to onnx as specified batch size
- transform_to_onnx(cfg_file, weight_file, batch_size)
- # Transform to onnx as demo
- onnx_path_demo = transform_to_onnx(cfg_file, weight_file, 1)
- session = onnxruntime.InferenceSession(onnx_path_demo)
- # session = onnx.load(onnx_path)
- print("The model expects input shape: ", session.get_inputs()[0].shape)
- image_src = cv2.imread(image_path)
- detect(session, image_src, namesfile)
- def detect(session, image_src, namesfile):
- IN_IMAGE_H = session.get_inputs()[0].shape[2]
- IN_IMAGE_W = session.get_inputs()[0].shape[3]
- # Input
- resized = cv2.resize(image_src, (IN_IMAGE_W, IN_IMAGE_H), interpolation=cv2.INTER_LINEAR)
- img_in = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
- img_in = np.transpose(img_in, (2, 0, 1)).astype(np.float32)
- img_in = np.expand_dims(img_in, axis=0)
- img_in /= 255.0
- print("Shape of the network input: ", img_in.shape)
- # Compute
- input_name = session.get_inputs()[0].name
- outputs = session.run(None, {input_name: img_in})
- boxes = post_processing(img_in, 0.4, 0.6, outputs)
- class_names = load_class_names(namesfile)
- plot_boxes_cv2(image_src, boxes[0], savename='predictions_onnx.jpg', class_names=class_names)
- if __name__ == '__main__':
- print("Converting to onnx and running demo ...")
- if len(sys.argv) == 6:
- cfg_file = sys.argv[1]
- namesfile = sys.argv[2]
- weight_file = sys.argv[3]
- image_path = sys.argv[4]
- batch_size = int(sys.argv[5])
- main(cfg_file, namesfile, weight_file, image_path, batch_size)
- else:
- print('Please run this way:\n')
- print(' python demo_onnx.py <cfgFile> <namesFile> <weightFile> <imageFile> <batchSize>')
|