File size: 2,426 Bytes
0199e9f 23d8387 0199e9f 23d8387 0199e9f 23d8387 0199e9f 23d8387 0199e9f 23d8387 |
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 83 84 |
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] |