JeVoisBase  1.22
JeVois Smart Embedded Machine Vision Toolkit Base Modules
Share this page:
Loading...
Searching...
No Matches
ppresnet.py
Go to the documentation of this file.
1# This file is part of OpenCV Zoo project.
2# It is subject to the license terms in the LICENSE file found in the same directory.
3#
4# Copyright (C) 2021, Shenzhen Institute of Artificial Intelligence and Robotics for Society, all rights reserved.
5# Third party copyrights are property of their respective owners.
6
7
8import numpy as np
9import cv2 as cv
10
12 def __init__(self, modelPath, labelPath, backendId=0, targetId=0):
13 self._modelPath = modelPath
14 self._labelPath = labelPath
15 self._backendId = backendId
16 self._targetId = targetId
17
18 self._model = cv.dnn.readNet(self._modelPath)
19 self._model.setPreferableBackend(self._backendId)
20 self._model.setPreferableTarget(self._targetId)
21
22 self._inputNames = ''
23 self._outputNames = ['save_infer_model/scale_0.tmp_0']
24 self._inputSize = [224, 224]
25 self._mean = np.array([0.485, 0.456, 0.406])[np.newaxis, np.newaxis, :]
26 self._std = np.array([0.229, 0.224, 0.225])[np.newaxis, np.newaxis, :]
27
28 # load labels
29 self._labels = self._load_labels()
30
31 def _load_labels(self):
32 labels = []
33 with open(self._labelPath, 'r') as f:
34 for line in f:
35 labels.append(line.strip())
36 return labels
37
38 @property
39 def name(self):
40 return self.__class__.__name__
41
42 def setBackend(self, backend_id):
43 self._backendId = backend_id
44 self._model.setPreferableBackend(self._backendId)
45
46 def setTarget(self, target_id):
47 self._targetId = target_id
48 self._model.setPreferableTarget(self._targetId)
49
50 def _preprocess(self, image):
51 image = image.astype(np.float32, copy=False) / 255.0
52 image -= self._mean
53 image /= self._std
54 return cv.dnn.blobFromImage(image)
55
56 def infer(self, image):
57 assert image.shape[0] == self._inputSize[1], '{} (height of input image) != {} (preset height)'.format(image.shape[0], self._inputSize[1])
58 assert image.shape[1] == self._inputSize[0], '{} (width of input image) != {} (preset width)'.format(image.shape[1], self._inputSize[0])
59
60 # Preprocess
61 inputBlob = self._preprocess(image)
62
63 # Forward
64 self._model.setInput(inputBlob, self._inputNames)
65 outputBlob = self._model.forward(self._outputNames)
66
67 # Postprocess
68 results = self._postprocess(outputBlob)
69
70 return results
71
72 def _postprocess(self, outputBlob):
73 class_id = np.argmax(outputBlob[0])
74 return self._labels[class_id]
75
_postprocess(self, outputBlob)
Definition ppresnet.py:72
setBackend(self, backend_id)
Definition ppresnet.py:42
__init__(self, modelPath, labelPath, backendId=0, targetId=0)
Definition ppresnet.py:12
setTarget(self, target_id)
Definition ppresnet.py:46
_preprocess(self, image)
Definition ppresnet.py:50
infer(self, image)
Definition ppresnet.py:56