Convert multi-class segmentation mask into integer indexed

zzzf: <p>I want to use ConfusionMatrix in pytorch_lightning.metrics: <a href="https://pytorch-lightning.readthedocs.io/en/latest/_modules/pytorch_lightning/metrics/functional/confusion_matrix.html#confusion_matrix" class="inline-onebox" rel="noopener nofollow ugc">pytorch_lightning.metrics.functional.confusion_matrix — PyTorch Lightning 1.1.6 documentation</a>.</p>
<p>I’m working on a multi-class segmentation problem where my mask originally is RGB valued (N, 3, H, W) and I converted it into one-hot (N, <span class="hashtag">#Class</span>, H, W) for training. However, pytorch_lightning.metrics.utils._input_format_classification used in ConfusionMatrix, requires that my mask is integer indexed: (N, H, W). I’m wondering how can I efficiently convert transform an RGB valued mask into integer indexed?</p>

KFrank: <p>Hi zzzf!</p>
<aside class="quote no-group quote-modified" data-username="zzzf" data-post="1" data-topic="110667" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/zzzf/48/14158_2.png" class="avatar"> zzzf:</div>
<blockquote>
<p>I converted it into one-hot (N, <span class="hashtag-raw">#Class</span>, H, W) for training.<br>
…<br>
how can I efficiently convert transform an RGB valued mask into integer indexed?</p>
</blockquote>
</aside>
<p>To convert a one-hot-encoded class-label mask into an integer-class-label<br>
mask (e.g., <code>shape = [N, nClass, H, W]</code> → <code>shape = [N, H, W]</code>), you<br>
take <code>torch.argmax (one_hot_mask, dim = 1)</code> along the class dimension.</p>
<p>As for your question about how to convert an RGB-valued mask, the issue<br>
is how your class labels are encoded into your color channels (which is<br>
something you haven’t told us).  You have said that you can convert RGB<br>
into one-hot, so you could apparently use a two-step procedure where you<br>
first convert RGB to one-hot, and then apply <code>argmax()</code>.</p>
<p>Best.</p>
<p>K. Frank</p>

zzzf: <p>Hi KFrank,</p>
<p>Thanks for your reply!</p>
<p>This two-step is indeed a solution but somewhat not optimal. Is there a way to directly convert the RGB valued mask to inter-indexed mask? Just for example, the class RGB values are (255, 0, 0), (0, 255, 0), (0, 0, 255), (0, 0, 0), (255, 255, 255).</p>

KFrank: <p>Hi zzzf!</p>
<aside class="quote no-group" data-username="zzzf" data-post="3" data-topic="110667" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/zzzf/48/14158_2.png" class="avatar"> zzzf:</div>
<blockquote>
<p>Is there a way to directly convert the RGB valued mask to inter-indexed mask?</p>
</blockquote>
</aside>
<p>There are a couple of ways of going about this.  Probably the most<br>
flexible and most logically straightforward will be to index into a lookup<br>
table with your three-channel colors.</p>
<p>Here is a demonstration script:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)
# convert rgb mask to integer class labels
#   red    =  class-0
#   green  =  class-1
#   blue   =  class-2
#   black  =  class-3
#   white  =  class-4
#   (bad)  =  class-5
# create example mask tensor of shape [nBatch, nRGB = 3, height]
nBatch = 2
height = 4
rgb_mask = 255 * torch.ones ((nBatch, 3, height), dtype = torch.int64)
rgb_mask[0, :, 0]       =    0   # black
rgb_mask[0, [1, 2], 1]  =    0   # red
rgb_mask[0, [0, 2], 2]  =    0   # green
rgb_mask[0, [0, 1], 3]  =    0   # blue
rgb_mask[1, :, 0]       =  255   # white
rgb_mask[1, :, 1]       =    0   # black
rgb_mask[1, [0, 1], 2]  =    0   # blue
rgb_mask[1, [1, 2], 3]  =    0   # red
print (rgb_mask.shape)
print (rgb_mask)
# build lookup table
lut_dim = torch.tensor (rgb_mask.shape)
lut_dim[1] = -1
lut = torch.tensor ([3, 0, 1, 5, 2, 5, 5, 4])  # color-class encoding
lut = lut.unsqueeze (-1).expand (lut_dim.tolist())
# convert rgb colors to 0-7 indices
label_mask = torch.sign (rgb_mask)
powers2 = torch.tensor ([2])**torch.arange (3)
label_mask = (powers2 * label_mask.transpose (1, -1)).transpose (1, -1).sum (dim = 1, keepdim = True)
# index into lut to get integer class labels
label_mask = torch.gather (lut, 1, label_mask).squeeze()
print (label_mask.shape)
print (label_mask)
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">1.7.1
torch.Size([2, 3, 4])
tensor([[[  0, 255,   0,   0],
         [  0,   0, 255,   0],
         [  0,   0,   0, 255]],

        [[255,   0,   0, 255],
         [255,   0,   0,   0],
         [255,   0, 255,   0]]])
torch.Size([2, 4])
tensor([[3, 0, 1, 2],
        [4, 3, 2, 0]])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>