from __future__ import absolute_import from __future__ import division from __future__ import print_function import torch.utils.data as data import os import re import torch from PIL import Image IMG_EXTENSIONS = ['.png', '.jpg', '.jpeg'] def natural_key(string_): """See http://www.codinghorror.com/blog/archives/001018.html""" return [int(s) if s.isdigit() else s for s in re.split(r'(\d+)', string_.lower())] def find_images_and_targets(folder, types=IMG_EXTENSIONS, class_to_idx=None, leaf_name_only=True, sort=True): if class_to_idx is None: class_to_idx = dict() build_class_idx = True else: build_class_idx = False labels = [] filenames = [] for root, subdirs, files in os.walk(folder, topdown=False): rel_path = os.path.relpath(root, folder) if (root != folder) else '' label = os.path.basename(rel_path) if leaf_name_only else rel_path.replace(os.path.sep, '_') if build_class_idx and not subdirs: class_to_idx[label] = None for f in files: base, ext = os.path.splitext(f) if ext.lower() in types: filenames.append(os.path.join(root, f)) labels.append(label) if build_class_idx: classes = sorted(class_to_idx.keys(), key=natural_key) for idx, c in enumerate(classes): class_to_idx[c] = idx images_and_targets = zip(filenames, [class_to_idx[l] for l in labels]) if sort: images_and_targets = sorted(images_and_targets, key=lambda k: natural_key(k[0])) if build_class_idx: return images_and_targets, classes, class_to_idx else: return images_and_targets class Dataset(data.Dataset): def __init__( self, root, load_bytes=False, transform=None): imgs, _, _ = find_images_and_targets(root) if len(imgs) == 0: raise(RuntimeError("Found 0 images in subfolders of: " + root + "\n" "Supported image extensions are: " + ",".join(IMG_EXTENSIONS))) self.root = root self.imgs = imgs self.load_bytes = load_bytes self.transform = transform def __getitem__(self, index): path, target = self.imgs[index] img = open(path, 'rb').read() if self.load_bytes else Image.open(path).convert('RGB') if self.transform is not None: img = self.transform(img) if target is None: target = torch.zeros(1).long() return img, target def __len__(self): return len(self.imgs) def filenames(self, indices=[], basename=False): if indices: if basename: return [os.path.basename(self.imgs[i][0]) for i in indices] else: return [self.imgs[i][0] for i in indices] else: if basename: return [os.path.basename(x[0]) for x in self.imgs] else: return [x[0] for x in self.imgs]