pad update

tungth07 · Oct 23, 2019 · 9a6f667 · 9a6f667
1 parent a19dc0f
commit 9a6f667
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/dataset.py b/dataset.py
@@ -274,8 +274,9 @@ def __call__(self, batch):
 
         if self.keep_ratio_with_pad:  # same concept with 'Rosetta' paper
             resized_max_w = self.imgW
-            transform = NormalizePAD((1, self.imgH, resized_max_w))
-
+            input_channel = 3 if images[0].mode == 'RGB' else 1
+            transform = NormalizePAD((input_channel, self.imgH, resized_max_w))
+
             resized_images = []
             for image in images:
                 w, h = image.size