-
Notifications
You must be signed in to change notification settings - Fork 3
/
dataset_ffhq.py
58 lines (48 loc) · 2.1 KB
/
dataset_ffhq.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
from io import BytesIO
import lmdb
from PIL import Image
from torch.utils.data import Dataset
import numpy as np
from torchvision import utils, transforms
import torch
import os
class MultiResolutionDataset(Dataset):
def __init__(self, path, transform, resolution, exp_setting):
self.env = lmdb.open(
path,
max_readers=32,
readonly=True,
lock=False,
readahead=False,
meminit=False,
)
if not self.env:
raise IOError('Cannot open lmdb dataset', path)
self.resolution = resolution
self.transform = transform
self.distance_dir = './distances'
self.exp_setting = exp_setting
id_path = os.path.join(self.distance_dir, f'index_{self.exp_setting}.txt')
self.id_list = np.loadtxt(id_path)
self.length = len(self.id_list)
self.attr_val = {}
attr_dict = {attr_name: f'{attr_name}_LARGE.txt' for attr_name in ['age', 'pose', 'glasses', 'hairlong', 'haircolor', 'smile', 'gender']}
for attr in attr_dict.keys():
self.attr_val[attr] = np.loadtxt(os.path.join(self.distance_dir, attr_dict[attr]))
def __len__(self):
return self.length
def __getitem__(self, index):
raw_index = index
index = int(self.id_list[index])
with self.env.begin(write=False) as txn:
key = f'{self.resolution}-{str(index).zfill(5)}'.encode('utf-8')
img_bytes = txn.get(key)
buffer = BytesIO(img_bytes)
img = Image.open(buffer)
img = self.transform(img)
if self.exp_setting == 'glassessmileage':
return img, index, self.attr_val['glasses'][index], self.attr_val['age'][index], self.attr_val['smile'][index]
elif self.exp_setting == 'haircolorhairlonggender':
return img, index, self.attr_val['haircolor'][index], self.attr_val['hairlong'][index], self.attr_val['gender'][index]
elif self.exp_setting == 'glassesagepose':
return img, index, self.attr_val['glasses'][index], self.attr_val['age'][index], self.attr_val['pose'][index]