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 import tarfile 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] def _extract_tar_info(tarfile): class_to_idx = {} files = [] labels = [] for ti in tarfile.getmembers(): if not ti.isfile(): continue dirname, basename = os.path.split(ti.path) label = os.path.basename(dirname) class_to_idx[label] = None ext = os.path.splitext(basename)[1] if ext.lower() in IMG_EXTENSIONS: files.append(ti) labels.append(label) for idx, c in enumerate(sorted(class_to_idx.keys(), key=natural_key)): class_to_idx[c] = idx tarinfo_and_targets = zip(files, [class_to_idx[l] for l in labels]) tarinfo_and_targets = sorted(tarinfo_and_targets, key=lambda k: natural_key(k[0].path)) return tarinfo_and_targets class DatasetTar(data.Dataset): def __init__(self, root, load_bytes=False, transform=None): assert os.path.isfile(root) self.root = root with tarfile.open(root) as tf: # cannot keep this open across processes, reopen later self.imgs = _extract_tar_info(tf) self.tarfile = None # lazy init in __getitem__ self.load_bytes = load_bytes self.transform = transform def __getitem__(self, index): if self.tarfile is None: self.tarfile = tarfile.open(self.root) tarinfo, target = self.imgs[index] iob = self.tarfile.extractfile(tarinfo) img = iob.read() if self.load_bytes else Image.open(iob).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) class AugMixDataset(torch.utils.data.Dataset): """Dataset wrapper to perform AugMix or other clean/augmentation mixes""" def __init__(self, dataset, num_aug=2): self.augmentation = None self.normalize = None self.dataset = dataset if self.dataset.transform is not None: self._set_transforms(self.dataset.transform) self.num_aug = num_aug def _set_transforms(self, x): assert isinstance(x, (list, tuple)) and len(x) == 3, 'Expecting a tuple/list of 3 transforms' self.dataset.transform = x[0] self.augmentation = x[1] self.normalize = x[2] @property def transform(self): return self.dataset.transform @transform.setter def transform(self, x): self._set_transforms(x) def _normalize(self, x): return x if self.normalize is None else self.normalize(x) def __getitem__(self, i): x, y = self.dataset[i] x_list = [self._normalize(x)] for n in range(self.num_aug): x_list.append(self._normalize(self.augmentation(x))) return tuple(x_list), y def __len__(self): return len(self.dataset)