You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
92 lines
3.0 KiB
92 lines
3.0 KiB
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]
|