Multiclass Segmentation

ru6928: <p>Hi, is there an example for creating a custom dataset and training for multiclass segmentation using U-Net? I find many examples for binary segmentation but yet to find something for multiclass segmentation. Thank you!</p>

ptrblck: <p>I assume you have already found suitable code snippets for a binary segmentation use case?<br>
If so, you could use it as a base line and make a few changes for a multi class segmentation use case:</p>
<ul>
<li>use <code>nn.CrossEntropyLoss</code> as your criterion</li>
<li>your model should output a tensor with the shape <code>[batch_size, nb_classes, height ,width]</code>
</li>
<li>the <code>target</code> should be a <code>LongTensor</code> with the shape <code>[batch_size, height, width]</code> and contain the class indices for each pixel location in the range <code>[0, nb_classes-1]</code>
</li>
</ul>
<p>Depending on the format of your segmentation mask images, you might need to create a mapping e.g. between color codes and the corresponding class indices.</p>
<p>Let us know, if and where you get stuck.</p>

ru6928: <p>Thank you very much for the reply. I have 2 follow up questions -<br>
1 - in binary, the output of the U-net model will is a softmax(). Will it be the same for multi-class?<br>
2 - my target has 3 classes. So, I should have 0 for background and 1,2,3 for the three classes at each pixel location, right?</p>

ptrblck: <ol>
<li>
<p>If you are using <code>nn.BCELoss</code>, the output should use <code>torch.sigmoid</code> as the activation function. Alternatively, you won’t use any activation function and pass raw logits to <code>nn.BCEWithLogitsLoss</code>. If you use <code>nn.CrossEntropyLoss</code> for the multi-class segmentation, you should also pass the raw logits without using any activation function.</p>
</li>
<li>
<p>Yes, but then you should deal with 4 classes (background + 3 classes), so the output of your model should be <code>[batch_size, 4, h, w]</code>.</p>
</li>
</ol>

Rishav_Sapahia: <p>I have a follow up question on this-<br>
1)Do we need to apply any normalisation on the masks?</p>
<p>Following is my custom data loader for image segmentation-</p>
<pre><code class="lang-auto">class DataLoaderSegmentation(data.Dataset):
    def __init__(self,folder_path,transform = None):
        super(DataLoaderSegmentation, self).__init__()
        self.img_files = glob.glob(os.path.join(folder_path,'images','*.tif'))
        self.mask_files = glob.glob(os.path.join(folder_path,'mask','*.bmp'))
        self.transform = transform
    def __getitem__(self, index):
        img_path = self.img_files[index]
        mask_path = self.mask_files[index]
        data = Image.open(img_path)
        label = Image.open(mask_path)
        if self.transform:
            data = self.transform(data)
        label = np.array(label)
        return data, torch.from_numpy(label).long()
    
    def __len__(self):
        return len(self.img_files)`Preformatted text`
</code></pre>
<p>Also when I don’t use <code>transform.ToTensor</code> normalisation but rather use <code>torch.from_numpy(label).long()</code>, the shape of target comes out to be <code>Batch_size,height,width,num_channels</code><br>
2)Why the difference and how to solve this problem?<br>
3)Also, I couldn’t understand this line much <code>Depending on the format of your segmentation mask images, you might need to create a mapping e.g. between color codes and the corresponding class indices.</code></p>

ptrblck: <p>Most likely you should not apply any normalization on your segmentation masks, as this will distort the class indices.<br>
Could you print the shape of <code>label</code> before passing it to <code>torch.from_numpy</code>?<br>
I would assume the channel in in dim0 or your images don’t have the channel dimension, if you are loading them with <code>PIL</code>.</p>
<p>Have a look at <a href="https://discuss.pytorch.org/t/error-in-python-s-multiprocessing-library/31355/6">this post</a> where I’ve explained the mapping a bit better.<br>
Basically your targets should contain class indices in <code>[0, nb_classes]</code>.<br>
However, sometimes your segmentation images use a color code for certain classes, e.g. red could be a car and blue could be a tree. Using a mapping, you would have to transform these color codes to class indices, e.g. red-&gt;0, blue-&gt;1, …</p>
<p>PS: If you are resizing the mask images, make sure to use nearest neighbor interpolation, as other interpolation techniques might distort the labels/colors.</p>

