Train and test prototypical networks on MiniImageNet

Author: jsalbert

dataloaders/__init__.py

@@ -0,0 +1,27 @@
+import torchvision.datasets.folder as datasets
+
+
+class ImageFolderFewShot(datasets.ImageFolder):
+    """A generic data loader where the images are arranged in this way: ::
+        root/dog/xxx.png
+        root/dog/xxy.png
+        root/dog/xxz.png
+        root/cat/123.png
+        root/cat/nsdf3.png
+        root/cat/asd932_.png
+    Args:
+        root (string): Root directory path.
+        transform (callable, optional): A function/transform that  takes in an PIL image
+            and returns a transformed version. E.g, ``transforms.RandomCrop``
+        target_transform (callable, optional): A function/transform that takes in the
+            target and transforms it.
+        loader (callable, optional): A function to load an image given its path.
+     Attributes:
+        classes (list): List of the class names.
+        class_to_idx (dict): Dict with items (class_name, class_index).
+        imgs (list): List of (image path, class_index) tuples
+    """
+
+    def __init__(self, root, transform=None, target_transform=None):
+        super(ImageFolderFewShot, self).__init__(root, transform=transform, target_transform=target_transform)
+        self.labels = [sample[1] for sample in self.samples]

dataloaders/__pycache__/__init__.cpython-35.pyc

@@ -0,0 +1,52 @@
+import os.path as osp
+from PIL import Image
+
+from torch.utils.data import Dataset
+from torchvision import transforms
+
+
+ROOT_PATH = './mini_imagenet/'
+
+
+class MiniImageNet(Dataset):
+
+    def __init__(self, setname):
+        csv_path = osp.join(ROOT_PATH, setname + '.csv')
+        lines = [x.strip() for x in open(csv_path, 'r').readlines()][1:]
+
+        data = []
+        label = []
+        lb = -1
+
+        self.wnids = []
+        self.class_idx_to_sample_idx = {}
+
+        for idx, l in enumerate(lines):
+            name, wnid = l.split(',')
+            path = osp.join(ROOT_PATH, 'images', name)
+            if wnid not in self.wnids:
+                self.wnids.append(wnid)
+                lb += 1
+                self.class_idx_to_sample_idx.update({lb: []})
+            data.append(path)
+            label.append(lb)
+            self.class_idx_to_sample_idx[lb].append(idx)
+
+        self.data = data
+        self.labels = label
+
+        self.transform = transforms.Compose([
+            transforms.Resize(84),
+            transforms.CenterCrop(84),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                 std=[0.229, 0.224, 0.225])
+        ])
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, i):
+        path, label = self.data[i], self.labels[i]
+        image = self.transform(Image.open(path).convert('RGB'))
+        return image, label