-
Notifications
You must be signed in to change notification settings - Fork 0
/
openvn_csrnet.py
63 lines (52 loc) Β· 2.53 KB
/
openvn_csrnet.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
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
import os.path
import cv2
import numpy as np
# Import OpenVINO Inference Engine
from openvino.inference_engine import IECore
from torchvision import transforms
transform = transforms.Compose([transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ])
def run_nn(model_xml_path, file_path):
ie = IECore()
net = ie.read_network(model=model_xml_path)
# Get names of input and output blobs
input_blob = next(iter(net.input_info))
out_blob = next(iter(net.outputs))
# # Get Output Layer Information
# print("Output Layer: ", out_blob)
# # Get Input Shape of Model
print("Input Shape: ", net.input_info[input_blob].input_data.shape)
# # Get Output Shape of Model
print("Output Shape: ", net.outputs[out_blob].shape)
# # Load IECore Object
# print("Available Devices: ", ie.available_devices)
# # ('Starting inference in synchronous mode')
# print('Starting inference in synchronous mode')
exec_net = ie.load_network(network=net, device_name=ie.available_devices[0])
# Following sample Without Image Reshape will make a broadcast error.
# rgb = Image.open(file_path).convert('RGB')
# img = transform(rgb)
# torch.reshape(img, (1, 3, 224, 224))
# output = exec_net.infer(inputs={input_blob: img.unsqueeze(0)})
# Reshape Input Size
# https://github.com/openvinotoolkit/openvino/blob/master/inference-engine/ie_bridges/python/sample/object_detection_sample_ssd/object_detection_sample_ssd.py
original_image = cv2.imread(file_path)
image = original_image.copy()
_, _, net_h, net_w = net.input_info[input_blob].input_data.shape
if image.shape[:-1] != (net_h, net_w):
image = cv2.resize(image, (net_w, net_h))
image = image.transpose((2, 0, 1))
image = np.expand_dims(image, axis=0)
output = exec_net.infer(inputs={input_blob: image})
prediction_arr = next(iter(output.values()))
prediction = int(np.sum(prediction_arr, dtype=np.float32))
print('Approx. Num of peoples: ', prediction)
if __name__ == '__main__':
file_path = './csr_rtn/images_a/IMG_100.jpg'
file_path = os.path.abspath(file_path)
model_xml_path = './model/csrnet-crowd-counting-caffe/shanghaia.xml'
run_nn(model_xml_path, file_path)
# Not sure why it makes diffrent result with same model,
# this model is converted through processes from pytorch to onnx then openvino.
# model_xml_path = '.\\model\\csrnet-crowd-counting-pytorch\\model_a.xml'
# run_nn(model_xml_path, file_path)