Merge pull request yunjey#159 from Michael-Hsu/patch-1

yunjey · web-flow · commit 32cf71657579 · 2020-01-28T00:04:54.000+09:00
Let sample.py can generate caption for gray image.
diff --git a/tutorials/03-advanced/image_captioning/sample.py b/tutorials/03-advanced/image_captioning/sample.py
@@ -14,7 +14,7 @@
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 def load_image(image_path, transform=None):
-    image = Image.open(image_path)
+    image = Image.open(image_path).convert('RGB')
     image = image.resize([224, 224], Image.LANCZOS)
     
     if transform is not None:
@@ -78,4 +78,4 @@ def main(args):
     parser.add_argument('--hidden_size', type=int , default=512, help='dimension of lstm hidden states')
     parser.add_argument('--num_layers', type=int , default=1, help='number of layers in lstm')
     args = parser.parse_args()
-    main(args)
+    main(args)