Rishav_Sapahia: <p>Yes, when I print the shape of image after converting it into numpy array, even reading from both <code>cv2</code> and <code>PIL</code>, the output shape is <code>3*512*512</code> which I corrected using <code>np.transpose(imgs,(2,0,1))</code>.</p>
<p>I got your point regarding the mapping but how to define the mapping for RGB images, my masks have 4 classes- Blue,Red,Green and black background.</p>
<p>I am not resizing the mask images but I am Random Cropping the original image to a smaller size and my final output segmented image is same size as of cropped input image.Is it the right way to do it, are there any performance issues with it?</p>

ptrblck: <p>Have a look at <a href="https://discuss.pytorch.org/t/training-semantic-segmentation/49275/4">this post</a> to see, how to create a custom mapping.</p>
<p><code>RandomCrop</code> should work in the default setup without padding. However, you should make sure to use the same “random” locations for both, your input and target as shown <a href="https://discuss.pytorch.org/t/torchvision-transfors-how-to-perform-identical-transform-on-both-image-and-target/10606/7">here</a>.</p>

Rishav_Sapahia: <p>okay, I followed the post and implemented the following dataloader class keeping in mind the same random locations for both input and target.<br>
First, I want to show the required transformation by me which I have applied in a normal way if I don’t need to apply them simultaneously-</p>
<pre><code class="lang-auto">data_transforms = transforms.Compose([transforms.RandomCrop((512,512)),
                                 transforms.Lambda(gaussian_blur),
                                 transforms.Lambda(elastic_transform),
                                 transforms.RandomRotation([+90,+180]),
                                 transforms.RandomRotation([+180,+270]),
                                 transforms.RandomHorizontalFlip(),
                                 transforms.ToTensor(),
                                 transforms.Normalize(mean=train_mean, std=train_std)
                               ])


</code></pre>
<p>Now, the thing is I can apply <code>RandomCrop, Random Rotation, ToTensor</code> in the class but I have no idea how to apply <code>gaussian_blur,elastic_transform and Normalization with custom mean</code><br>
Also, below is my dataloader class</p>
<pre><code class="lang-auto">class DataLoaderSegmentation(data.Dataset):
    def __init__(self,folder_path):
        super(DataLoaderSegmentation, self).__init__()
        self.img_files = glob.glob(os.path.join(folder_path,'images','*.tif'))
        self.mask_files = glob.glob(os.path.join(folder_path,'mask','*.bmp'))
  
    def mask_to_class(self,mask):
        target = torch.from_numpy(mask)
        h,w = target.shape[0],target.shape[1]
        masks = torch.empty(h, w, dtype=torch.long)
        colors = torch.unique(target.view(-1,target.size(2)),dim=0).numpy()
        target = target.permute(2, 0, 1).contiguous()
        mapping = {tuple(c): t for c, t in zip(colors.tolist(), range(len(colors)))}
        for k in mapping:
            idx = (target==torch.tensor(k, dtype=torch.uint8).unsqueeze(1).unsqueeze(2))
            validx = (idx.sum(0) == 3) 
            masks[validx] = torch.tensor(mapping[k], dtype=torch.long)
        return masks
    
    def transform(self,image,mask):
        i, j, h, w = transforms.RandomCrop.get_params(
        image, output_size=(512, 512))
        image = TF.crop(image, i, j, h, w)
        mask = TF.crop(mask, i, j, h, w)

        # Random horizontal flipping
        if random.random() &gt; 0.5:
            image = TF.hflip(image)
            mask = TF.hflip(mask)
        
        image = TF.rotate(image,90)
        mask = TF.rotate(mask,90)
        image = TF.rotate(image,180)
        mask = TF.rotate(mask,180)
        image = TF.rotate(image,270)
        mask = TF.rotate(mask,270)

        # Transform to tensor
        image = TF.to_tensor(image)
        mask = TF.to_tensor(mask)
        return image, mask
    
    def __getitem__(self, index):
        img_path = self.img_files[index]
        mask_path = self.mask_files[index]
        data = Image.open(img_path)
        label = Image.open(mask_path)
        data,label = self.transform(data,label)
        label = np.array(label)
        mask = self.mask_to_class(label)
        return data,mask
           
    def __len__(self):
        return len(self.img_files)
