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.
pytorch-image-models/timm/data/dataset.py

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]