import os.path
import torchvision.transforms as transforms
from data.base_dataset import BaseDataset, get_transform
from data.image_folder import make_dataset
from PIL import Image
import PIL
import random
class UnalignedTripletDataset(BaseDataset):
def initialize(self, opt):
self.opt = opt
self.root = opt.dataroot
self.dir_A = os.path.join(opt.dataroot, opt.phase + 'A')
self.dir_B = os.path.join(opt.dataroot, opt.phase + 'B')
self.A_paths = make_dataset(self.dir_A)
self.B_paths = make_dataset(self.dir_B)
self.A_paths = sorted(self.A_paths)
self.B_paths = sorted(self.B_paths)
self.A_size = len(self.A_paths)
self.B_size = len(self.B_paths)
# self.transform = get_transform(opt)
transform_list = [transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5),
(0.5, 0.5, 0.5))]
self.transform = transforms.Compose(transform_list)
def __getitem__(self, index):
A_path = self.A_paths[index % self.A_size]
index_A = index % self.A_size
if self.opt.serial_batches:
index_B = index % self.B_size
else:
index_B = random.randint(0, self.B_size - 1)
B_path = self.B_paths[index_B]
# print('(A, B) = (%d, %d)' % (index_A, index_B))
# read the triplet from A and B --
A_img = Image.open(A_path).convert('RGB')
B_img = Image.open(B_path).convert('RGB')
# A = self.transform(A_img)
# B = self.transform(B_img)
# get the triplet from A
A_img = A_img.resize((self.opt.loadSize * 3, self.opt.loadSize), Image.BICUBIC)
A_img = self.transform(A_img)
w_total = A_img.size(2)
w = int(w_total / 3)
h = A_img.size(1)
w_offset = random.randint(0, max(0, w - self.opt.fineSize - 1))
h_offset = random.randint(0, max(0, h - self.opt.fineSize - 1))
A0 = A_img[:, h_offset:h_offset + self.opt.fineSize,
w_offset:w_offset + self.opt.fineSize]
A1 = A_img[:, h_offset:h_offset + self.opt.fineSize,
w + w_offset:w + w_offset + self.opt.fineSize]
A2 = A_img[:, h_offset:h_offset + self.opt.fineSize,
2 * w + w_offset:2 * w + w_offset + self.opt.fineSize]
## -- get the triplet from B
B_img = B_img.resize((self.opt.loadSize * 3, self.opt.loadSize), Image.BICUBIC)
B_img = self.transform(B_img)
w_total = B_img.size(2)
w = int(w_total / 3)
h = B_img.size(1)
w_offset = random.randint(0, max(0, w - self.opt.fineSize - 1))
h_offset = random.randint(0, max(0, h - self.opt.fineSize - 1))
B0 = B_img[:, h_offset:h_offset + self.opt.fineSize,
w_offset:w_offset + self.opt.fineSize]
B1 = B_img[:, h_offset:h_offset + self.opt.fineSize,
w + w_offset:w + w_offset + self.opt.fineSize]
B2 = B_img[:, h_offset:h_offset + self.opt.fineSize,
2 * w + w_offset:2 * w + w_offset + self.opt.fineSize]
#######
input_nc = self.opt.input_nc
output_nc = self.opt.output_nc
# if input_nc == 1: # RGB to gray
# tmp = A[0, ...] * 0.299 + A[1, ...] * 0.587 + A[2, ...] * 0.114
# A = tmp.unsqueeze(0)
# if output_nc == 1: # RGB to gray
# tmp = B[0, ...] * 0.299 + B[1, ...] * 0.587 + B[2, ...] * 0.114
# B = tmp.unsqueeze(0)
return {'A0': A0, 'A1': A1, 'A2': A2, 'B0': B0, 'B1': B1, 'B2': B2,
'A_paths': A_path, 'B_paths': B_path}
def __len__(self):
return max(self.A_size, self.B_size)
def name(self):
return 'UnalignedTripletDataset'
