[Fix] Fix astype error in function tensor2img (open-mmlab#429)

* fix astype error in function tensor2img

* Test out dtype for tensor2img

Signed-off-by: lizz <[email protected]>

Co-authored-by: yangyuantao <[email protected]>
Co-authored-by: lizz <[email protected]>

Author: 3 people

mmedit/core/misc.py

@@ -67,7 +67,7 @@ def tensor2img(tensor, out_type=np.uint8, min_max=(0, 1)):
         if out_type == np.uint8:
             # Unlike MATLAB, numpy.unit8() WILL NOT round by default.
             img_np = (img_np * 255.0).round()
-        img_np.astype(out_type)
+        img_np = img_np.astype(out_type)
         result.append(img_np)
     result = result[0] if len(result) == 1 else result
     return result

tests/test_utils/test_tensor2img.py

@@ -28,28 +28,33 @@ def test_tensor2img():
 
     # 4d
     rlt = tensor2img(tensor_4d_1, out_type=np.uint8, min_max=(0, 1))
+    assert rlt.dtype == np.uint8
     tensor_4d_1_np = make_grid(tensor_4d_1, nrow=1, normalize=False).numpy()
     tensor_4d_1_np = np.transpose(tensor_4d_1_np[[2, 1, 0], :, :], (1, 2, 0))
     np.testing.assert_almost_equal(rlt, (tensor_4d_1_np * 255).round())
 
     rlt = tensor2img(tensor_4d_2, out_type=np.uint8, min_max=(0, 1))
+    assert rlt.dtype == np.uint8
     tensor_4d_2_np = tensor_4d_2.squeeze().numpy()
     tensor_4d_2_np = np.transpose(tensor_4d_2_np[[2, 1, 0], :, :], (1, 2, 0))
     np.testing.assert_almost_equal(rlt, (tensor_4d_2_np * 255).round())
 
     rlt = tensor2img(tensor_4d_3, out_type=np.uint8, min_max=(0, 1))
+    assert rlt.dtype == np.uint8
     tensor_4d_3_np = make_grid(tensor_4d_3, nrow=1, normalize=False).numpy()
     tensor_4d_3_np = np.transpose(tensor_4d_3_np[[2, 1, 0], :, :], (1, 2, 0))
     np.testing.assert_almost_equal(rlt, (tensor_4d_3_np * 255).round())
 
     rlt = tensor2img(tensor_4d_4, out_type=np.uint8, min_max=(0, 1))
+    assert rlt.dtype == np.uint8
     tensor_4d_4_np = tensor_4d_4.squeeze().numpy()
     np.testing.assert_almost_equal(rlt, (tensor_4d_4_np * 255).round())
 
     # 3d
     rlt = tensor2img([tensor_3d_1, tensor_3d_2],
                      out_type=np.uint8,
                      min_max=(0, 1))
+    assert rlt[0].dtype == np.uint8
     tensor_3d_1_np = tensor_3d_1.numpy()
     tensor_3d_1_np = np.transpose(tensor_3d_1_np[[2, 1, 0], :, :], (1, 2, 0))
     tensor_3d_2_np = tensor_3d_2.numpy()
@@ -58,16 +63,20 @@ def test_tensor2img():
     np.testing.assert_almost_equal(rlt[1], (tensor_3d_2_np * 255).round())
 
     rlt = tensor2img(tensor_3d_3, out_type=np.uint8, min_max=(0, 1))
+    assert rlt.dtype == np.uint8
     tensor_3d_3_np = tensor_3d_3.squeeze().numpy()
     np.testing.assert_almost_equal(rlt, (tensor_3d_3_np * 255).round())
 
     # 2d
     rlt = tensor2img(tensor_2d, out_type=np.uint8, min_max=(0, 1))
+    assert rlt.dtype == np.uint8
     tensor_2d_np = tensor_2d.numpy()
     np.testing.assert_almost_equal(rlt, (tensor_2d_np * 255).round())
     rlt = tensor2img(tensor_2d, out_type=np.float32, min_max=(0, 1))
+    assert rlt.dtype == np.float32
     np.testing.assert_almost_equal(rlt, tensor_2d_np)
 
     rlt = tensor2img(tensor_2d, out_type=np.float32, min_max=(0.1, 0.5))
+    assert rlt.dtype == np.float32
     tensor_2d_np = (np.clip(tensor_2d_np, 0.1, 0.5) - 0.1) / 0.4
     np.testing.assert_almost_equal(rlt, tensor_2d_np)