From 16691c00f83332267fcf180fd92e84ce1d1d915d Mon Sep 17 00:00:00 2001 From: Michael Hsu Date: Sun, 24 Feb 2019 10:02:07 +0800 Subject: [PATCH] Let sample.py can generate caption for gray image. The sample.py cannot be used to generate caption for gray image.In the data_loader.py, all of the images have been converted to 'RGB' format. But when generating captions using sample.py for single image, this conversion is missed. So for gray image, it will have such RuntimeError: Given groups=1, weight of size [64, 3, 7, 7], expected input[1, 1, 224, 224] to have 3 channels, but got 1 channels instead. --- tutorials/03-advanced/image_captioning/sample.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tutorials/03-advanced/image_captioning/sample.py b/tutorials/03-advanced/image_captioning/sample.py index 23e07ef..b41cc6f 100644 --- a/tutorials/03-advanced/image_captioning/sample.py +++ b/tutorials/03-advanced/image_captioning/sample.py @@ -14,7 +14,7 @@ from PIL import Image 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 @@ if __name__ == '__main__': 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) \ No newline at end of file + main(args)