support nc=1, add new leaerning rate policy and new initialization

1 files changed, 15 insertions, 0 deletions

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}