File size: 2,566 Bytes
e0b3895 ec17a5b e0b3895 ec17a5b e0b3895 ec17a5b e0b3895 ec17a5b e0b3895 ec17a5b |
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 64 65 66 67 68 69 70 71 72 73 74 75 76 |
# This file is part of OpenCV Zoo project.
# It is subject to the license terms in the LICENSE file found in the same directory.
#
# Copyright (C) 2021, Shenzhen Institute of Artificial Intelligence and Robotics for Society, all rights reserved.
# Third party copyrights are property of their respective owners.
import numpy as np
import cv2 as cv
class PPResNet:
def __init__(self, modelPath, labelPath, backendId=0, targetId=0):
self._modelPath = modelPath
self._labelPath = labelPath
self._backendId = backendId
self._targetId = targetId
self._model = cv.dnn.readNet(self._modelPath)
self._model.setPreferableBackend(self._backendId)
self._model.setPreferableTarget(self._targetId)
self._inputNames = ''
self._outputNames = ['save_infer_model/scale_0.tmp_0']
self._inputSize = [224, 224]
self._mean = np.array([0.485, 0.456, 0.406])[np.newaxis, np.newaxis, :]
self._std = np.array([0.229, 0.224, 0.225])[np.newaxis, np.newaxis, :]
# load labels
self._labels = self._load_labels()
def _load_labels(self):
labels = []
with open(self._labelPath, 'r') as f:
for line in f:
labels.append(line.strip())
return labels
@property
def name(self):
return self.__class__.__name__
def setBackend(self, backend_id):
self._backendId = backend_id
self._model.setPreferableBackend(self._backendId)
def setTarget(self, target_id):
self._targetId = target_id
self._model.setPreferableTarget(self._targetId)
def _preprocess(self, image):
image = image.astype(np.float32, copy=False) / 255.0
image -= self._mean
image /= self._std
return cv.dnn.blobFromImage(image)
def infer(self, image):
assert image.shape[0] == self._inputSize[1], '{} (height of input image) != {} (preset height)'.format(image.shape[0], self._inputSize[1])
assert image.shape[1] == self._inputSize[0], '{} (width of input image) != {} (preset width)'.format(image.shape[1], self._inputSize[0])
# Preprocess
inputBlob = self._preprocess(image)
# Forward
self._model.setInput(inputBlob, self._inputNames)
outputBlob = self._model.forward(self._outputNames)
# Postprocess
results = self._postprocess(outputBlob)
return results
def _postprocess(self, outputBlob):
class_id = np.argmax(outputBlob[0])
return self._labels[class_id]
|