diff options
Diffstat (limited to 'util/util.py')
| -rw-r--r-- | util/util.py | 2 |
1 files changed, 2 insertions, 0 deletions
diff --git a/util/util.py b/util/util.py index 781239f..4de0a74 100644 --- a/util/util.py +++ b/util/util.py @@ -11,6 +11,8 @@ import collections # |imtype|: the desired type of the converted numpy array def tensor2im(image_tensor, imtype=np.uint8): image_numpy = image_tensor[0].cpu().float().numpy() + if image_numpy.shape[0] == 1: + image_numpy = np.tile(image_numpy, (3, 1, 1)) image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + 1) / 2.0 * 255.0 return image_numpy.astype(imtype) |