169 lines
5.6 KiB
Python
169 lines
5.6 KiB
Python
import io
|
|
import os.path
|
|
import pickle
|
|
import string
|
|
from collections.abc import Iterable
|
|
from pathlib import Path
|
|
from typing import Any, Callable, cast, List, Optional, Tuple, Union
|
|
|
|
from PIL import Image
|
|
|
|
from .utils import iterable_to_str, verify_str_arg
|
|
from .vision import VisionDataset
|
|
|
|
|
|
class LSUNClass(VisionDataset):
|
|
def __init__(
|
|
self, root: str, transform: Optional[Callable] = None, target_transform: Optional[Callable] = None
|
|
) -> None:
|
|
import lmdb
|
|
|
|
super().__init__(root, transform=transform, target_transform=target_transform)
|
|
|
|
self.env = lmdb.open(root, max_readers=1, readonly=True, lock=False, readahead=False, meminit=False)
|
|
with self.env.begin(write=False) as txn:
|
|
self.length = txn.stat()["entries"]
|
|
cache_file = "_cache_" + "".join(c for c in root if c in string.ascii_letters)
|
|
if os.path.isfile(cache_file):
|
|
self.keys = pickle.load(open(cache_file, "rb"))
|
|
else:
|
|
with self.env.begin(write=False) as txn:
|
|
self.keys = [key for key in txn.cursor().iternext(keys=True, values=False)]
|
|
pickle.dump(self.keys, open(cache_file, "wb"))
|
|
|
|
def __getitem__(self, index: int) -> Tuple[Any, Any]:
|
|
img, target = None, None
|
|
env = self.env
|
|
with env.begin(write=False) as txn:
|
|
imgbuf = txn.get(self.keys[index])
|
|
|
|
buf = io.BytesIO()
|
|
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) -> int:
|
|
return self.length
|
|
|
|
|
|
class LSUN(VisionDataset):
|
|
"""`LSUN <https://www.yf.io/p/lsun>`_ dataset.
|
|
|
|
You will need to install the ``lmdb`` package to use this dataset: run
|
|
``pip install lmdb``
|
|
|
|
Args:
|
|
root (str or ``pathlib.Path``): Root directory for the database files.
|
|
classes (string or list): One of {'train', 'val', 'test'} or a list of
|
|
categories to load. e,g. ['bedroom_train', 'church_outdoor_train'].
|
|
transform (callable, optional): A function/transform that takes in a 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.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
root: Union[str, Path],
|
|
classes: Union[str, List[str]] = "train",
|
|
transform: Optional[Callable] = None,
|
|
target_transform: Optional[Callable] = None,
|
|
) -> None:
|
|
super().__init__(root, transform=transform, target_transform=target_transform)
|
|
self.classes = self._verify_classes(classes)
|
|
|
|
# for each class, create an LSUNClassDataset
|
|
self.dbs = []
|
|
for c in self.classes:
|
|
self.dbs.append(LSUNClass(root=os.path.join(root, f"{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: Union[str, List[str]]) -> List[str]:
|
|
categories = [
|
|
"bedroom",
|
|
"bridge",
|
|
"church_outdoor",
|
|
"classroom",
|
|
"conference_room",
|
|
"dining_room",
|
|
"kitchen",
|
|
"living_room",
|
|
"restaurant",
|
|
"tower",
|
|
]
|
|
dset_opts = ["train", "val", "test"]
|
|
|
|
try:
|
|
classes = cast(str, classes)
|
|
verify_str_arg(classes, "classes", dset_opts)
|
|
if classes == "test":
|
|
classes = [classes]
|
|
else:
|
|
classes = [c + "_" + classes for c in categories]
|
|
except ValueError:
|
|
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_type = "Expected type str for elements in argument classes, but got type {}."
|
|
for c in classes:
|
|
verify_str_arg(c, custom_msg=msg_fmtstr_type.format(type(c)))
|
|
c_short = c.split("_")
|
|
category, dset_opt = "_".join(c_short[:-1]), c_short[-1]
|
|
|
|
msg_fmtstr = "Unknown value '{}' for {}. Valid values are {{{}}}."
|
|
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)
|
|
|
|
return classes
|
|
|
|
def __getitem__(self, index: int) -> Tuple[Any, Any]:
|
|
"""
|
|
Args:
|
|
index (int): Index
|
|
|
|
Returns:
|
|
tuple: Tuple (image, target) where target is the index of the target category.
|
|
"""
|
|
target = 0
|
|
sub = 0
|
|
for ind in self.indices:
|
|
if index < ind:
|
|
break
|
|
target += 1
|
|
sub = ind
|
|
|
|
db = self.dbs[target]
|
|
index = index - sub
|
|
|
|
if self.target_transform is not None:
|
|
target = self.target_transform(target)
|
|
|
|
img, _ = db[index]
|
|
return img, target
|
|
|
|
def __len__(self) -> int:
|
|
return self.length
|
|
|
|
def extra_repr(self) -> str:
|
|
return "Classes: {classes}".format(**self.__dict__)
|