JeVoisBase  1.21
JeVois Smart Embedded Machine Vision Toolkit Base Modules
Share this page:
Loading...
Searching...
No Matches
classification.py
Go to the documentation of this file.
1import os
2
3import numpy as np
4import cv2 as cv
5
6from .base_dataloader import _BaseImageLoader
7from ..factory import DATALOADERS
8
9@DATALOADERS.register
11 def __init__(self, **kwargs):
12 super().__init__(**kwargs)
13
14 self._to_rgb = kwargs.pop('toRGB', False)
15 self._center_crop = kwargs.pop('centerCrop', None)
16
17 def _toRGB(self, image):
18 return cv.cvtColor(image, cv.COLOR_BGR2RGB)
19
20 def _centerCrop(self, image):
21 h, w, _ = image.shape
22 w_crop = int((w - self._center_crop) / 2.)
23 assert w_crop >= 0
24 h_crop = int((h - self._center_crop) / 2.)
25 assert h_crop >= 0
26 return image[w_crop:w-w_crop, h_crop:h-h_crop, :]
27
28 def __iter__(self):
29 for filename in self._files:
30 image = cv.imread(os.path.join(self._path, filename))
31
32 if self._to_rgb:
33 image = self._toRGB(image)
34
35 if [0, 0] in self._sizes:
36 yield filename, image
37 else:
38 for size in self._sizes:
39 image = cv.resize(image, size)
40 if self._center_crop:
41 image = self._centerCrop(image)
42 yield filename, image