Spaces:
Runtime error
Runtime error
File size: 1,186 Bytes
1b2a9b1 |
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 |
import random
import os.path
from swapae.data.base_dataset import BaseDataset
from swapae.data.lmdb_dataset import LMDBDataset
import swapae.util
class UnalignedLMDBDataset(BaseDataset):
def __init__(self, opt):
super().__init__(opt)
self.dir_A = os.path.join(opt.dataroot, opt.phase + 'A') # create a path '/path/to/data/trainA'
self.dir_B = os.path.join(opt.dataroot, opt.phase + 'B') # create a path '/path/to/data/trainB'
self.dataset_A = LMDBDataset(util.copyconf(opt, dataroot=self.dir_A))
self.dataset_B = LMDBDataset(util.copyconf(opt, dataroot=self.dir_B))
self.B_indices = list(range(len(self.dataset_B)))
def __len__(self):
return max(len(self.dataset_A), len(self.dataset_B))
def __getitem__(self, index):
if index == 0 and self.opt.isTrain:
random.shuffle(self.B_indices)
result = self.dataset_A.__getitem__(index % len(self.dataset_A))
B_index = self.B_indices[index % len(self.dataset_B)]
B_result = self.dataset_B.__getitem__(B_index)
result["real_B"] = B_result["real_A"]
result["path_B"] = B_result["path_A"]
return result
|