File size: 1,115 Bytes
0199e9f 596a24b 0199e9f |
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 |
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:
if os.path.exists(os.path.join(self._path, '{}.txt'.format(filename[:-4]))):
labels[filename] = np.loadtxt(os.path.join(self._path, '{}.txt'.format(filename[:-4])), ndmin=2)
else:
labels[filename] = None
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] |