Class TinyImageNet

tensorpack/dataflow/dataset/ilsvrc.py:293–348

The TinyImageNet classification dataset, with 200 classes and 500 images per class. See https://tiny-imagenet.herokuapp.com/. It produces [image, label] where image is a 64x64x3(BGR) image, label is an integer in [0, 200).

class TinyImageNet(RNGDataFlow):
    """
    The TinyImageNet classification dataset, with 200 classes and 500 images
    per class. See https://tiny-imagenet.herokuapp.com/.

    It produces [image, label] where image is a 64x64x3(BGR) image, label is an
    integer in [0, 200).
    """
    def __init__(self, dir, name, shuffle=None):
        """
        Args:
            dir (str): a directory
            name (str): one of 'train' or 'val'
            shuffle (bool): shuffle the dataset.
                Defaults to True if name=='train'.
        """
        assert name in ['train', 'val'], name
        dir = Path(os.path.expanduser(dir))
        assert os.path.isdir(dir), dir
        self.full_dir = dir / name
        if shuffle is None:
            shuffle = name == 'train'
        self.shuffle = shuffle

        with open(dir / "wnids.txt") as f:
            wnids = [x.strip() for x in f.readlines()]
        cls_to_id = {name: id for id, name in enumerate(wnids)}
        assert len(cls_to_id) == 200

        self.imglist = []
        if name == 'train':
            for clsid, cls in enumerate(wnids):
                cls_dir = self.full_dir / cls / "images"
                for img in cls_dir.iterdir():
                    self.imglist.append((str(img), clsid))
        else:
            with open(self.full_dir / "val_annotations.txt") as f:
                for line in f:
                    line = line.strip().split()
                    img, cls = line[0], line[1]
                    img = self.full_dir / "images" / img
                    clsid = cls_to_id[cls]
                    self.imglist.append((str(img), clsid))

    def __len__(self):
        return len(self.imglist)

    def __iter__(self):
        idxs = np.arange(len(self.imglist))
        if self.shuffle:
            self.rng.shuffle(idxs)
        for k in idxs:
            fname, label = self.imglist[k]
            im = cv2.imread(fname, cv2.IMREAD_COLOR)
            assert im is not None, fname
            yield [im, label]