Convert segmentation mask of shape [224,224,3] to mask [224,224,classes]

Amasu: <p>Hi I am having problem while converting rgb mask of shape [224,224,3] to mask of shape [224,224,3]. I have attached the code below.</p>
<h2>I am getting masks of shape [224,224,classes] but lose information of classes in channels only one channel has some mask will others don’t</h2>
<p>Code:</p>
<pre><code>class CamVid_Dataset():

   def __init__(self, img_pth, mask_pth, transform):

      self.img_pth = img_pth
     self.mask_pth = mask_pth
     self.transform = transform
     all_imgs = os.listdir(self.img_pth)
     all_masks = os.listdir(self.mask_pth)
     self.total_imgs = natsort.natsorted(all_imgs)
     self.total_masks = natsort.natsorted(all_masks)

def __len__(self):
     return len(self.total_imgs)

def __getitem__(self, idx):
     img_loc = os.path.join(self.img_pth, self.total_imgs[idx])
     image = Image.open(img_loc).convert("RGB")
     tensor_image = self.transform(image)

     mask_loc = os.path.join(self.mask_pth, self.total_masks[idx])
     mask = Image.open(mask_loc).convert("RGB")
     tensor_mask = self.transform(mask)
     tensor_mask = rgb_to_mask(np.array(tensor_mask).transpose(1,2,0), id2code)
     return tensor_image, tensor_mask
</code></pre>
<p><span class="hashtag">#Define</span> transforms for the training data and validation data</p>
<pre><code>train_transforms = transforms.Compose([ transforms.Resize((224,224)), transforms.ToTensor()])
</code></pre>
<p>Pass transform here-in</p>
<pre><code>train_data = CamVid_Dataset(img_pth = path + 'train/', mask_pth = path + 'train_labels/', transform = train_transforms)
</code></pre>
<p>Data loaders</p>
<pre><code>trainloader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, shuffle=False)

inputs, mask = next(iter(trainloader))
</code></pre>
<p>RGB2MASK:</p>
<pre><code>def rgb_to_mask(img, color_map):

''' 

    Converts a RGB image mask of shape [batch_size, h, w, 3] to Binary Mask of shape [batch_size, classes, h, w]

    Parameters:

        img: A RGB img mask

        color_map: Dictionary representing color mappings

    returns:

        out: A Binary Mask of shape [batch_size, classes, h, w]

'''

     num_classes = len(color_map)
     shape = img.shape[:2]+(num_classes,)
     out = np.zeros(shape, dtype=np.int8)
     for i, cls in enumerate(color_map):
         out[:,:,i] = np.all(img.reshape( (-1,3) ) == color_map[i], axis=1).reshape(shape[:2])
     return out</code></pre>

ptrblck: <p>I assume your mask images contain a specific color code, which would be used to create the corresponding class indices.<br>
If that’s the case, be careful with <code>transforms.Resize((224,224))</code>, as the default interpolation would be a bilinear interpolation, which could corrupt some mask colors.<br>
To resize a mask image you should this use <code>PIL.Image.NEAREST</code>.<br>
Let me know, if this helps.</p>

Amasu: <p>Hi, Thanks for responding.</p>
<ul>
<li>Well I tried what you suggested but didn’t work. But it seems that using transform.PILToTensor (for masks) instead of transforms.Tensor works fine as the masks don’t lose information then. Can you tell me why ?</li>
<li>Another thing I am working on CamVid dataset with 32 classes. Initial mask provided was of size [h,w,3] I converted it into [h,w,32] where each channel represent corresponding class. Is this approach correct ?</li>
<li>Also, I am not sure about if I should use Sigmoid or Softmax function. And Currently I am using loss BCE<br>
The prediction on test is not good<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/b/a/ba0fdc6cf98f705456487932e56c16e44a5ffab6.png" alt="download" data-base62-sha1="qxYN4Zi6noOfMectxu2KoxrTLJY" width="262" height="252"></li>
</ul>
<p>UNet:</p>
<pre><code>def Double_Conv(input_channel, output_channel):
     return nn.Sequential(
          nn.Conv2d(input_channel, output_channel, kernel_size=3, padding=1),
          nn.BatchNorm2d(output_channel),
          nn.ReLU(inplace=True),
          nn.Conv2d(output_channel, output_channel, kernel_size=3, padding=1),
          nn.BatchNorm2d(output_channel),
          nn.ReLU(inplace=True))

class UNet(nn.Module):
      def __init__(self, no_of_cls):
        super().__init__()

        self.down1 = Double_Conv(3, 64)
        self.down2 = Double_Conv(64, 128)
        self.down3 = Double_Conv(128, 256)
        self.down4 = Double_Conv(256, 512)
        self.down5 = Double_Conv(512, 1024)

        self.max_pool = nn.MaxPool2d(2)
        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)

        self.up4 = Double_Conv(512 + 1024, 512)
        self.up3 = Double_Conv(256 + 512, 256)
        self.up2 = Double_Conv(128 + 256, 128)
        self.up1 = Double_Conv(64 + 128, 64)

        self.lastConv = nn.Conv2d(64, no_of_cls, 1)

def forward(self, x):
  conv1 = self.down1(x)
  x = self.max_pool(conv1)

  conv2 = self.down2(x)
  x = self.max_pool(conv2)

  conv3 = self.down3(x)
  x = self.max_pool(conv3)

  conv4 = self.down4(x)
  x = self.max_pool(conv4)

  x = self.down5(x)

  x = self.upsample(x)
  x = torch.cat([x, conv4], dim=1)

  x = self.up4(x)
  x = self.upsample(x)
  x = torch.cat([x, conv3], dim=1)

  x = self.up3(x)
  x = self.upsample(x)
  x = torch.cat([x, conv2], dim=1)

  x = self.up2(x)
  x = self.upsample(x)
  x = torch.cat([x, conv1], dim=1)

  x = self.up1(x)
  out = self.lastConv(x)
  return torch.sigmoid(out)
