Spaces:
Sleeping
Sleeping
| 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) | |