</code></pre>
<p>I am obtaining the following error -</p>
<pre><code class="lang-auto">RuntimeError: Assertion `input0 == target0 &amp;&amp; input2 == target1 &amp;&amp; input3 == target2' failed. size mismatch (got input: 5x4x504x504, target: 5x584x565) at ../aten/src/THNN/generic/SpatialClassNLLCriterion.c:59

</code></pre>
<p>Because size of my label is now <code>batch_size*height*width</code>, so does this implies I have to do certain change in my network too or is someplace else I am screwing up?</p>

ptrblck: <p>The number of dimensions look alright. However the spatial size is different.<br>
While your data is <code>[504 x 504]</code>, your target is <code>[584 x 565]</code>, which will throw this error.</p>
<p>Also, be careful about using <code>TF.to_tensor</code> on your mask, as this might normalize the class indices into float values in the range <code>[0, 1]</code>.</p>

Rishav_Sapahia: <p>Thanks a lot, there was a problem in the architecture which I resolved.<br>
Yes, I removed <code>TF.to_tensor</code> on the mask though yet to figure out how to do selective transformations on it.<br>
Although I am encountering the following error on running the current network-</p>
<pre><code class="lang-auto">RuntimeError: Assertion `cur_target &gt;= 0 &amp;&amp; cur_target &lt; n_classes' failed.  at/aten/src/THNN/generic/SpatialClassNLLCriterion.c:109
</code></pre>
<p>Now, I am not able to figure out where I am doing wrong?</p>

ptrblck: <p>Make sure the target has the shape <code>[N, H, W]</code> and contains class indices in the range <code>[0, nb_classes-1]</code>, while your output should have the shape <code>[N, nb_classes, H, W]</code>.<br>
Apparently some target tensor contains values outside of this range.<br>
You could use a print statement while iterating your <code>DataLoader</code> to track down the problematic batch.</p>

