File size: 1,862 Bytes
0f52c9d | 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 | import os.path
from data.base_dataset import BaseDataset, get_transform, get_params
from data.image_folder import make_dataset
from PIL import Image
import random
class UnalignedDayNightDataset(BaseDataset):
@staticmethod
def modify_commandline_options(parser, is_train):
parser.add_argument('--dataroot', required=True,
help='path to images (should have subfolders train/val containing day/night)')
parser.set_defaults(preprocess_mode='resize_and_crop', load_size=286, crop_size=256)
if not is_train:
parser.set_defaults(no_flip=True)
return parser
def __init__(self, opt):
BaseDataset.__init__(self)
self.opt = opt
root = opt.dataroot
phase = opt.phase
self.dir_A = os.path.join(root, phase, 'day')
self.dir_B = os.path.join(root, phase, 'night')
self.A_paths = sorted(make_dataset(self.dir_A, recursive=True))
self.B_paths = sorted(make_dataset(self.dir_B, recursive=True))
self.A_size = len(self.A_paths)
self.B_size = len(self.B_paths)
if self.A_size == 0 or self.B_size == 0:
raise (RuntimeError(f"Found 0 images in one of the data directories: {self.dir_A} or {self.dir_B}"))
def __getitem__(self, index):
A_path = self.A_paths[index % self.A_size]
index_B = random.randint(0, self.B_size - 1)
B_path = self.B_paths[index_B]
A_img = Image.open(A_path).convert('RGB')
B_img = Image.open(B_path).convert('RGB')
params = get_params(self.opt, A_img.size)
transform = get_transform(self.opt, params)
A = transform(A_img)
B = transform(B_img)
return {'day': A, 'night': B, 'cpath': A_path, 'spath_night': B_path}
def __len__(self):
return max(self.A_size, self.B_size) |