File size: 1,225 Bytes
35cdf61 |
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 |
import logging
import glob
import torch
class FolderData(torch.utils.data.Dataset):
def __init__(self, path, transforms, extensions=['.jpg', '.png'], recursive=False, verbose=False):
self.verbose = verbose
if self.verbose:
logger = logging.getLogger(__name__)
if len(extensions) == 0:
if self.verbose:
logger.error("Expected at list one extension, but none was received.")
raise ValueError
if self.verbose:
logger.info("Constructing the list of images.")
additional_pattern = '/**/*' if recursive else '/*'
files = []
for extension in extensions:
files.extend(glob.glob(path + additional_pattern + extension, recursive=recursive))
if self.verbose:
logger.info("Finished searching for images. %s images found", len(files))
logger.info("Preparing to run the detection.")
self.files = files
self.transforms = transforms
def __getitem__(self, idx):
image_path = self.files[idx]
image = self.transforms(image_path)
return image_path, image
def __len__(self):
return len(self.files) |