Adithia_Jo: <p>from  <strong>future</strong>  import print_function<br>
from keras.preprocessing.image import ImageDataGenerator<br>
import numpy as np<br>
import os<br>
import glob<br>
import <a href="http://skimage.io/" rel="nofollow noopener">skimage.io</a> as io<br>
import skimage.transform as trans</p>
<p>def adjustData(img,mask,flag_multi_class,num_class):<br>
if(flag_multi_class):<br>
img = img / 255<br>
mask = mask[:,:,:,0] if(len(mask.shape) == 4) else mask[:,:,0]<br>
new_mask = np.zeros(mask.shape + (num_class,))<br>
for i in range(num_class):<br>
<span class="hashtag">#for</span> one pixel in the image, find the class in mask and convert it into one-hot vector<br>
<span class="hashtag">#index</span> = np.where(mask == i)<br>
<span class="hashtag">#index_mask</span> = (index[0],index[1],index[2],np.zeros(len(index[0]),dtype = np.int64) + i) if (len(mask.shape) == 4) else (index[0],index[1],np.zeros(len(index[0]),dtype = np.int64) + i)<br>
<span class="hashtag">#new_mask</span>[index_mask] = 1<br>
new_mask[mask == i,i] = 1<br>
new_mask = np.reshape(new_mask,(new_mask.shape[0],new_mask.shape[1]*new_mask.shape[2],new_mask.shape[3])) if flag_multi_class else np.reshape(new_mask,(new_mask.shape[0]*new_mask.shape[1],new_mask.shape[2]))<br>
mask = new_mask<br>
elif(np.max(img) &gt; 1):<br>
img = img / 255<br>
mask = mask /255<br>
mask[mask &gt; 0.5] = 1<br>
mask[mask &lt;= 0.5] = 0<br>
return (img,mask)</p>
<p>def trainGenerator(batch_size,train_path,image_folder,mask_folder,aug_dict,image_color_mode = “grayscale”,<br>
mask_color_mode = “grayscale”,image_save_prefix = “image”,mask_save_prefix = “mask”,<br>
flag_multi_class = False,num_class = 2,save_to_dir = None,target_size = (256,256),seed = 1):<br>
‘’’<br>
can generate image and mask at the same time<br>
use the same seed for image_datagen and mask_datagen to ensure the transformation for image and mask is the same<br>
if you want to visualize the results of generator, set save_to_dir = “your path”<br>
‘’’<br>
image_datagen = ImageDataGenerator(**aug_dict)<br>
mask_datagen = ImageDataGenerator(**aug_dict)<br>
image_generator = image_datagen.flow_from_directory(<br>
train_path,<br>
classes = [image_folder],<br>
class_mode = None,<br>
color_mode = image_color_mode,<br>
target_size = target_size,<br>
batch_size = batch_size,<br>
save_to_dir = save_to_dir,<br>
save_prefix = image_save_prefix,<br>
seed = seed)<br>
mask_generator = mask_datagen.flow_from_directory(<br>
train_path,<br>
classes = [mask_folder],<br>
class_mode = None,<br>
color_mode = mask_color_mode,<br>
target_size = target_size,<br>
batch_size = batch_size,<br>
save_to_dir = save_to_dir,<br>
save_prefix = mask_save_prefix,<br>
seed = seed)<br>
train_generator = zip(image_generator, mask_generator)<br>
for (img,mask) in train_generator:<br>
img,mask = adjustData(img,mask,flag_multi_class,num_class)<br>
yield (img,mask)</p>
<p>def valGenerator(batch_size,train_path,image_folder,mask_folder,image_color_mode = “grayscale”,<br>
mask_color_mode = “grayscale”,image_save_prefix = “image”,mask_save_prefix = “mask”,<br>
flag_multi_class = False,num_class = 2,save_to_dir = None,target_size = (256,256),seed = 1):<br>
‘’’<br>
can generate image and mask at the same time<br>
use the same seed for image_datagen and mask_datagen to ensure the transformation for image and mask is the same<br>
if you want to visualize the results of generator, set save_to_dir = “your path”<br>
‘’’<br>
image_datagen = ImageDataGenerator()<br>
mask_datagen = ImageDataGenerator()<br>
image_generator = image_datagen.flow_from_directory(<br>
train_path,<br>
classes = [image_folder],<br>
class_mode = None,<br>
color_mode = image_color_mode,<br>
target_size = target_size,<br>
batch_size = batch_size,<br>
save_to_dir = save_to_dir,<br>
save_prefix = image_save_prefix,<br>
seed = seed)<br>
mask_generator = mask_datagen.flow_from_directory(<br>
train_path,<br>
classes = [mask_folder],<br>
class_mode = None,<br>
color_mode = mask_color_mode,<br>
target_size = target_size,<br>
batch_size = batch_size,<br>
save_to_dir = save_to_dir,<br>
save_prefix = mask_save_prefix,<br>
seed = seed)<br>
train_generator = zip(image_generator, mask_generator)<br>
for (img,mask) in train_generator:<br>
img,mask = adjustData(img,mask,flag_multi_class,num_class)<br>
yield (img,mask)</p>
<p>def testGenerator(test_path,num_image = 18,target_size = (256,256),flag_multi_class = False,as_gray = True):<br>
for i in test_path:<br>
img = io.imread(i, as_gray = as_gray)<br>
print(i)</p>
<h1>img = img / 255</h1>
<pre><code class="lang-auto">    img = trans.resize(img,target_size)
    img = np.reshape(img,img.shape+(1,)) if (not flag_multi_class) else img
    img = np.reshape(img,(1,)+img.shape)
    yield img
