-
Notifications
You must be signed in to change notification settings - Fork 1
/
main.py
35 lines (21 loc) · 932 Bytes
/
main.py
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
import cv2
import numpy as np
import preprocess_yolov4 as pre
import postprocess_yolov4 as post
from PIL import Image
input_size = 416
original_image = cv2.imread("kite.jpg")
original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
original_image_size = original_image.shape[:2]
image_data = pre.image_preprocess(np.copy(original_image), [input_size, input_size])
image_data = image_data[np.newaxis, ...].astype(np.float32)
print("Preprocessed image shape:",image_data.shape) # shape of the preprocessed input
import onnxruntime as rt
sess = rt.InferenceSession("yolov4.onnx")
output_name = sess.get_outputs()[0].name
input_name = sess.get_inputs()[0].name
detections = sess.run([output_name], {input_name: image_data})[0]
print("Output shape:", detections.shape)
image = post.image_postprocess(original_image, input_size, detections)
image = Image.fromarray(image)
image.save("kite-with-objects.jpg")