Model not predicting any values for certain mask in multi-class segmentation

nalinbot: <p>I have trained a multi-class segmentation model on some data and my data was in the form (4, 320, 320)/(4, 512, 512) which denotes 4 classes with each (1, 512, 512) being a binary segmentation of the one class which is how I’m pretty sure it works, if not please correct me. I got my predictions back and everything looks like its working(albeit the predictions aren’t too good which doesn’t matter to me right now) and I my predictions are also in the form of (4, 512, 512) and I can extract each binary mask with slicing eg. <code>img[1, :, :]</code> and I’m wondering how I can combine all of these masks to make a final mask for the images.</p>
<p>Here is an example of the predictions</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/9/29f40c5c6baf2ddcf16bc5db1d9d55e59963a44e.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/29f40c5c6baf2ddcf16bc5db1d9d55e59963a44e" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/9/29f40c5c6baf2ddcf16bc5db1d9d55e59963a44e_2_690x271.jpeg" alt="image" data-base62-sha1="5Z8rJwNDHriBR9Czthg1eEWfofA" width="690" height="271" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/9/29f40c5c6baf2ddcf16bc5db1d9d55e59963a44e_2_690x271.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/2/9/29f40c5c6baf2ddcf16bc5db1d9d55e59963a44e_2_1035x406.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/2/9/29f40c5c6baf2ddcf16bc5db1d9d55e59963a44e_2_1380x542.jpeg 2x" data-dominant-color="886893"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1548×608 37.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
and there are 4 of these all in seperate classes that I would like to combine into one image that looks something like this:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/d/0d08054aa8d83b433ecb95aaf8825deb10ae0dad.png" data-download-href="https://discuss.pytorch.org/uploads/default/0d08054aa8d83b433ecb95aaf8825deb10ae0dad" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/0/d/0d08054aa8d83b433ecb95aaf8825deb10ae0dad.png" alt="image" data-base62-sha1="1RhogfeWotEO0AmokzNpMd2QOv3" width="500" height="500" data-dominant-color="A46359"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">512×512 7.09 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Is there a way to achieve this, if so what can I do?</p>

ptrblck: <p>To get the predicted class index for each pixel position you could use:</p>
<pre><code class="lang-python">output # [batch_size, nb_classes, height, width]
pred = torch.argmax(output, dim=1) # [batch_size, height, width]
</code></pre>
<p>and map each class index to a color code or directly plot it via e.g. <code>plt.imshow</code>.</p>

nalinbot: <p>I understand the argmax part but how would I map the pixel values to RGB values?</p>

ptrblck: <p>You could either pick a <code>colormap</code> in <code>plt.imshow</code> which would work for you or you could map each class index to an RGB value before visualizing the output.</p>