From 7800d516596f1a25986b458cddf8b8785bcc7df8 Mon Sep 17 00:00:00 2001 From: junyanz Date: Fri, 6 Oct 2017 10:46:43 -0700 Subject: support nc=1, add new leaerning rate policy and new initialization --- data/aligned_dataset.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) (limited to 'data/aligned_dataset.py') diff --git a/data/aligned_dataset.py b/data/aligned_dataset.py index 0f45c40..bccd6fc 100644 --- a/data/aligned_dataset.py +++ b/data/aligned_dataset.py @@ -40,12 +40,27 @@ class AlignedDataset(BaseDataset): B = AB[:, h_offset:h_offset + self.opt.fineSize, w + w_offset:w + w_offset + self.opt.fineSize] + if self.opt.which_direction == 'BtoA': + input_nc = self.opt.output_nc + output_nc = self.opt.input_nc + else: + input_nc = self.opt.input_nc + output_nc = self.opt.output_nc + if (not self.opt.no_flip) and random.random() < 0.5: idx = [i for i in range(A.size(2) - 1, -1, -1)] idx = torch.LongTensor(idx) A = A.index_select(2, idx) B = B.index_select(2, idx) + 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 {'A': A, 'B': B, 'A_paths': AB_path, 'B_paths': AB_path} -- cgit v1.2.3-70-g09d2