Decoupling dataloader and benchmark runner (#16)

* impl dataloader factory

* set a type to init all instances for benchmarks

benchmark/benchmark.py

@@ -7,17 +7,23 @@ import cv2 as cv
 
 # from ..models import MODELS
 from models import MODELS
-from utils import METRICS
+from utils import METRICS, DATALOADERS
 
 parser = argparse.ArgumentParser("Benchmarks for OpenCV Zoo.")
 parser.add_argument('--cfg', '-c', type=str,
                     help='Benchmarking on the given config.')
 args = parser.parse_args()
 
-def build_from_cfg(cfg, registery, key='name'):
-    obj_name = cfg.pop(key)
-    obj = registery.get(obj_name)
-    return obj(**cfg)
+def build_from_cfg(cfg, registery, key=None, name=None):
+    if key is not None:
+        obj_name = cfg.pop(key)
+        obj = registery.get(obj_name)
+        return obj(**cfg)
+    elif name is not None:
+        obj = registery.get(name)
+        return obj(**cfg)
+    else:
+        raise NotImplementedError()
 
 def prepend_pythonpath(cfg):
     for k, v in cfg.items():
@@ -27,62 +33,26 @@ def prepend_pythonpath(cfg):
             if 'path' in k.lower():
                 cfg[k] = os.path.join(os.environ['PYTHONPATH'], v)
 
-class Data:
-    def __init__(self, **kwargs):
-        self._path = kwargs.pop('path', None)
-        assert self._path, 'Benchmark[\'data\'][\'path\'] cannot be empty.'
-
-        self._files = kwargs.pop('files', None)
-        if not self._files:
-            print('Benchmark[\'data\'][\'files\'] is empty, loading all images by default.')
-            self._files = list()
-            for filename in os.listdir(self._path):
-                if filename.endswith('jpg') or filename.endswith('png'):
-                    self._files.append(filename)
-
-        self._use_label = kwargs.pop('useLabel', False)
-        if self._use_label:
-            self._labels = self._load_label()
-
-        self._to_rgb = kwargs.pop('toRGB', False)
-        self._resize = tuple(kwargs.pop('resize', []))
-        self._center_crop = kwargs.pop('centerCrop', None)
-
-    def _load_label(self):
-        labels = dict.fromkeys(self._files, None)
-        for filename in self._files:
-            labels[filename] = np.loadtxt(os.path.join(self._path, '{}.txt'.format(filename[:-4])), ndmin=2)
-        return labels
-
-    def __getitem__(self, idx):
-        image = cv.imread(os.path.join(self._path, self._files[idx]))
-
-        if self._to_rgb:
-            image = cv.cvtColor(image, cv.COLOR_BGR2RGB)
-        if self._resize:
-            image = cv.resize(image, self._resize)
-        if self._center_crop:
-            h, w, _ = image.shape
-            w_crop = int((w - self._center_crop) / 2.)
-            assert w_crop >= 0
-            h_crop = int((h - self._center_crop) / 2.)
-            assert h_crop >= 0
-            image = image[w_crop:w-w_crop, h_crop:h-h_crop, :]
-
-        if self._use_label:
-            return self._files[idx], image, self._labels[self._files[idx]]
-        else:
-            return self._files[idx], image
-
 class Benchmark:
     def __init__(self, **kwargs):
+        self._type = kwargs.pop('type', None)
+        if self._type is None:
+            self._type = 'Base'
+            print('Benchmark[\'type\'] is omitted, set to \'Base\' by default.')
+
         self._data_dict = kwargs.pop('data', None)
         assert self._data_dict, 'Benchmark[\'data\'] cannot be empty and must have path and files.'
-        self._data = Data(**self._data_dict)
+        if 'type' in self._data_dict:
+            self._dataloader = build_from_cfg(self._data_dict, registery=DATALOADERS, key='type')
+        else:
+            self._dataloader = build_from_cfg(self._data_dict, registery=DATALOADERS, name=self._type)
 
         self._metric_dict = kwargs.pop('metric', None)
-        # self._metric = Metric(**self._metric_dict)
-        self._metric = build_from_cfg(self._metric_dict, registery=METRICS, key='type')
+        assert self._metric_dict, 'Benchmark[\'metric\'] cannot be empty.'
+        if 'type' in self._metric_dict:
+            self._metric = build_from_cfg(self._metric_dict, registery=METRICS, key='type')
+        else:
+            self._metric = build_from_cfg(self._metric_dict, registery=METRICS, name=self._type)
 
         backend_id = kwargs.pop('backend', 'default')
         available_backends = dict(
@@ -115,8 +85,15 @@ class Benchmark:
         model.setBackend(self._backend)
         model.setTarget(self._target)
 
-        for data in self._data:
-            self._benchmark_results[data[0]] = self._metric.forward(model, *data[1:])
+        if 'video' in self._dataloader.name.lower():
+            model.init(self._dataloader.getROI())
+
+        for data in self._dataloader:
+            filename, img = data[:2]
+            size = [img.shape[1], img.shape[0]]
+            if filename not in self._benchmark_results:
+                self._benchmark_results[filename] = dict()
+            self._benchmark_results[filename][str(size)] = self._metric.forward(model, *data[1:])
 
     def printResults(self):
         for imgName, results in self._benchmark_results.items():
@@ -138,7 +115,7 @@ if __name__ == '__main__':
     benchmark = Benchmark(**cfg['Benchmark'])
 
     # Instantiate model
-    model = build_from_cfg(cfg=cfg['Model'], registery=MODELS)
+    model = build_from_cfg(cfg=cfg['Model'], registery=MODELS, key='name')
 
     # Run benchmarking
     print('Benchmarking {}:'.format(model.name))

benchmark/config/face_detection_yunet.yaml

 Benchmark:
   name: "Face Detection Benchmark"
+  type: "Detection"
   data:
     path: "benchmark/data/face/detection"
     files: ["group.jpg", "concerts.jpg", "dance.jpg"]
-  metric:
-    type: "Detection"
     sizes: # [[w1, h1], ...], Omit to run at original scale
       - [160, 120]
       - [640, 480]
+  metric:
     warmup: 30
     repeat: 10
     reduction: "median"

benchmark/config/face_recognition_sface.yaml

 Benchmark:
   name: "Face Recognition Benchmark"
+  type: "Recognition"
   data:
     path: "benchmark/data/face/recognition"
     files: ["Aaron_Tippin_0001.jpg", "Alvaro_Uribe_0028.jpg", "Alvaro_Uribe_0029.jpg", "Jose_Luis_Rodriguez_Zapatero_0001.jpg"]
-    useLabel: True
   metric: # 'sizes' is omitted since this model requires input of fixed size
-    type: "Recognition"
     warmup: 30
     repeat: 10
     reduction: "median"

benchmark/config/human_segmentation_pphumanseg.yaml

 Benchmark:
   name: "Human Segmentation Benchmark"
+  type: "Base"
   data:
     path: "benchmark/data/human_segmentation"
     files: ["messi5.jpg", "100040721_1.jpg", "detect.jpg"]
+    sizes: [[192, 192]]
     toRGB: True
-    resize: [192, 192]
   metric:
-    type: "Base"
     warmup: 30
     repeat: 10
     reduction: "median"

benchmark/config/image_classification_ppresnet.yaml

 Benchmark:
   name: "Image Classification Benchmark"
+  type: "Classification"
   data:
     path: "benchmark/data/image_classification"
     files: ["coffee_mug.jpg", "umbrella.jpg", "wall_clock.jpg"]
+    sizes: [[256, 256]]
     toRGB: True
-    resize: [256, 256]
     centerCrop: 224
   metric:
-    type: "Base"
     warmup: 30
     repeat: 10
     reduction: "median"

benchmark/config/qrcode_wechatqrcode.yaml

 Benchmark:
   name: "QRCode Detection and Decoding Benchmark"
+  type: "Detection"
   data:
     path: "benchmark/data/qrcode"
     files: ["opencv.png", "opencv_zoo.png"]
-  metric:
-    type: "Detection"
     sizes:
       - [100, 100]
       - [300, 300]
+  metric:
     warmup: 30
     repeat: 10
     reduction: "median"

benchmark/config/text_detection_db.yaml

 Benchmark:
   name: "Text Detection Benchmark"
+  type: "Detection"
   data:
     path: "benchmark/data/text"
     files: ["1.jpg", "2.jpg", "3.jpg"]
-  metric:
-    type: "Detection"
     sizes: # [[w1, h1], ...], Omit to run at original scale
       - [640, 480]
+  metric:
     warmup: 30
     repeat: 10
     reduction: "median"

benchmark/config/text_recognition_crnn.yaml

 Benchmark:
   name: "Text Recognition Benchmark"
+  type: "Recognition"
   data:
     path: "benchmark/data/text"
     files: ["1.jpg", "2.jpg", "3.jpg"]
-    useLabel: True
   metric: # 'sizes' is omitted since this model requires input of fixed size
-    type: "Recognition"
     warmup: 30
     repeat: 10
     reduction: "median"

benchmark/utils/__init__.py

 from .factory import (METRICS, DATALOADERS)
 from .metrics import *
+from .dataloaders import *
 
 __all__ = ['METRICS', 'DATALOADERS']
\ No newline at end of file

benchmark/utils/dataloaders/__init__.py

+from .base import BaseImageLoader, BaseVideoLoader
+from .classification import ClassificationImageLoader
+from .recognition import RecognitionImageLoader
+from .tracking import TrackingVideoLoader
+
+__all__ = ['BaseImageLoader', 'BaseVideoLoader', 'ClassificationImageLoader', 'RecognitionImageLoader', 'TrackingVideoLoader']
\ No newline at end of file

benchmark/utils/dataloaders/base.py

+from .base_dataloader import _BaseImageLoader, _BaseVideoLoader
+from ..factory import DATALOADERS
+
+@DATALOADERS.register
+class BaseImageLoader(_BaseImageLoader):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+@DATALOADERS.register
+class BaseVideoLoader(_BaseVideoLoader):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
\ No newline at end of file

benchmark/utils/dataloaders/base_dataloader.py

+import os
+
+import cv2 as cv
+
+class _BaseImageLoader:
+    def __init__(self, **kwargs):
+        self._path = kwargs.pop('path', None)
+        assert self._path, 'Benchmark[\'data\'][\'path\'] cannot be empty.'
+
+        self._files = kwargs.pop('files', None)
+        assert self._files, 'Benchmark[\'data\'][\'files\'] cannot be empty'
+        self._len_files = len(self._files)
+
+        self._sizes = kwargs.pop('sizes', [[0, 0]])
+        self._len_sizes = len(self._sizes)
+
+    @property
+    def name(self):
+        return self.__class__.__name__
+
+    def __len__(self):
+        return self._len_files * self._len_sizes
+
+    def __iter__(self):
+        for filename in self._files:
+            image = cv.imread(os.path.join(self._path, filename))
+            if [0, 0] in self._sizes:
+                yield filename, image
+            else:
+                for size in self._sizes:
+                    image_r = cv.resize(image, size)
+                    yield filename, image_r
+
+class _VideoStream:
+    def __init__(self, filepath):
+        self._filepath = filepath
+        self._video = cv.VideoCapture(filepath)
+
+    def __iter__(self):
+        while True:
+            has_frame, frame = self._video.read()
+            if has_frame:
+                yield frame
+            else:
+                break
+
+    def reload(self):
+        self._video = cv.VideoCapture(filepath)
+
+
+class _BaseVideoLoader:
+    def __init__(self, **kwargs):
+        self._path = kwargs.pop('path', None)
+        assert self._path, 'Benchmark[\'data\'][\'path\'] cannot be empty.'
+
+        self._files = kwargs.pop('files', None)
+        assert self._files,'Benchmark[\'data\'][\'files\'] cannot be empty.'
+
+    @property
+    def name(self):
+        return self.__class__.__name__
+
+    def __len__(self):
+        return len(self._files)
+
+    def __getitem__(self, idx):
+        return self._files[idx], _VideoStream(os.path.join(self._path, self._files[idx]))
\ No newline at end of file

benchmark/utils/dataloaders/classification.py

+import os
+
+import numpy as np
+import cv2 as cv
+
+from .base_dataloader import _BaseImageLoader
+from ..factory import DATALOADERS
+
+@DATALOADERS.register
+class ClassificationImageLoader(_BaseImageLoader):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+        self._to_rgb = kwargs.pop('toRGB', False)
+        self._center_crop = kwargs.pop('centerCrop', None)
+
+    def _toRGB(self, image):
+        return cv.cvtColor(image, cv.COLOR_BGR2RGB)
+
+    def _centerCrop(self, image):
+        h, w, _ = image.shape
+        w_crop = int((w - self._center_crop) / 2.)
+        assert w_crop >= 0
+        h_crop = int((h - self._center_crop) / 2.)
+        assert h_crop >= 0
+        return image[w_crop:w-w_crop, h_crop:h-h_crop, :]
+
+    def __iter__(self):
+        for filename in self._files:
+            image = cv.imread(os.path.join(self._path, filename))
+
+            if self._to_rgb:
+                image = self._toRGB(image)
+
+            if [0, 0] in self._sizes:
+                yield filename, image
+            else:
+                for size in self._sizes:
+                    image = cv.resize(image, size)
+                    if self._center_crop:
+                        image = self._centerCrop(image)
+                    yield filename, image
\ No newline at end of file

benchmark/utils/dataloaders/recognition.py

+import os
+
+import numpy as np
+import cv2 as cv
+
+from .base_dataloader import _BaseImageLoader
+from ..factory import DATALOADERS
+
+@DATALOADERS.register
+class RecognitionImageLoader(_BaseImageLoader):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+        self._labels = self._load_label()
+
+    def _load_label(self):
+        labels = dict.fromkeys(self._files, None)
+        for filename in self._files:
+            labels[filename] = np.loadtxt(os.path.join(self._path, '{}.txt'.format(filename[:-4])), ndmin=2)
+        return labels
+
+    def __iter__(self):
+        for filename in self._files:
+            image = cv.imread(os.path.join(self._path, filename))
+            if [0, 0] in self._sizes:
+                yield filename, image, self._labels[filename]
+            else:
+                for size in self._sizes:
+                    image_r = cv.resize(image, size)
+                    yield filename, image_r, self._labels[filename]
\ No newline at end of file

benchmark/utils/dataloaders/tracking.py

+import numpy as np
+
+from .base_dataloader import _BaseVideoLoader
+from ..factory import DATALOADERS
+
+@DATALOADERS.register
+class TrackingVideoLoader(_BaseVideoLoader):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+        self._rois = self._load_roi()
+
+        unsupported_keys = []
+        for k, _ in kwargs.items():
+            unsupported_keys.append(k)
+        print('Keys ({}) are not supported in Benchmark[\'data\'].'.format(str(unsupported_keys)))
+
+    def _load_roi(self):
+        rois = dict.fromkeys(self._files, None)
+        for filename in self._files:
+            rois[filename] = np.loadtxt(os.path.join(self._path, '{}.txt'.format(filename[:-4])), ndmin=2)
+        return rois
+
+    def getROI(self):
+        return self._rois
\ No newline at end of file

benchmark/utils/factory.py

@@ -4,10 +4,17 @@ class Registery:
         self._dict = dict()
 
     def get(self, key):
-        return self._dict[key]
+        if key in self._dict:
+            return self._dict[key]
+        else:
+            return self._dict['Base']
 
     def register(self, item):
         self._dict[item.__name__] = item
+        # renaming *ImageLoader/*VideoLoader
+        if 'ImageLoader' in item.__name__:
+            name = item.__name__.replace('ImageLoader', '')
+            self._dict[name] = item
 
 METRICS = Registery('Metrics')
 DATALOADERS = Registery('DataLoaders')
\ No newline at end of file

benchmark/utils/metrics/base.py

@@ -10,20 +10,15 @@ class Base(BaseMetric):
 
     def forward(self, model, *args, **kwargs):
         img = args[0]
-        if not self._sizes:
-            h, w, _ = img.shape
-            self._sizes.append([w, h])
 
-        results = dict()
+        # warmup
+        for _ in range(self._warmup):
+            model.infer(img)
+        # repeat
         self._timer.reset()
-        for size in self._sizes:
-            input_data = cv.resize(img, size)
-            for _ in range(self._warmup):
-                model.infer(input_data)
-            for _ in range(self._repeat):
-                self._timer.start()
-                model.infer(input_data)
-                self._timer.stop()
-            results[str(size)] = self._getResult()
+        for _ in range(self._repeat):
+            self._timer.start()
+            model.infer(img)
+            self._timer.stop()
 
-        return results
\ No newline at end of file
+        return self._getResult()
\ No newline at end of file

benchmark/utils/metrics/base_metric.py

@@ -4,9 +4,6 @@ from ..timer import Timer
 
 class BaseMetric:
     def __init__(self, **kwargs):
-        self._sizes = kwargs.pop('sizes', None)
-        if self._sizes is None:
-            self._sizes = []
         self._warmup = kwargs.pop('warmup', 3)
         self._repeat = kwargs.pop('repeat', 10)
         self._reduction = kwargs.pop('reduction', 'median')

benchmark/utils/metrics/detection.py

@@ -10,24 +10,20 @@ class Detection(BaseMetric):
 
     def forward(self, model, *args, **kwargs):
         img = args[0]
-        if not self._sizes:
-            h, w, _ = img.shape
-            self._sizes.append([w, h])
+        size = [img.shape[1], img.shape[0]]
+        try:
+            model.setInputSize(size)
+        except:
+            pass
 
-        results = dict()
+        # warmup
+        for _ in range(self._warmup):
+            model.infer(img)
+        # repeat
         self._timer.reset()
-        for size in self._sizes:
-            input_data = cv.resize(img, size)
-            try:
-                model.setInputSize(size)
-            except:
-                pass
-            for _ in range(self._warmup):
-                model.infer(input_data)
-            for _ in range(self._repeat):
-                self._timer.start()
-                model.infer(input_data)
-                self._timer.stop()
-            results[str(size)] = self._getResult()
+        for _ in range(self._repeat):
+            self._timer.start()
+            model.infer(img)
+            self._timer.stop()
 
-        return results
\ No newline at end of file
+        return self._getResult()
\ No newline at end of file

benchmark/utils/metrics/recognition.py

@@ -10,11 +10,7 @@ class Recognition(BaseMetric):
 
     def forward(self, model, *args, **kwargs):
         img, bboxes = args
-        if not self._sizes:
-            h, w, _ = img.shape
-            self._sizes.append([w, h])
 
-        results = dict()
         self._timer.reset()
         for idx, bbox in enumerate(bboxes):
             for _ in range(self._warmup):
@@ -23,6 +19,5 @@ class Recognition(BaseMetric):
                 self._timer.start()
                 model.infer(img, bbox)
                 self._timer.stop()
-            results['bbox{}'.format(idx)] = self._getResult()
 
-        return results
\ No newline at end of file
+        return self._getResult()
\ No newline at end of file