I have a UNET segmentation model with 5 classes and I am having trouble trying to save the image predictions.
Here is the code:
def save_predictions_as_imgs(
    loader, model, folder="saved_images/", device="cuda"
):
    model.eval()
    for idx, (x, y) in enumerate(loader):
        x = x.to(device=device)
        with torch.no_grad():
            preds = torch.sigmoid(model(x))
            preds = (preds > 0.5).float()
        torchvision.utils.save_image(
            preds, f"{folder}/pred_{idx}.png"
        )
        torchvision.utils.save_image(y.unsqueeze(1), f"{folder}{idx}.png")

    model.train()

My tensor preds has shape [1, 5, 100, 100] because I have 5 classes and the image size is 100x100.
I get the following error message:
Traceback (most recent call last):
  File "/REDACTED/REDACTED/REDACTED/pkg/anaconda/anaconda3-2021a/lib/python3.8/site-packages/PIL/Image.py", line 2749, in fromarray
    mode, rawmode = _fromarray_typemap[typekey]
KeyError: ((1, 1, 5), '|u1')

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "train.py", line 125, in <module>
    main()
  File "train.py", line 119, in main
    save_predictions_as_imgs(
  File "/home/REDACTED/REDACTED/unet-segmentation/utils.py", line 116, in save_predictions_as_imgs
    torchvision.utils.save_image(
  File "/REDACTED/REDACTED/REDACTED/pkg/anaconda/anaconda3-2021a/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/REDACTED/REDACTED/REDACTED/pkg/anaconda/anaconda3-2021a/lib/python3.8/site-packages/torchvision/utils.py", line 134, in save_image
    im = Image.fromarray(ndarr)
  File "/REDACTED/REDACTED/REDACTED/pkg/anaconda/anaconda3-2021a/lib/python3.8/site-packages/PIL/Image.py", line 2751, in fromarray
    raise TypeError("Cannot handle this data type: %s, %s" % typekey) from e
TypeError: Cannot handle this data type: (1, 1, 5), |u1

Any reason why this might be occurring? I tried the following suggestions but they do not seem to work:

  
      stackoverflow.com
  
  
      
    
  

  PIL TypeError: Cannot handle this data type 1



  python, python-imaging-library



  asked by
  
  
    Jerome
  
  on 01:26AM - 24 Mar 19 UTC


  
  
    
    
  
  



  
      stackoverflow.com
  
  
      
    
  

  PIL TypeError: Cannot handle this data type: (1, 1, 1), |u1 1



  arrays, image, numpy, python-imaging-library



  asked by
  
  
    kkgarg
  
  on 10:19PM - 11 Aug 20 UTC


  
  
    
    
  
  


Thank you!