Merge pull request #34 from bruce-willis/tensor2img-grayscale-patch

fix tensor2img for grayscale images

core/util.py

@@ -12,7 +12,7 @@ def tensor2img(tensor, out_type=np.uint8, min_max=(-1, 1)):
     Input: 4D(B,(3/1),H,W), 3D(C,H,W), or 2D(H,W), any range, RGB channel order
     Output: 3D(H,W,C) or 2D(H,W), [0,255], np.uint8 (default)
     '''
-    tensor = tensor.squeeze().clamp_(*min_max)  # clamp
+    tensor = tensor.clamp_(*min_max)  # clamp
     n_dim = tensor.dim()
     if n_dim == 4:
         n_img = len(tensor)
@@ -28,7 +28,7 @@ def tensor2img(tensor, out_type=np.uint8, min_max=(-1, 1)):
     if out_type == np.uint8:
         img_np = ((img_np+1) * 127.5).round()
         # Important. Unlike matlab, numpy.unit8() WILL NOT round by default.
-    return img_np.astype(out_type)
+    return img_np.astype(out_type).squeeze()
 
 def postprocess(images):
 	return [tensor2img(image) for image in images]