JeVoisBase  1.18
JeVois Smart Embedded Machine Vision Toolkit Base Modules
Share this page:
quantize-inc.py
Go to the documentation of this file.
1 import os
2 import sys
3 import numpy as np
4 import cv2 as cv
5 
6 import onnx
7 from neural_compressor.experimental import Quantization, common
8 
9 class 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 
30 class 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 
40  self.image_list = self.load_image_list(self.root)
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 
73 models=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 
86 if __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 
quantize-inc.Dataset.__getitem__
def __getitem__(self, idx)
Definition: quantize-inc.py:50
demo.str
str
Definition: demo.py:35
quantize-inc.Dataset.image_list
image_list
Definition: quantize-inc.py:40
quantize-inc.Dataset.mean
mean
Definition: quantize-inc.py:35
quantize-inc.Dataset.root
root
Definition: quantize-inc.py:32
quantize-inc.Quantize.custom_dataset
custom_dataset
Definition: quantize-inc.py:13
quantize-inc.Dataset.std
std
Definition: quantize-inc.py:36
quantize-inc.Dataset.__init__
def __init__(self, root, size=None, dim='chw', mean=0.0, std=1.0, swapRB=False, toFP32=False)
Definition: quantize-inc.py:31
quantize-inc.Dataset
Definition: quantize-inc.py:30
quantize-inc.Dataset.load_image_list
def load_image_list(self, path)
Definition: quantize-inc.py:42
quantize-inc.Dataset.swapRB
swapRB
Definition: quantize-inc.py:37
quantize-inc.Quantize
Definition: quantize-inc.py:9
quantize-inc.Dataset.__len__
def __len__(self)
Definition: quantize-inc.py:70
quantize-inc.Dataset.dim
dim
Definition: quantize-inc.py:34
quantize-inc.Quantize.__init__
def __init__(self, model_path, config_path, custom_dataset=None)
Definition: quantize-inc.py:10
quantize-inc.Quantize.config_path
config_path
Definition: quantize-inc.py:12
quantize-inc.Quantize.model_path
model_path
Definition: quantize-inc.py:11
quantize-inc.Quantize.run
def run(self)
Definition: quantize-inc.py:15
quantize-inc.Dataset.size
size
Definition: quantize-inc.py:33
quantize-inc.Dataset.toFP32
toFP32
Definition: quantize-inc.py:38