Spaces:
Build error
Build error
File size: 4,216 Bytes
8f87579 |
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 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 |
'''
Variants of pytorch's ImageFolder for loading image datasets with more
information, such as parallel feature channels in separate files,
cached files with lists of filenames, etc.
'''
import os, torch, re
import torch.utils.data as data
from torchvision.datasets.folder import default_loader
from PIL import Image
from collections import OrderedDict
from .progress import default_progress
def grayscale_loader(path):
with open(path, 'rb') as f:
return Image.open(f).convert('L')
class ParallelImageFolders(data.Dataset):
"""
A data loader that looks for parallel image filenames, for example
photo1/park/004234.jpg
photo1/park/004236.jpg
photo1/park/004237.jpg
photo2/park/004234.png
photo2/park/004236.png
photo2/park/004237.png
"""
def __init__(self, image_roots,
transform=None,
loader=default_loader,
stacker=None,
intersection=False,
verbose=None,
size=None):
self.image_roots = image_roots
self.images = make_parallel_dataset(image_roots,
intersection=intersection, verbose=verbose)
if len(self.images) == 0:
raise RuntimeError("Found 0 images within: %s" % image_roots)
if size is not None:
self.image = self.images[:size]
if transform is not None and not hasattr(transform, '__iter__'):
transform = [transform for _ in image_roots]
self.transforms = transform
self.stacker = stacker
self.loader = loader
def __getitem__(self, index):
paths = self.images[index]
sources = [self.loader(path) for path in paths]
# Add a common shared state dict to allow random crops/flips to be
# coordinated.
shared_state = {}
for s in sources:
s.shared_state = shared_state
if self.transforms is not None:
sources = [transform(source)
for source, transform in zip(sources, self.transforms)]
if self.stacker is not None:
sources = self.stacker(sources)
else:
sources = tuple(sources)
return sources
def __len__(self):
return len(self.images)
def is_npy_file(path):
return path.endswith('.npy') or path.endswith('.NPY')
def is_image_file(path):
return None != re.search(r'\.(jpe?g|png)$', path, re.IGNORECASE)
def walk_image_files(rootdir, verbose=None):
progress = default_progress(verbose)
indexfile = '%s.txt' % rootdir
if os.path.isfile(indexfile):
basedir = os.path.dirname(rootdir)
with open(indexfile) as f:
result = sorted([os.path.join(basedir, line.strip())
for line in progress(f.readlines(),
desc='Reading %s' % os.path.basename(indexfile))])
return result
result = []
for dirname, _, fnames in sorted(progress(os.walk(rootdir),
desc='Walking %s' % os.path.basename(rootdir))):
for fname in sorted(fnames):
if is_image_file(fname) or is_npy_file(fname):
result.append(os.path.join(dirname, fname))
return result
def make_parallel_dataset(image_roots, intersection=False, verbose=None):
"""
Returns [(img1, img2), (img1, img2)..]
"""
image_roots = [os.path.expanduser(d) for d in image_roots]
image_sets = OrderedDict()
for j, root in enumerate(image_roots):
for path in walk_image_files(root, verbose=verbose):
key = os.path.splitext(os.path.relpath(path, root))[0]
if key not in image_sets:
image_sets[key] = []
if not intersection and len(image_sets[key]) != j:
raise RuntimeError(
'Images not parallel: %s missing from one dir' % (key))
image_sets[key].append(path)
tuples = []
for key, value in image_sets.items():
if len(value) != len(image_roots):
if intersection:
continue
else:
raise RuntimeError(
'Images not parallel: %s missing from one dir' % (key))
tuples.append(tuple(value))
return tuples
|