JeVoisBase  1.18
JeVois Smart Embedded Machine Vision Toolkit Base Modules
Share this page:
classification.py
Go to the documentation of this file.
1 import os
2 
3 import numpy as np
4 import cv2 as cv
5 
6 from .base_dataloader import _BaseImageLoader
7 from ..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
utils.dataloaders.classification.ClassificationImageLoader._center_crop
_center_crop
Definition: classification.py:15
demo.int
int
Definition: demo.py:37
utils.dataloaders.classification.ClassificationImageLoader.__iter__
def __iter__(self)
Definition: classification.py:28
utils.dataloaders.classification.ClassificationImageLoader.__init__
def __init__(self, **kwargs)
Definition: classification.py:11
utils.dataloaders.base_dataloader._BaseImageLoader._files
_files
Definition: base_dataloader.py:10
utils.dataloaders.base_dataloader._BaseImageLoader
Definition: base_dataloader.py:5
utils.dataloaders.base_dataloader._BaseImageLoader._sizes
_sizes
Definition: base_dataloader.py:14
utils.dataloaders.classification.ClassificationImageLoader._centerCrop
def _centerCrop(self, image)
Definition: classification.py:20
utils.dataloaders.classification.ClassificationImageLoader
Definition: classification.py:10
utils.dataloaders.base_dataloader._BaseImageLoader._path
_path
Definition: base_dataloader.py:7
utils.dataloaders.classification.ClassificationImageLoader._to_rgb
_to_rgb
Definition: classification.py:14
utils.dataloaders.classification.ClassificationImageLoader._toRGB
def _toRGB(self, image)
Definition: classification.py:17