</code></pre>
<p>def testGenerator2(batch_size,test_path,target_size = (256,256), image_color_mode = “grayscale”, num_class = 2,flag_multi_class = False,as_gray = True):<br>
image_datagen = ImageDataGenerator()<br>
image_generator = image_datagen.flow_from_directory(<br>
test_path,<br>
class_mode = None,<br>
color_mode = image_color_mode,<br>
target_size = target_size,<br>
batch_size = batch_size<br>
)<br>
return image_generator</p>
<p>def geneTrainNpy(image_path,mask_path,flag_multi_class = False,num_class = 2,image_prefix = “image”,mask_prefix = “mask”,image_as_gray = True,mask_as_gray = True):<br>
image_name_arr = glob.glob(os.path.join(image_path,"%s*.png"%image_prefix))<br>
image_arr = []<br>
mask_arr = []<br>
for index,item in enumerate(image_name_arr):<br>
img = io.imread(item,as_gray = image_as_gray)<br>
img = np.reshape(img,img.shape + (1,)) if image_as_gray else img<br>
mask = io.imread(item.replace(image_path,mask_path).replace(image_prefix,mask_prefix),as_gray = mask_as_gray)<br>
mask = np.reshape(mask,mask.shape + (1,)) if mask_as_gray else mask<br>
img,mask = adjustData(img,mask,flag_multi_class,num_class)<br>
image_arr.append(img)<br>
mask_arr.append(mask)<br>
image_arr = np.array(image_arr)<br>
mask_arr = np.array(mask_arr)<br>
return image_arr,mask_arr</p>
<p>def labelVisualize(num_class,color_dict,img):<br>
img = img[:,:,0] if len(img.shape) == 3 else img<br>
img_out = np.zeros(img.shape + (3,))<br>
for i in range(num_class):<br>
img_out[img == i,:] = color_dict[i]<br>
return img_out</p>
<p>def labelVisualizeBinary(img):<br>
img[img &gt; 0.5] = 1<br>
img[img&lt;= 0.5] = 0<br>
return img</p>
<p>def saveResult(save_path,results_filename,npyfile,flag_multi_class = False,num_class = 2):<br>
print(results_filename[0])<br>
for i,item in enumerate(npyfile):<br>
img = labelVisualize(num_class,COLOR_DICT,item) if flag_multi_class else item[:,:,0]<br>
print(i)<br>
io.imsave(os.path.join(save_path,results_filename[i]),img)</p>
<h4>my problem</h4>
<p>this is my data code , i have problem how to change it to multiclass segmentation,<br>
i divided codes with 3 part , 1 for data(to transform (num classes), 1 for train(like epoch,path_to_train,etc) , and 1 for model (unet)</p>
<p>in my model i was change the loss function to sparse_categorical_crossentropy, and the last layer to softmax,i know my problem in the data codes (how to transform it multiclass (num classes), i has 3 output (black,white, red) that mean my num_classes = 3), but i dont know how to change it? i just change the num_classes but nothing happend , maybe you can solve my problem thank you <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"> (codes up here) thank you</p>

ptrblck: <p>It seems you are using Keras not PyTorch, so I think you will get a better answer on StackOverflow.<br>
I’m not familiar enough with the data loading in Keras. <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=9" title=":confused:" class="emoji" alt=":confused:"></p>

Adithia_Jo: <p>okay , sir do you have complete codes multiclass segmentation ? (data,model,train) ? can you give it to me?  thank you sir <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

Arun_Mohan: <p>Could you please explain how there are 4 channels(classes) in output. Because if we have 3 channels in output and if 3 channels are empty(that means no segmentation).We can automatically consider this case as background.right?So it would work with 3 channels.ie, output can be [batch_size, 3, h, w] for 3 classes + BG class. right? please correct me if iam wrong.</p>

ptrblck: <p>This wouldn’t work for a multi-class segmentation with <code>nn.CrossEntropyLoss</code> or <code>nn.NLLLoss</code>, since the target has to contain a class index.<br>
With three classes the class indices would be <code>[0, 1, 2]</code> and that for each pixel one of these classes would be active.</p>
<p>If you consider your use case a multi-label segmentation, you could use <code>nn.BCEWithLogitsLoss</code>, which would allow you to define zero, one, or more active classes for each pixel.<br>
This approach would work for 3 explicit channels, as no active class could be considered the “background class”. However, since each pixel could also have all classes set to active, this would be a multi-label segmentation.</p>

sky: <p>Thanks you. But I’m not understand when we deal with 4 classes (background + 3 classes),the output of model should be <code>[batch_size, 4, h, w]</code>，I what to know  1) if the target should use the onehot encode and 2) what the target should like ? ([batch_size, height, width] using onehot encode or no use)</p>

ptrblck: <p>If you are using <code>nn.CrossEntropyLoss</code> and your model outputs the mentioned shape, your target should have the shape <code>[batch_size, height, width]</code> and contain class indices in the range <code>[0, 3]</code>, i.e. it should not be one-hot encoded.</p>

Julien_Guegan: <p>What if I use the DiceLoss ? Should it be one-hot encoded in this case ?</p>