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

hanoody: <p>I have a UNET segmentation model with 5 classes and I am having trouble trying to save the image predictions.</p>
<p>Here is the code:</p>
<pre><code class="lang-auto">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 &gt; 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()
</code></pre>
<p>My tensor preds has shape [1, 5, 100, 100] because I have 5 classes and the image size is 100x100.</p>
<p>I get the following error message:</p>
<pre><code class="lang-auto">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 &lt;module&gt;
    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
</code></pre>
<p>Any reason why this might be occurring? I tried the following suggestions but they do not seem to work:</p>
<aside class="onebox stackexchange">
  <header class="source">
      <a href="https://stackoverflow.com/questions/55319949/pil-typeerror-cannot-handle-this-data-type" target="_blank" rel="noopener nofollow ugc">stackoverflow.com</a>
  </header>
  <article class="onebox-body">
      <a href="https://stackoverflow.com/users/10360699/jerome" target="_blank" rel="noopener nofollow ugc">
    <img alt="Jerome" src="https://www.gravatar.com/avatar/a3f1c965ea8b54ea6f666825e9eeaa1d?s=128&amp;d=identicon&amp;r=PG&amp;f=1" class="thumbnail onebox-avatar" width="128" height="128">
  </a>
<h4>
  <a href="https://stackoverflow.com/questions/55319949/pil-typeerror-cannot-handle-this-data-type" target="_blank" rel="noopener nofollow ugc">PIL TypeError: Cannot handle this data type</a>
</h4>

<div class="tags">
  <strong>python, python-imaging-library</strong>
</div>

<div class="date">
  asked by
  
  <a href="https://stackoverflow.com/users/10360699/jerome" target="_blank" rel="noopener nofollow ugc">
    Jerome
  </a>
  on <a href="https://stackoverflow.com/questions/55319949/pil-typeerror-cannot-handle-this-data-type" target="_blank" rel="noopener nofollow ugc">01:26AM - 24 Mar 19 UTC</a>
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<aside class="onebox stackexchange">
  <header class="source">
      <a href="https://stackoverflow.com/questions/63367136/pil-typeerror-cannot-handle-this-data-type-1-1-1-u1" target="_blank" rel="noopener nofollow ugc">stackoverflow.com</a>
  </header>
  <article class="onebox-body">
      <a href="https://stackoverflow.com/users/3306097/kkgarg" target="_blank" rel="noopener nofollow ugc">
    <img alt="kkgarg" src="https://www.gravatar.com/avatar/40c7885a69cb1cfe34056f5e51b58fa9?s=128&amp;d=identicon&amp;r=PG&amp;f=1" class="thumbnail onebox-avatar" width="60" height="60">
  </a>
<h4>
  <a href="https://stackoverflow.com/questions/63367136/pil-typeerror-cannot-handle-this-data-type-1-1-1-u1" target="_blank" rel="noopener nofollow ugc">PIL TypeError: Cannot handle this data type: (1, 1, 1), |u1</a>
</h4>

<div class="tags">
  <strong>arrays, image, numpy, python-imaging-library</strong>
</div>

<div class="date">
  asked by
  
  <a href="https://stackoverflow.com/users/3306097/kkgarg" target="_blank" rel="noopener nofollow ugc">
    kkgarg
  </a>
  on <a href="https://stackoverflow.com/questions/63367136/pil-typeerror-cannot-handle-this-data-type-1-1-1-u1" target="_blank" rel="noopener nofollow ugc">10:19PM - 11 Aug 20 UTC</a>
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>Thank you!</p>

ptrblck: <p>Based on the error message it seems that <code>PIL</code> doesn’t recognize the used “image” format with 5 channels as seen here:</p>
<pre><code class="lang-python">x = torch.empty(1, 1, 5).uniform_().byte()
PIL.Image.fromarray(x.numpy())
&gt; TypeError: Cannot handle this data type: (1, 1, 5), |u1
</code></pre>
<p>so you would have to make sure the model output is using an expected shape (e.g. 3  or 1 channel(s)).</p>

hanoody: <p>Hi, thank you for your response!</p>
<p>I am confused because I thought that the output of a UNET should be <code>[batch_size, # of classes, height, width]</code>.</p>
<p>When I print <code>preds.shape</code> right after <code>preds = (preds &gt; 0.5).float()</code>, I get<br>
= <code>Prediction shape: torch.Size([1, 5, 100, 100])</code>. Am I doing something wrong?</p>

hanoody: <p>Update. I realized I was calling my UNET model incorrectly.</p>
<p>Previously I had <code>model = UNET(num_classes=5, in_channels=3, out_channels=5).to(DEVICE)</code> but it should be <code>model = UNET(num_classes=5, in_channels=3, out_channels=3).to(DEVICE)</code>.</p>

hanoody: <p>Actually sorry the issue has not been resolved.</p>
<p>For a UNET multi-class segmentation model, the output should be <code>[batch_size, # of classes, height, width]</code>.</p>
<p>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.</p>
<p>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.</p>

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

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

hanoody: <p>Thank you for your response as always!</p>
<p>I believe this post has an implementation of your predicted class indices suggestion: <a href="https://discuss.pytorch.org/t/how-to-save-multi-class-segmentation-prediction-as-image/123536/2" class="inline-onebox">How to save multi-class segmentation prediction as image? - #2 by eqy</a>.</p>
<p>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.</p>

Bruce_Kieu: <p>I followed your suggested post, then it works well to me, thanks for raising up the issue</p>
<pre><code class="lang-auto">for idx, (x, y) in enumerate(loader):
        x = x.to(device=device)
        with torch.no_grad():
            preds = torch.sigmoid(model(x))
            out = (preds &gt; 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")
</code></pre>