Commit 9dfd3ba1 authored by Amelie Royer's avatar Amelie Royer

Adding support for images with less than 3 channels

parent 739193ea
......@@ -503,8 +503,12 @@ with tf.Session() as sess:
### Apply the model on one image
else:
image = imread(os.path.abspath(args.input))
w, h, _ = image.shape
image = imresize(image, (WIDTH, HEIGHT))[:, :, 0][:, :, None]
if len(image.shape) == 3:
w, h, _ = image.shape
image = imresize(image, (WIDTH, HEIGHT))[:, :, 0][:, :, None]
else:
w, h = image.shape
image = imresize(image, (WIDTH, HEIGHT))[:, :, None]
base, ext = os.path.basename(args.input).rsplit('.', 1)
out_path = os.path.join(log_dir, "%s_colorized.%s" % (base, ext))
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment