File size: 2,030 Bytes
d33294b 86cd32d d33294b 86cd32d d33294b |
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 |
import numpy as np
import cv2 as cv
class MobileNetV1:
def __init__(self, modelPath, labelPath, backendId=0, targetId=0):
self.model_path = modelPath
self.label_path = labelPath
self.backend_id = backendId
self.target_id = targetId
self.model = cv.dnn.readNet(self.model_path)
self.model.setPreferableBackend(self.backend_id)
self.model.setPreferableTarget(self.target_id)
self.input_names = ''
self.output_names = ''
self.input_size = [224, 224]
self.mean=[0.485, 0.456, 0.406]
self.std=[0.229, 0.224, 0.225]
# load labels
self.labels = self._load_labels()
def _load_labels(self):
labels = []
with open(self.label_path, 'r') as f:
for line in f:
labels.append(line.strip())
return labels
@property
def name(self):
return self.__class__.__name__
def setBackend(self, backendId):
self.backend_id = backendId
self.model.setPreferableBackend(self.backend_id)
def setTarget(self, targetId):
self.target_id = targetId
self.model.setPreferableTarget(self.target_id)
def _preprocess(self, image):
input_blob = (image / 255.0 - self.mean) / self.std
input_blob = input_blob.transpose(2, 0, 1)
input_blob = input_blob[np.newaxis, :, :, :]
input_blob = input_blob.astype(np.float32)
return input_blob
def infer(self, image):
# Preprocess
input_blob = self._preprocess(image)
# Forward
self.model.setInput(input_blob, self.input_names)
output_blob = self.model.forward(self.output_names)
# Postprocess
results = self._postprocess(output_blob)
return results
def _postprocess(self, output_blob):
predicted_labels = []
for o in output_blob:
class_id = np.argmax(o)
predicted_labels.append(self.labels[class_id])
return predicted_labels
|