Error when saving UNET image predictions to folder: TypeError: Cannot handle this data type: (1, 1, 5), |u1

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"
    for idx, (x, y) in enumerate(loader):
        x =
        with torch.no_grad():
            preds = torch.sigmoid(model(x))
            preds = (preds > 0.5).float()
            preds, f"{folder}/pred_{idx}.png"
        torchvision.utils.save_image(y.unsqueeze(1), f"{folder}{idx}.png")


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/", 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 "", line 125, in <module>
  File "", line 119, in main
  File "/home/REDACTED/REDACTED/unet-segmentation/", line 116, in save_predictions_as_imgs
  File "/REDACTED/REDACTED/REDACTED/pkg/anaconda/anaconda3-2021a/lib/python3.8/site-packages/torch/autograd/", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/REDACTED/REDACTED/REDACTED/pkg/anaconda/anaconda3-2021a/lib/python3.8/site-packages/torchvision/", line 134, in save_image
    im = Image.fromarray(ndarr)
  File "/REDACTED/REDACTED/REDACTED/pkg/anaconda/anaconda3-2021a/lib/python3.8/site-packages/PIL/", 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:

Thank you!

Based on the error message it seems that PIL doesn’t recognize the used “image” format with 5 channels as seen here:

x = torch.empty(1, 1, 5).uniform_().byte()
> TypeError: Cannot handle this data type: (1, 1, 5), |u1

so you would have to make sure the model output is using an expected shape (e.g. 3 or 1 channel(s)).

Hi, thank you for your response!

I am confused because I thought that the output of a UNET should be [batch_size, # of classes, height, width].

When I print preds.shape right after preds = (preds > 0.5).float(), I get
= Prediction shape: torch.Size([1, 5, 100, 100]). Am I doing something wrong?

Update. I realized I was calling my UNET model incorrectly.

Previously I had model = UNET(num_classes=5, in_channels=3, out_channels=5).to(DEVICE) but it should be model = UNET(num_classes=5, in_channels=3, out_channels=3).to(DEVICE).

Actually sorry the issue has not been resolved.

For a UNET multi-class segmentation model, the output should be [batch_size, # of classes, height, width].

So with 5 classes, a batch size of 1 and an image size of 100x100, the output should be [1, 5, 100, 100] which is correct.

The issue is trying to save this prediction directly as an image because PIL only recognizes 1 or 3 channels which makes sense. How can I save my model output as an image if the shape is technically incorrect for an image? Thank you.

You won’t be able to save it directly as an image, if it’s not a valid format.
In case you want to store the output directly, you could use On the other hand, you might want to store the predictions for each class as a color-coded image.
In that case, you could get the predicted class indices via preds = torch.argmax(output, dim=1) (assuming output contains the logits/probabilities) and use a color mapping to create an RGB image (e.g. class0 -> RED, class1 -> BLUE, etc.).

1 Like

I am not sure
Please can you confirm the number of channels in UNET output. I have a doubt regarding that it would not be 5.
Is not it there is argmax which give a semantic map.
How you are defining the target here (mask are generally of one channel, so how you are reading target)

Thank you for your response as always!

I believe this post has an implementation of your predicted class indices suggestion: How to save multi-class segmentation prediction as image? - #2 by eqy.

I believe the idea is correct but I am having trouble getting it to run. If you could take a look at that then that would be really amazing! Thank you very much.

I followed your suggested post, then it works well to me, thanks for raising up the issue

for idx, (x, y) in enumerate(loader):
        x =
        with torch.no_grad():
            preds = torch.sigmoid(model(x))
            out = (preds > 0.5).float()        
            class_to_color = [torch.tensor([0.0, 0.0, 0.0]), 
                torch.tensor([14, 1, 133]),  torch.tensor([33, 255, 1]), 
                torch.tensor([243, 5, 247]), torch.tensor([(255, 0, 0)])] #colors' num is equal to out_channels
            output = torch.zeros(1, 3, out.size(-2), out.size(-1), dtype=torch.float)
            for class_idx, color in enumerate(class_to_color):
                mask = out[:,class_idx,:,:] == torch.max(out, dim=1)[0]
                mask = mask.unsqueeze(1) # should have shape 1, 1, 180, 100
                curr_color = color.reshape(1, 3, 1, 1)
                segment = mask*curr_color # should have shape 1, 3, 180, 100
                output += segment
        print('saved pic shape {}, origin shape {}'.format(output.shape, y.shape))
        torchvision.utils.save_image(output, f"{folder}/pred_{idx}.jpg")