How to save multi-class segmentation prediction as image?

hanoody: <p>Hi,</p>
<p>My multi-class UNET model output is of the following shape: [1, 6, 100, 100] which is expected because the batch size is 1, I have 6 classes, and the image size is 100x100.</p>
<p>How can I save a prediction as an image which contains all 6 classes using <code>torchvision.utils.</code> <code>save_image</code>? The tensor in save_image expects either 1 or 3 channels but my model technically outputs 6 channels for the 6 classes.</p>
<p>Thank you.</p>

eqy: <p>There might be a utility function somewhere that does this, but you can write your own function like:</p>
<pre><code class="lang-auto">    class_to_color = [torch.tensor([1.0, 0.0, 0.0]), ...]
    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)
        mask = mask.unsqueeze(1) # should have shape 1, 1, 100, 100
        curr_color = color.reshape(1, 3, 1, 1)
        segment = mask*color # should have shape 1, 3, 100, 100
        output += segment
</code></pre>

hanoody: <p>Thank you for your response. I get the following error when executing that code:</p>
<pre><code class="lang-auto">mask = mask.unsqueeze(1) # should have shape 1, 1, 100, 100
AttributeError: 'bool' object has no attribute 'unsqueeze'
</code></pre>
<p>since <code>mask</code> is a bool from <code>mask = out[:,class_idx,:,:] == torch.max(out, dim=1)</code>.</p>
<p>What are we trying to accomplish from this line? How should we rewrite it to make mask a tensor and not a bool? Thank you!</p>

eqy: <p>Ah, that’s because max actually gives back the values and the indices; changing it to <code>mask = out[:,class_idx,:,:] == torch.max(out, dim=1)[0]</code> should fix that issue.</p>
<pre><code class="lang-auto">    class_to_color = [torch.tensor([1.0, 0.0, 0.0]), ...]
    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, 100, 100
        curr_color = color.reshape(1, 3, 1, 1)
        segment = mask*color # should have shape 1, 3, 100, 100
        output += segment
</code></pre>

hanoody: <p>Thank you. That solved that issue but now I am running into a new one.</p>
<pre><code class="lang-auto">segment = mask*color # should have shape 1, 3, 100, 100
RuntimeError: The size of tensor a (100) must match the size of tensor b (3) at non-singleton dimension 3
</code></pre>
<p>Also, the shape of mask after <code>mask = mask.unsqueeze(1)</code> is <code>[1, 3, 100, 100]</code> instead of <code>[1, 1, 100, 100]</code> but I am not sure why.</p>

eqy: <p>I cannot reproduce that issue.<br>
Can you check the shape of your output is expected?</p>
<pre><code class="lang-auto">&gt;&gt;&gt; out = torch.randn(1, 6, 100, 100)
&gt;&gt;&gt; mask = out[:,0,:,:] == torch.max(out, dim=1)[0]
&gt;&gt;&gt; mask.shape
torch.Size([1, 100, 100])
&gt;&gt;&gt; mask.unsqueeze(1).shape
torch.Size([1, 1, 100, 100])
</code></pre>

hanoody: <p>Yes the shape is correct.</p>
<p>Here is the code:</p>
<pre><code class="lang-auto">            pred = torch.sigmoid(model(x))
            out = (pred &gt; 0.5).float()
            print(f"out shape: {out.shape}\n")
            class_to_color = [torch.tensor([0.0, 0.0, 0.0]), torch.tensor([10, 133, 1]), torch.tensor([14, 1, 133]),  torch.tensor([33, 255, 1]), torch.tensor([243, 5, 247]), torch.tensor([(255, 0, 0)])]
            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]
                print(f"{mask}\n")
                mask = mask.unsqueeze(1) # should have shape 1, 1, 100, 100
                print(f"mask shape: {mask.shape}\n")
                curr_color = color.reshape(1, 3, 1, 1)
                print(f"color shape: {color.shape}\n")
                segment = mask*color # should have shape 1, 3, 100, 100
                output += segment
            torchvision.utils.save_image(output, f"{folder}/pred_{idx}.png")
</code></pre>
<p>Here is the corresponding output:</p>
<pre><code class="lang-auto">out shape: torch.Size([1, 6, 100, 100])

tensor([[[True, True, True,  ..., True, True, True],
         [True, True, True,  ..., True, True, True],
         [True, True, True,  ..., True, True, True],
         ...,
         [True, True, True,  ..., True, True, True],
         [True, True, True,  ..., True, True, True],
         [True, True, True,  ..., True, True, True]]], device='cuda:0')

mask shape: torch.Size([1, 1, 100, 100])

color shape: torch.Size([3])
</code></pre>
<p>And the error:</p>
<pre><code class="lang-auto">segment = mask*color # should have shape 1, 3, 100, 100
RuntimeError: The size of tensor a (100) must match the size of tensor b (3) at non-singleton dimension 3
</code></pre>

eqy: <p>Ah, <code>curr_color</code> should be used instead:</p>
<pre><code class="lang-auto">    class_to_color = [torch.tensor([1.0, 0.0, 0.0]), ...]
    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, 100, 100
        curr_color = color.reshape(1, 3, 1, 1)
        segment = mask*curr_color # should have shape 1, 3, 100, 100
        output += segment
</code></pre>
<p>However, note that you should make sure the color formatting is consistent (e.g., either floating point values between 0.0 and 1.0 or integers between 0 and 255).</p>

hanoody: <p>Thank you! That solved that issue but now I am getting the following issue <img src="https://discuss.pytorch.org/images/emoji/apple/face_with_monocle.png?v=9" title=":face_with_monocle:" class="emoji" alt=":face_with_monocle:"></p>
<pre><code class="lang-auto">segment = mask*curr_color # should have shape 1, 3, 100, 100
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
</code></pre>
<p>Did one of mask or curr_color get implicitly allocated on the CPU? I thought everything gets allocated on the GPU by default.</p>