</code></pre>

ptrblck: <ol>
<li>
<p>I’m not aware of the <code>transforms.PILToTensor</code> or <code>transforms.Tensor</code> methods. Do you mean <code>transforms.ToTensor</code>? If so, what are you comparing it against?</p>
</li>
<li>
<p>That’s incorrect for a multi-class segmentation, where each pixel belongs to one specific class only. In that case the target tensor should have the shape <code>[batch_size, height, width]</code> and contain the class indices in the range <code>[0, nb_classes-1]</code>. Since your tensor seems to be a one-hot encoded mask, you could simply call <code>target = one_hot_target.argmax(2)</code>.</p>
</li>
<li>
<p>For a multi-class segmentation you should use <code>nn.CrossEntropyLoss</code>, remove the last sigmoid and pass the raw logits in the shape <code>[batch_size, nb_classes, height, width]</code> to the criterion.</p>
</li>
</ol>

Amasu: <ol>
<li>Sorry, not transforms.Tensor it is transform.ToTensor to convert PIL Image or numpy.ndarray (H x W x C) in the range [0, 255] to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0]</li>
<li>Can you explain a bit about how to deal with masks in case of multi-class segmentation. I am confused about this particularly. Currently my each of 32 channel has pixel value 0-1. and each channel of mask shows a particular class.</li>
<li>I’ll do that</li>
</ol>

ptrblck: <p>Since your current mask is a one-hot encoded tensor (each channel represents a class, where <code>1</code> denotes an active class), you could transform it into the desired class mask via:</p>
<pre><code class="lang-python">target = torch.argmax(mask, dim=1)
</code></pre>
<p>Note that <code>dim=1</code> is used, if your channel dimension (class channels) is in <code>dim1</code>.</p>

Amasu: <p>Yes, thank you. Finally, I got it.<br>
Can you tell me how can I improve performance of my Network. These are the current results that I a getting.</p><aside class="quote quote-modified" data-post="1" data-topic="93407">
  <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/amasu/48/31695_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/unet-segmentation-results-not-good/93407">UNet Segmentation Results not good</a> <a class="badge-category__wrapper " href="/c/vision/5"><span data-category-id="5" style="--category-badge-color: #AB9364; --category-badge-text-color: #FFFFFF;" data-drop-close="true" class="badge-category " title="Topics related to either pytorch/vision or vision research related topics"><span class="badge-category__name">vision</span></span></a>
  </div>
  <blockquote>
    Hi, I am trying to work on Segmentation with Unet on CamVid Dataset. Data has 32 classes, masks available are of [n,3,h,w]. I converted the masks in [n,32,h,w]. While during training using CrossEntropyloss Ii game target to criterion as target=target.argmax(1) which makes target of shape [n,1,h,w] while my model predictions are of shape [n,32,h,w]. But when test on test data result is not good. 
[download (1)] 
[download (2)] 
My code is Given: 
def Double_Conv(input_channel, output_channel):
re…
  </blockquote>
</aside>


sigma_x: <p>I’m still a bit confused about the difference between <code>torchvision.transforms.Resize</code> and <code>torch.functional.interpolate</code> for masks (assuming masks have uint8 format and contains only integers, obviosuly).</p>
<p>The problem is, I kept getting suspicious results using Mask R-CNN:</p>
<p><code>mask = np.array(PILImage.open(mask_name).resize(self.img_size, resample=PILImage.NEAREST))</code></p>
<p>It seems that <code>transforms.Resize</code> is identical to <code>PIL.Image.Resize</code>. On the other hand, <code>functional.interpolate(mode=NEAREST, align_corners=True)</code> uses opencv or some other solution, and in fact seems to give the correct resizing. I think so because <code>torchvision.detection.transforms</code> uses this method both for the image <code>(mode=bilinear, align_corners=False)</code> and masks <code>(mode=nearest, align_corneres=True)</code>.</p>
<p>So I don’t quite understand either the difference, nor the correct use of <code>transforms.Resize</code> for masks, apart from the fact that <code>transforms.Resize</code> is applied to <code>PIL.Image</code> and <code>functional.interpolate</code> to tensors…</p>

ptrblck: <p>There might be differences in the interpolation results between <code>PIL</code>'s implementation and the native one (I haven’t checked it recently). To make sure the same algorithm is used, you could stick to one library for both, the input and mask.</p>

sigma_x: <p>Thanks, does it mean that <code>F.interpolate</code> uses an algorithm different to PIL.Image.resize? I looked at the source code, but couldn’t identify any.</p>
<p>Does this also mean that if I don’t use masks (e.g. for classification), I can stick to transforms.Resize?</p>

ptrblck: <p>It could use a different algorithms and you could compare the outputs for different interpolation methods. If I remember it correctly, at least OpenCV and PIL used different methods and I guess that the internal PyTorch method might have used one of them.</p>

sigma_x: <p>Thanks. My main concern was related to masks, not images, as I got the impression algorithms other that <code>F.interpolate</code> produced suspicious results, incl. <code>transforms.Resize (also numpy.resize, PILImage.resize</code>, etc)</p>

ptrblck: <p>Which concern do you have about the mentioned methods and how do these suspicious results look like or what makes them suspicious?</p>

sigma_x: <p>Other methods somehow returned other pixel values, e.g. [3,4] instead of [4]. It’s algood now.</p>