Add model file collector, fix some bugs and add some features (#123)

add model path searcher and remove hardwritten file paths from benchmark configs; pack ppresnet, mobilenet & crnn with labels; fix palm det data; add flags to enable models of different precision separately (#123)

demo.py

@@ -3,8 +3,7 @@ import argparse
 import numpy as np
 import cv2 as cv
 
-from mobilenet_v1 import MobileNetV1
-from mobilenet_v2 import MobileNetV2
+from mobilenet import MobileNet
 
 def str2bool(v):
     if v.lower() in ['on', 'yes', 'true', 'y', 't']:
@@ -26,24 +25,23 @@ try:
 except:
     print('This version of OpenCV does not support TIM-VX and NPU. Visit https://github.com/opencv/opencv/wiki/TIM-VX-Backend-For-Running-OpenCV-On-NPU for more information.')
 
+all_mobilenets = [
+    'image_classification_mobilenetv1_2022apr.onnx',
+    'image_classification_mobilenetv2_2022apr.onnx',
+    'image_classification_mobilenetv1_2022apr-int8-quantized.onnx',
+    'image_classification_mobilenetv2_2022apr-int8-quantized.onnx'
+]
+
 parser = argparse.ArgumentParser(description='Demo for MobileNet V1 & V2.')
 parser.add_argument('--input', '-i', type=str, help='Usage: Set input path to a certain image, omit if using camera.')
-parser.add_argument('--model', '-m', type=str, choices=['v1', 'v2', 'v1-q', 'v2-q'], default='v1', help='Usage: Set model type, defaults to image_classification_mobilenetv1_2022apr.onnx (v1).')
+parser.add_argument('--model', '-m', type=str, choices=all_mobilenets, default=all_mobilenets[0], help='Usage: Set model type, defaults to image_classification_mobilenetv1_2022apr.onnx (v1).')
 parser.add_argument('--backend', '-b', type=int, default=backends[0], help=help_msg_backends.format(*backends))
 parser.add_argument('--target', '-t', type=int, default=targets[0], help=help_msg_targets.format(*targets))
-parser.add_argument('--label', '-l', type=str, default='./imagenet_labels.txt', help='Usage: Set path to the different labels that will be used during the detection. Default list found in imagenet_labels.txt')
 args = parser.parse_args()
 
 if __name__ == '__main__':
-    # Instantiate ResNet
-    models = {
-        'v1': MobileNetV1(modelPath='./image_classification_mobilenetv1_2022apr.onnx', labelPath=args.label, backendId=args.backend, targetId=args.target),
-        'v2': MobileNetV2(modelPath='./image_classification_mobilenetv2_2022apr.onnx', labelPath=args.label, backendId=args.backend, targetId=args.target),
-        'v1-q': MobileNetV1(modelPath='./image_classification_mobilenetv1_2022apr-int8-quantized.onnx', labelPath=args.label, backendId=args.backend, targetId=args.target),
-        'v2-q': MobileNetV2(modelPath='./image_classification_mobilenetv2_2022apr-int8-quantized.onnx', labelPath=args.label, backendId=args.backend, targetId=args.target)
-
-    }
-    model = models[args.model]
+    # Instantiate MobileNet
+    model = MobileNet(modelPath=args.model, backendId=args.backend, targetId=args.target)
 
     # Read image and get a 224x224 crop from a 256x256 resized
     image = cv.imread(args.input)
@@ -56,4 +54,3 @@ if __name__ == '__main__':
 
     # Print result
     print('label: {}'.format(result))
-

imagenet_labels.txt → mobilenet.py

@@ -1,4 +1,83 @@
-tench
+import numpy as np
+import cv2 as cv
+
+class MobileNet:
+    '''
+    Works with MobileNet V1 & V2.
+    '''
+
+    def __init__(self, modelPath, topK=1, backendId=0, targetId=0):
+        self.model_path = modelPath
+        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):
+        return self.LABELS_IMAGENET_1K.splitlines()
+
+    @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
+
+    LABELS_IMAGENET_1K = '''tench
 goldfish
 great white shark
 tiger shark
@@ -997,4 +1076,4 @@ earthstar
 hen-of-the-woods
 bolete
 ear
-toilet tissue
+toilet tissue'''

mobilenet_v1.py

@@ -1,81 +0,0 @@
-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
-

mobilenet_v2.py

@@ -1,81 +0,0 @@
-import numpy as np
-import cv2 as cv
-
-class MobileNetV2:
-    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
-