eqy: <p>Nope, things are allocated on the CPU by default. You can simply add <code>device='cuda'</code> to the <code>torch.tensor(...)</code> calls to fix this.</p>

hanoody: <p>Works like a charm. I also had to add <code>output</code> to the GPU.</p>
<p>Thank you very much I really appreciate your prompt and helpful responses! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

hanoody: <p>Hi <a class="mention" href="/u/eqy">@eqy</a>, related to this, do you know how I can use this approach of iterating class by class so that I can compute the accuracy for each class?</p>
<p>This is what I currently have. <code>check_accuracy</code> gets executed for each epoch during training.</p>
<pre><code class="lang-auto">def check_accuracy(loader, model, device="cuda"):
    model.eval()

    with torch.no_grad():
        for x, y, _ in loader:
            x = x.to(device)
            y = y.to(device).unsqueeze(1)
            preds = model(x)
            num_classes = preds.shape[1]
            for class_idx in range(num_classes):
                curr_class = (preds[:,class_idx,:,:] == torch.max(preds, dim=1)[0]).float()
                curr_class = curr_class.unsqueeze(1)
                num_correct = (curr_class == y).float()
                class_accuracy = num_correct.sum()/num_correct.numel()
                class_accuracy = class_accuracy*100
                print(f"{class_accuracy}% accuracy for class {class_idx}\n")
    model.train()
</code></pre>
<p>Is this the correct way to do it?</p>
<p>Thank you very much!</p>

eqy: <p>I don’t see anything obviously wrong with this immediately, but I also don’t recall the exact data layout and format used here. The best practice would be to write some predictable test cases if you aren’t sure about the correctness (e.g., try generating synthetic model outputs and labels like an entire image with just one class, half and half). You could also check that random outputs get roughly the accuracy of random guessing, etc…</p>

hanoody: <p>Hi,</p>
<p>I just have a follow up question. Since it’s been almost 2 months, here is a refresher.</p>
<p>We wanted to visualize the output of a UNET multi-class segmentation model. The model outputs a tensor of the following form: [batch size, number of classes, h, w] where in each pixel of the tensor, it is an index in [0, number of classes). Obviously if the number of classes isn’t 1 or 3, then this tensor will not be a valid image (also by the pixel values too). So what we do is we map each class index to its corresponding color and then append each class to an output of the form [1, 3, h, w] so that way we can save <code>output</code> as an image in order to visualize it.</p>
<pre><code class="lang-auto">for idx, (x, y) in enumerate(loader):
    x = x.to(device=device)
    with torch.no_grad():
        pred = torch.sigmoid(model(x))
        out = model(x)
        class_to_color = [torch.tensor([1.0, 0.0, 0.0]), ...]
        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, 100, 100
            curr_color = color.reshape(1, 3, 1, 1)
            segment = mask*curr_color # should have shape 1, 3, 100, 100
            output += segment
</code></pre>
<p>This solution works perfectly fine for batch sizes of 1. But I tried a batch size of 8 and then <code>mask</code> has the shape [8, 1, 100, 100] and thus something like <code>segment = mask*curr_color</code> won’t work. Do you know how we can modify this so that way it can work for batch sizes that aren’t 1? Thank you!</p>

eqy: <p>Can you describe in more detail what the problem is? I’m not sure I understand where having a batch size larger than one breaks things, as the max is across the class dimensions and the remaining computations are elementwise or broadcasting across all dimensions.<br>
Out of the box, the code still runs on my end:</p>
<pre><code class="lang-auto">&gt;&gt;&gt; import torch
&gt;&gt;&gt; out = torch.randn(8, 3, 100, 100)
&gt;&gt;&gt; mask = out[:, 1, :, :] == torch.max(out, dim=1)[0]
&gt;&gt;&gt; mask = mask.unsqueeze(1)
&gt;&gt;&gt; curr_color = torch.randn(1, 3, 1, 1)
&gt;&gt;&gt; segment = mask*curr_color
&gt;&gt;&gt; segment.shape
torch.Size([8, 3, 100, 100])
&gt;&gt;&gt;
</code></pre>

hanoody: <p>Actually what I got was the following error, sorry for not being specific:</p>
<pre><code class="lang-auto">output += segment
RuntimeError: output with shape [1, 3, 1000, 1000] doesn't match the broadcast shape [8, 3, 1000, 1000]
</code></pre>
<p>I was able to fix this issue by changing</p>
<pre><code class="lang-auto">output = torch.zeros(1, 3, out.size(-2), out.size(-1), dtype=torch.float)
</code></pre>
<p>to</p>
<pre><code class="lang-auto">output = torch.zeros(out.shape[0], 3, ... ) #this ensures that regardless of the batch size, output image shape will match up
</code></pre>
<p>But now my output image consists of a grid of size batch (8)… So output contains 8 images, do you know how I can iterate through the first dimension and get each of the 8 images from it? Thank you very much!</p>

eqy: <p>Yes, you would need to increase the output size.</p>
<p>You can access at 1,3,1000,1000 image by using an indexing operation</p>
<pre><code class="lang-auto">image0 = output[0,:,:,:]
image1 = output[1,:,:,:]
...
</code></pre>
<p>Saving the images can also be done in this way (a sketch)</p>
<pre><code class="lang-auto">for i in range(8):
    curr_image = output[i,:,:,:]
    save_image(curr_image, i,)
    ....
</code></pre>

hanoody: <p>Works like a charm! Thank you very much!</p>