-
Notifications
You must be signed in to change notification settings - Fork 2
/
predict.py
38 lines (27 loc) · 1.08 KB
/
predict.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
import onnxruntime as ort
import utils.process as process
class FlowerClassifier:
def __init__(self, configs):
if configs['device'] == 'CUDA':
providers = ['CUDAExecutionProvider']
else:
providers = ['CPUExecutionProvider']
self.configs = configs
self.session = ort.InferenceSession(f'weights/classify-{self.precision}.onnx', providers=providers)
def __call__(self, image):
inputs = process.preprocess(image, size=224, padding_color=127, precision=self.precision)
outputs = self.predict(inputs)
outputs = self.reshape(outputs)
class_index, confidences = process.parse_outputs(outputs)
return self.classes[class_index], f'{confidences[class_index]:.3f}'
@property
def precision(self):
return self.configs['precision']
@property
def classes(self):
return self.configs['classes']
def predict(self, inputs):
return self.session.run(None, {'input': inputs})
@staticmethod
def reshape(outputs):
return outputs[0].squeeze()