JeVoisBase  1.21
JeVois Smart Embedded Machine Vision Toolkit Base Modules
Share this page:
Loading...
Searching...
No Matches
quantize-inc.py
Go to the documentation of this file.
1import os
2import sys
3import numpy as np
4import cv2 as cv
5
6import onnx
7from neural_compressor.experimental import Quantization, common
8
9class Quantize:
10 def __init__(self, model_path, config_path, custom_dataset=None):
11 self.model_path = model_path
12 self.config_path = config_path
13 self.custom_dataset = custom_dataset
14
15 def run(self):
16 print('Quantizing (int8) with Intel\'s Neural Compressor:')
17 print('\tModel: {}'.format(self.model_path))
18 print('\tConfig: {}'.format(self.config_path))
19
20 output_name = '{}-int8-quantized.onnx'.format(self.model_path[:-5])
21
22 model = onnx.load(self.model_path)
23 quantizer = Quantization(self.config_path)
24 if self.custom_dataset is not None:
25 quantizer.calib_dataloader = common.DataLoader(self.custom_dataset)
26 quantizer.model = common.Model(model)
27 q_model = quantizer()
28 q_model.save(output_name)
29
30class Dataset:
31 def __init__(self, root, size=None, dim='chw', mean=0.0, std=1.0, swapRB=False, toFP32=False):
32 self.root = root
33 self.size = size
34 self.dim = dim
35 self.mean = mean
36 self.std = std
37 self.swapRB = swapRB
38 self.toFP32 = toFP32
39
41
42 def load_image_list(self, path):
43 image_list = []
44 for f in os.listdir(path):
45 if not f.endswith('.jpg'):
46 continue
47 image_list.append(os.path.join(path, f))
48 return image_list
49
50 def __getitem__(self, idx):
51 img = cv.imread(self.image_list[idx])
52
53 if self.swapRB:
54 img = cv.cvtColor(img, cv.COLOR_BGR2RGB)
55
56 if self.size:
57 img = cv.resize(img, dsize=self.size)
58
59 if self.toFP32:
60 img = img.astype(np.float32)
61
62 img = img - self.mean
63 img = img / self.std
64
65 if self.dim == 'chw':
66 img = img.transpose(2, 0, 1) # hwc -> chw
67
68 return img, 1
69
70 def __len__(self):
71 return len(self.image_list)
72
73models=dict(
74 mobilenetv1=Quantize(model_path='../../models/image_classification_mobilenet/image_classification_mobilenetv1_2022apr.onnx',
75 config_path='./inc_configs/mobilenet.yaml'),
76 mobilenetv2=Quantize(model_path='../../models/image_classification_mobilenet/image_classification_mobilenetv2_2022apr.onnx',
77 config_path='./inc_configs/mobilenet.yaml'),
78 mp_palmdet=Quantize(model_path='../../models/palm_detection_mediapipe/palm_detection_mediapipe_2022may.onnx',
79 config_path='./inc_configs/mp_palmdet.yaml',
80 custom_dataset=Dataset(root='../../benchmark/data/palm_detection', dim='hwc', swapRB=True, mean=127.5, std=127.5, toFP32=True)),
81 lpd_yunet=Quantize(model_path='../../models/license_plate_detection_yunet/license_plate_detection_lpd_yunet_2022may.onnx',
82 config_path='./inc_configs/lpd_yunet.yaml',
83 custom_dataset=Dataset(root='../../benchmark/data/license_plate_detection', size=(320, 240), dim='chw', toFP32=True)),
84)
85
86if __name__ == '__main__':
87 selected_models = []
88 for i in range(1, len(sys.argv)):
89 selected_models.append(sys.argv[i])
90 if not selected_models:
91 selected_models = list(models.keys())
92 print('Models to be quantized: {}'.format(str(selected_models)))
93
94 for selected_model_name in selected_models:
95 q = models[selected_model_name]
96 q.run()
97
__init__(self, root, size=None, dim='chw', mean=0.0, std=1.0, swapRB=False, toFP32=False)
load_image_list(self, path)
__init__(self, model_path, config_path, custom_dataset=None)