How to set mask labels for mask R-CNN so that I can fine-tune it into a 3-classes classification and segementation model?

username: crissallan
cooked: <p>Hello, I recetently followed the tutorial on the Pytorch Official Webset（<a href="https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html" rel="nofollow noopener">Link</a>） to fine-tune a Mask R-CNN.</p>
<p>In the tutorial,  the author just need to seperate the pedstrain and the background. So the definition of the Dataset class looks like this:<br>
<em>class PennFudanDataset(object):</em></p>
<ul>
<li>
<p>def <strong>init</strong>(self, root, transforms):*</p>
</li>
<li>
<pre><code>   self.root = root*
</code></pre>
</li>
<li>
<pre><code>   self.transforms = transforms*
</code></pre>
</li>
<li>
<pre><code>   # load all image files, sorting them to*
</code></pre>
</li>
<li>
<pre><code>   # ensure that they are aligned*
</code></pre>
</li>
<li>
<pre><code>   self.imgs = list(sorted(os.listdir(os.path.join(root, "PNGImages"))))*
</code></pre>
</li>
<li>
<pre><code>   self.masks = list(sorted(os.listdir(os.path.join(root, "PedMasks"))))*
</code></pre>
</li>
<li>
<p>def <strong>getitem</strong>(self, idx):*</p>
</li>
<li>
<pre><code>   # load images ad masks*
</code></pre>
</li>
<li>
<pre><code>   img_path = os.path.join(self.root, "PNGImages", self.imgs[idx])*
</code></pre>
</li>
<li>
<pre><code>   mask_path = os.path.join(self.root, "PedMasks", self.masks[idx])*
</code></pre>
</li>
<li>
<pre><code>   img = Image.open(img_path).convert("RGB")*
</code></pre>
</li>
<li>
<pre><code>   # note that we haven't converted the mask to RGB,*
</code></pre>
</li>
<li>
<pre><code>   # because each color corresponds to a different instance*
</code></pre>
</li>
<li>
<pre><code>   # with 0 being background*
</code></pre>
</li>
<li>
<pre><code>   mask = Image.open(mask_path)*
</code></pre>
</li>
<li>
<pre><code>   # convert the PIL Image into a numpy array*
</code></pre>
</li>
<li>
<pre><code>   mask = np.array(mask)*
</code></pre>
</li>
<li>
<pre><code>   # instances are encoded as different colors*
</code></pre>
</li>
<li>
<pre><code>   obj_ids = np.unique(mask)*
</code></pre>
</li>
<li>
<pre><code>   # first id is the background, so remove it*
</code></pre>
</li>
<li>
<pre><code>   obj_ids = obj_ids[1:]*
</code></pre>
</li>
<li>
<pre><code>   # split the color-encoded mask into a set*
</code></pre>
</li>
<li>
<pre><code>   # of binary masks*
</code></pre>
</li>
<li>
<pre><code>   masks = mask == obj_ids[:, None, None]*
</code></pre>
</li>
<li>
<pre><code>   # get bounding box coordinates for each mask*
</code></pre>
</li>
<li>
<pre><code>   num_objs = len(obj_ids)*
</code></pre>
</li>
<li>
<pre><code>   boxes = []*
</code></pre>
</li>
<li>
<pre><code>   for i in range(num_objs):*
</code></pre>
</li>
<li>
<pre><code>       pos = np.where(masks[i])*
</code></pre>
</li>
<li>
<pre><code>       xmin = np.min(pos[1])*
</code></pre>
</li>
<li>
<pre><code>       xmax = np.max(pos[1])*
</code></pre>
</li>
<li>
<pre><code>       ymin = np.min(pos[0])*
</code></pre>
</li>
<li>
<pre><code>       ymax = np.max(pos[0])*
</code></pre>
</li>
<li>
<pre><code>       boxes.append([xmin, ymin, xmax, ymax])*
</code></pre>
</li>
<li>
<pre><code>   # convert everything into a torch.Tensor*
</code></pre>
</li>
<li>
<pre><code>   boxes = torch.as_tensor(boxes, dtype=torch.float32)*
</code></pre>
</li>
<li>
<pre><code>   # there is only one class*
</code></pre>
</li>
<li>
<pre><code>   labels = torch.ones((num_objs,), dtype=torch.int64)*
</code></pre>
</li>
<li>
<pre><code>   masks = torch.as_tensor(masks, dtype=torch.uint8)*
</code></pre>
</li>
<li>
<pre><code>   image_id = torch.tensor([idx])*
</code></pre>
</li>
<li>
<pre><code>   area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])*
</code></pre>
</li>
<li>
<pre><code>   # suppose all instances are not crowd*
</code></pre>
</li>
<li>
<pre><code>   iscrowd = torch.zeros((num_objs,), dtype=torch.int64)*
</code></pre>
</li>
<li>
<pre><code>   target = {}*
</code></pre>
</li>
<li>
<pre><code>   target["boxes"] = boxes*
</code></pre>
</li>
<li>
<pre><code>   target["labels"] = labels*
</code></pre>
</li>
<li>
<pre><code>   target["masks"] = masks*
</code></pre>
</li>
<li>
<pre><code>   target["image_id"] = image_id*
</code></pre>
</li>
<li>
<pre><code>   target["area"] = area*
</code></pre>
</li>
<li>
<pre><code>   target["iscrowd"] = iscrowd*
</code></pre>
</li>
<li>
<pre><code>   if self.transforms is not None:*
</code></pre>
</li>
<li>
<pre><code>       img, target = self.transforms(img, target)*
</code></pre>
</li>
<li>
<pre><code>   return img, target*
</code></pre>
</li>
<li>
<p>def <strong>len</strong>(self):*</p>
</li>
<li>
<pre><code>   return len(self.imgs)*
</code></pre>
</li>
</ul>
<p>However, in my case, I need to seperate the nose, mouth  and background from face images.</p>
<p>I’d like to know that how should I change definition of Dataset class,(especially the <strong>getitem</strong> method), then my module could seperate the nose and mouth as different classes after training?</p>

