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:
            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]