lsun.py 5.33 KB
Newer Older
1
from .vision import VisionDataset
soumith's avatar
soumith committed
2
3
4
from PIL import Image
import os
import os.path
soumith's avatar
soumith committed
5
import six
soumith's avatar
soumith committed
6
7
import string
import sys
8
9
10
11
12

if sys.version_info < (3, 3):
    from collections import Iterable
else:
    from collections.abc import Iterable
13

soumith's avatar
soumith committed
14
15
16
17
18
if sys.version_info[0] == 2:
    import cPickle as pickle
else:
    import pickle

19
20
from .utils import verify_str_arg, iterable_to_str

21

22
class LSUNClass(VisionDataset):
Jason Park's avatar
Jason Park committed
23
    def __init__(self, root, transform=None, target_transform=None):
soumith's avatar
soumith committed
24
        import lmdb
25
26
        super(LSUNClass, self).__init__(root, transform=transform,
                                        target_transform=target_transform)
Jason Park's avatar
Jason Park committed
27
28

        self.env = lmdb.open(root, max_readers=1, readonly=True, lock=False,
soumith's avatar
soumith committed
29
                             readahead=False, meminit=False)
soumith's avatar
soumith committed
30
31
        with self.env.begin(write=False) as txn:
            self.length = txn.stat()['entries']
32
        cache_file = '_cache_' + ''.join(c for c in root if c in string.ascii_letters)
soumith's avatar
soumith committed
33
        if os.path.isfile(cache_file):
34
            self.keys = pickle.load(open(cache_file, "rb"))
soumith's avatar
soumith committed
35
36
        else:
            with self.env.begin(write=False) as txn:
37
38
                self.keys = [key for key, _ in txn.cursor()]
            pickle.dump(self.keys, open(cache_file, "wb"))
soumith's avatar
soumith committed
39
40
41
42
43
44
45

    def __getitem__(self, index):
        img, target = None, None
        env = self.env
        with env.begin(write=False) as txn:
            imgbuf = txn.get(self.keys[index])

soumith's avatar
soumith committed
46
        buf = six.BytesIO()
soumith's avatar
soumith committed
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
        buf.write(imgbuf)
        buf.seek(0)
        img = Image.open(buf).convert('RGB')

        if self.transform is not None:
            img = self.transform(img)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target

    def __len__(self):
        return self.length

62

63
class LSUN(VisionDataset):
soumith's avatar
soumith committed
64
    """
Gerald Baier's avatar
Gerald Baier committed
65
    `LSUN <https://www.yf.io/p/lsun>`_ dataset.
66
67

    Args:
Jason Park's avatar
Jason Park committed
68
        root (string): Root directory for the database files.
69
        classes (string or list): One of {'train', 'val', 'test'} or a list of
70
            categories to load. e,g. ['bedroom_train', 'church_outdoor_train'].
71
72
73
74
        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.
soumith's avatar
soumith committed
75
    """
76

77
78
79
    def __init__(self, root, classes='train', transform=None, target_transform=None):
        super(LSUN, self).__init__(root, transform=transform,
                                   target_transform=target_transform)
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
        self.classes = self._verify_classes(classes)

        # for each class, create an LSUNClassDataset
        self.dbs = []
        for c in self.classes:
            self.dbs.append(LSUNClass(
                root=root + '/' + c + '_lmdb',
                transform=transform))

        self.indices = []
        count = 0
        for db in self.dbs:
            count += len(db)
            self.indices.append(count)

        self.length = count

    def _verify_classes(self, classes):
soumith's avatar
soumith committed
98
99
100
101
        categories = ['bedroom', 'bridge', 'church_outdoor', 'classroom',
                      'conference_room', 'dining_room', 'kitchen',
                      'living_room', 'restaurant', 'tower']
        dset_opts = ['train', 'val', 'test']
Jason Park's avatar
Jason Park committed
102

103
104
        try:
            verify_str_arg(classes, "classes", dset_opts)
soumith's avatar
soumith committed
105
106
107
108
            if classes == 'test':
                classes = [classes]
            else:
                classes = [c + '_' + classes for c in categories]
109
        except ValueError:
110
111
112
113
114
115
116
117
            if not isinstance(classes, Iterable):
                msg = ("Expected type str or Iterable for argument classes, "
                       "but got type {}.")
                raise ValueError(msg.format(type(classes)))

            classes = list(classes)
            msg_fmtstr = ("Expected type str for elements in argument classes, "
                          "but got type {}.")
soumith's avatar
soumith committed
118
            for c in classes:
119
                verify_str_arg(c, custom_msg=msg_fmtstr.format(type(c)))
soumith's avatar
soumith committed
120
                c_short = c.split('_')
121
122
                category, dset_opt = '_'.join(c_short[:-1]), c_short[-1]

123
                msg_fmtstr = "Unknown value '{}' for {}. Valid values are {{{}}}."
124
125
126
127
128
129
                msg = msg_fmtstr.format(category, "LSUN class",
                                        iterable_to_str(categories))
                verify_str_arg(category, valid_values=categories, custom_msg=msg)

                msg = msg_fmtstr.format(dset_opt, "postfix", iterable_to_str(dset_opts))
                verify_str_arg(dset_opt, valid_values=dset_opts, custom_msg=msg)
soumith's avatar
soumith committed
130

131
        return classes
soumith's avatar
soumith committed
132
133

    def __getitem__(self, index):
134
135
136
137
138
139
140
        """
        Args:
            index (int): Index

        Returns:
            tuple: Tuple (image, target) where target is the index of the target category.
        """
soumith's avatar
soumith committed
141
142
143
144
145
146
        target = 0
        sub = 0
        for ind in self.indices:
            if index < ind:
                break
            target += 1
Zhou Le's avatar
Zhou Le committed
147
            sub = ind
soumith's avatar
soumith committed
148
149
150
151
152
153
154

        db = self.dbs[target]
        index = index - sub

        if self.target_transform is not None:
            target = self.target_transform(target)

soumith's avatar
soumith committed
155
156
        img, _ = db[index]
        return img, target
soumith's avatar
soumith committed
157
158
159
160

    def __len__(self):
        return self.length

161
162
    def extra_repr(self):
        return "Classes: {classes}".format(**self.__dict__)