File size: 2,590 Bytes
d33294b 28adb60 d33294b 28adb60 d33294b 86cd32d d33294b 90edc6d d33294b 28adb60 d33294b 86cd32d d33294b 28adb60 d33294b 28adb60 90edc6d 28adb60 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 71 72 73 74 75 76 77 78 79 80 81 82 |
import numpy as np
import cv2 as cv
class MobileNetV1:
def __init__(self, modelPath, labelPath=None, topK=1, backendId=0, targetId=0):
self.model_path = modelPath
self.label_path = labelPath
assert topK >= 1
self.top_k = topK
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 = []
if self.label_path is not None:
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):
batched_class_id_list = []
for o in output_blob:
class_id_list = o.argsort()[::-1][:self.top_k]
batched_class_id_list.append(class_id_list)
if len(self._labels) > 0:
batched_predicted_labels = []
for class_id_list in batched_class_id_list:
predicted_labels = []
for class_id in class_id_list:
predicted_labels.append(self._labels[class_id])
batched_predicted_labels.append(predicted_labels)
return batched_predicted_labels
else:
return batched_class_id_list
|