Spaces:
Runtime error
Runtime error
File size: 2,027 Bytes
128757a |
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 |
import os
import os.path
import json
from PIL import Image
import torch.utils.data as data
def pil_loader(path):
# open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835)
with open(path, 'rb') as f:
img = Image.open(f)
return img.convert('RGB')
class ImageNet(data.Dataset):
""" ImageNet
Args:
root (string): Root directory where images are downloaded to.
annFile (string): Path to json annotation file.
transform (callable, optional): A function/transform that takes in an PIL image
and returns a transformed version. E.g, ``transforms.ToTensor``
"""
def __init__(self, ann_file, root, remove_images_without_annotations=None, transforms=None):
self.root = root
self.transform = transforms
meta_file = os.path.join(root, ann_file)
assert os.path.exists(meta_file), 'meta file %s under root %s not found' % (os.path.basename(meta_file), root)
with open(meta_file, 'r') as f:
meta = json.load(f)
self.classes = meta['classes']
self.class_to_idx = meta['class_to_idx']
self.samples = meta['samples']
self.num_sample = len(self.samples)
self.allsamples = self.samples
def select_class(self, cls):
new_samples = [sample for sample in self.allsamples if sample[-1] in cls]
self.samples = new_samples
self.num_sample = len(self.samples)
def __getitem__(self, index):
"""
Args:
index (int): Index
Returns:
tuple: (sample, target) where target is class_index of the target class.
"""
img_path, target = self.samples[index]
sample = pil_loader(self.root + '/' + img_path)
if self.transform is not None:
sample = self.transform(sample)
return sample, target, index
def __len__(self):
return len(self.samples) |