username: ptrblck
cooked: <p>How did you store the target masks?<br>
If each class is using a separate class index, you could reuse most of the code and create the corresponding <code>labels</code> tensor based on the current class index.</p>

username: crissallan
cooked: <p>For my mask data, each region of the face has its own binary mask image.(e.g. if the image name is 0.jpg, then the related masks are: 0_nose.png, 0_mouth.png).<br>
In the Dataset class, I followed the methodology of the tutuorial. I <strong>convert each binary mask image into a [1, 512, 512] tensor，then concatenate the mouth and nose mask tensor along with 0 dim i.e my target masks is a [2, 512, 512] tensor.</strong> Then I set the labels tensor as torch.tensor([1, 2]).</p>
<p>Is it right for me to stack different classes masks along with 0-dim?</p>

username: ptrblck
cooked: <p>If you are using <code>nn.CrossEntropyLoss</code> (or <code>nn.NLLLoss</code>), then the target mask shape should be <code>[batch_size, height, width]</code>. I would therefore recommend to transform the binary masks to masks containing the class index (e.g. 0 for background, 1 for mouth and 2 for nose) and add them to a single target mask.</p>

username: crissallan
cooked: <aside class="quote no-group" data-username="ptrblck" data-post="4" data-topic="66164">
<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/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>add them to a single target mask.</p>
</blockquote>
</aside>
<p>does this mean the channel dim of the mask tensor is just 1? I thought the size of the mask shape is [batch_size, channel, height, width]</p>

username: ptrblck
cooked: <p>The target for the mentioned criteria does not have a channel dimension, as it’s not one-hot encoded.<br>
Instead it contains the class indices directly, so for a segmentation use case, your target should have the shape <code>[batch_size, height, width]</code>.</p>

username: maren11
cooked: <p>Hi there,<br>
I have the same problem as you. Can you tell me which codelines you changed?</p>