Vision: Required range is confusing in torchvision.utils.save_image

[chinglamchoi]

https://discuss.pytorch.org/t/float-vs-int-in-torchvision-utils-save-image/68596

[pmeier]

From the referenced community post by @chinglamchoi:

I predict a binary segmentation mask using an array with values 1 and 0, multiply this by 255, and try to save the resulting array as an image. However, I’m getting the following error:

Traceback (most recent call last):
File “test.py”, line 71, in
torchvision.utils.save_image(predicted, path_ + idx[0])
File “C:\Users\CCL\Anaconda3\lib\site-packages\torchvision\utils.py”, line 107, in save_image
ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to(‘cpu’, torch.uint8).numpy() RuntimeError: result type Float can’t be cast to the desired output type Int

Here is my code:

mask_pred = net(img)
predicted = torch.sigmoid(mask_pred) > 0.5
predicted = torch.tensor(predicted, dtype=torch.int32)*255
torchvision.utils.save_image(predicted, path_ + idx[0])

I’m using torch version 1.4.0 and torchvision version 0.5.0. The above code works on torch v1.1.0 and torchvision v0.3.0.

[pmeier]

The native "image data type" for torchvision is a torch.FloatTensor within the range [0.0, 1.0]. In your example this means you can eliminate the casting you performed and simply cast it to torch.float.

mask_pred = net(img)
predicted = torch.sigmoid(mask_pred) > 0.5
# predicted = torch.tensor(predicted, dtype=torch.int32)*255
predicted = predicted.float()
torchvision.utils.save_image(predicted, path_ + idx[0])

Here is a minimal working example:

import torch
import torchvision

mask_pred = torch.randn(3, 1, 256, 256)
predicted = torch.sigmoid(mask_pred) > 0.5
torchvision.utils.save_image(predicted.float(), "test_segmentation.png")

[fmassa]

Thanks for the help @pmeier !