| 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(self._filepath) |
|
|
| def __iter__(self): |
| while True: |
| has_frame, frame = self._video.read() |
| if has_frame: |
| yield frame |
| else: |
| break |
|
|
| def __next__(self): |
| while True: |
| has_frame, frame = self._video.read() |
| if has_frame: |
| return frame |
| else: |
| break |
|
|
| def reload(self): |
| self._video = cv.VideoCapture(self._filepath) |
|
|
| def getFrameSize(self): |
| w = int(self._video.get(cv.CAP_PROP_FRAME_WIDTH)) |
| h = int(self._video.get(cv.CAP_PROP_FRAME_HEIGHT)) |
| return [w, h] |
|
|
|
|
| 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.' |
|
|
| self._streams = dict() |
| for filename in self._files: |
| self._streams[filename] = _VideoStream(os.path.join(self._path, filename)) |
|
|
| @property |
| def name(self): |
| return self.__class__.__name__ |
|
|
| def __len__(self): |
| return len(self._files) |
|
|
| def __getitem__(self, idx): |
| return self._files[idx], self._streams[idx] |
