IndexError: Target 255 is out of bounds

username: Hassan_Khanzae
cooked: <p>I am training a Unet model for the segmentation of medical images. I know there is some issue in getting the loss due to the difference in target and input labels. how can I fix this error?</p>
<p>class Segmenter(pl.LightningModule):<br>
def <strong>init</strong>(self):<br>
super().<strong>init</strong>()</p>
<pre><code>    self.model = UNet()
    
    self.optimizer = torch.optim.Adam(self.model.parameters(), lr=1e-4)
    self.loss_fn = torch.nn.CrossEntropyLoss()

def forward(self, data):
    pred = self.model(data)
    return pred

def training_step(self, batch, batch_idx):
    # You can obtain the raw volume arrays by accessing the data attribute of the subject
    img = batch["MR"]["data"]
    mask = batch["Label"]["data"][:,0]  # Remove single channel as CrossEntropyLoss expects NxHxW
    mask = mask.long()
    
    pred = self(img)
    loss = self.loss_fn(pred, mask)
    
    # Logs
    self.log("Train Loss", loss)
    if batch_idx % 50 == 0:
        self.log_images(img.cpu(), pred.cpu(), mask.cpu(), "Train")
    return loss
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/b/cb04e795802207c548d939455809dea51dabad95.png" data-download-href="https://discuss.pytorch.org/uploads/default/cb04e795802207c548d939455809dea51dabad95" title="321"><img src="https://discuss.pytorch.org/uploads/default/original/3X/c/b/cb04e795802207c548d939455809dea51dabad95.png" alt="321" data-base62-sha1="sXZqWE7FPmhYH5ctnwjYJavfwLX" width="690" height="259" data-dominant-color="F4F5F4"><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">321</span><span class="informations">1117×420 25 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>

username: ptrblck
cooked: <p><code>nn.CrossEntropyLoss</code> expects targets containing class indices in <code>[0, nb_classes-1]</code>. In your use case it seems 255 is not a valid class index and the model does not output logits for &gt;=256 classes so you should double check how the label is created.<br>
Based on its value I would guess you are simply loading the target image in <code>uint8</code> without assigning class indices to each pixel.</p>

username: Hassan_Khanzae
cooked: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/9/49a683b709343efa5f21b61c32d093018480fa5a.png" data-download-href="https://discuss.pytorch.org/uploads/default/49a683b709343efa5f21b61c32d093018480fa5a" title="label"><img src="https://discuss.pytorch.org/uploads/default/original/3X/4/9/49a683b709343efa5f21b61c32d093018480fa5a.png" alt="label" data-base62-sha1="avxBSPCDThrV3BkqnsBZHRpJrd0" width="690" height="230" data-dominant-color="F2F2F2"><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">label</span><span class="informations">1077×360 13.1 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>so, how to resolve this?</p>

username: Hassan_Khanzae
cooked: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/7/6763379d42438a443a86d0e38ca40ae22dc0aab7.png" data-download-href="https://discuss.pytorch.org/uploads/default/6763379d42438a443a86d0e38ca40ae22dc0aab7" title="label2"><img src="https://discuss.pytorch.org/uploads/default/original/3X/6/7/6763379d42438a443a86d0e38ca40ae22dc0aab7.png" alt="label2" data-base62-sha1="eKBKAsmb7dK1bgOotGn6luHLSOb" width="690" height="260" data-dominant-color="F1F1F1"><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">label2</span><span class="informations">872×329 10.3 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>

username: ptrblck
cooked: <p>Please post a minimal and executable code snippet by wrapping it into three backticks ``` to reproduce the issue instead of screenshot.</p>

username: Hassan_Khanzae
cooked: <pre><code class="lang-auto">path = Path("T1/image/")
subjects_paths = list(path.glob("image_*"))
subjects = []

for subject_path in subjects_paths:
    label_path = change_img_to_label_path(subject_path)
    subject = tio.Subject({"MR":tio.ScalarImage(subject_path), "Label":tio.LabelMap(label_path)})
    subjects.append(subject)


train_dataset = tio.SubjectsDataset(subjects[:450], transform=train_transform)
val_dataset = tio.SubjectsDataset(subjects[450:], transform=val_transform)

sampler = tio.data.LabelSampler(patch_size=96, label_name="Label", label_probabilities={0:0.2, 1:0.3, 2:0.5}) 


</code></pre>

username: Hassan_Khanzae
cooked: <p>yes, I am loading the target image in <code>uint8</code> , so how can I assign class indices to each pixel?</p>

username: ptrblck
cooked: <p>Your code is unfortunately not executable.</p>
<aside class="quote no-group" data-username="Hassan_Khanzae" data-post="7" data-topic="192468">
<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/hassan_khanzae/48/65715_2.png" class="avatar"> Hassan_Khanzae:</div>
<blockquote>
<p>so how can I assign class indices to each pixel?</p>
</blockquote>
</aside>
<p>Depending on your approach to load data either assign labels inside the <code>__getitem__</code> or use a pre-dwfined dataset, such as <code>ImageFolder</code>.</p>

username: Hassan_Khanzae
cooked: <p>i upload the full code. maybe then you can understand.</p>
<pre><code class="lang-auto">def change_img_to_label_path(path):
    """
    Replace data with mask to get the masks
    """
    parts = list(path.parts)
    parts[parts.index("image")] = "mask"
    return Path(*parts)
path = Path("T1/image/")
subjects_paths = list(path.glob("image_*"))
subjects = []

for subject_path in subjects_paths:
    label_path = change_img_to_label_path(subject_path)
    
    subject = tio.Subject({"MR":tio.ScalarImage(subject_path), "Label":tio.LabelMap(label_path)})
    
    subjects.append(subject)
for subject in subjects:
    assert subject["MR"].orientation == ("L", "P", "S")
process = tio.Compose([
            tio.CropOrPad((256, 256, 200)),
            tio.RescaleIntensity((-1, 1))
            ])


augmentation = tio.RandomAffine(scales=(0.9, 1.1), degrees=(-10, 10))


val_transform = process
train_transform = tio.Compose([process, augmentation])

train_dataset = tio.SubjectsDataset(subjects[:650], transform=train_transform)
val_dataset = tio.SubjectsDataset(subjects[650:], transform=val_transform)

sampler = tio.data.LabelSampler(patch_size=96, label_name="Label", label_probabilities={0:0.2, 1:0.3, 2:0.5})
train_patches_queue = tio.Queue(
     train_dataset,
     max_length=40,
     samples_per_volume=5,
     sampler=sampler,
     num_workers=4,
    )

val_patches_queue = tio.Queue(
     val_dataset,
     max_length=40,
     samples_per_volume=5,
     sampler=sampler,
     num_workers=4,
    )
batch_size = 2

train_loader = torch.utils.data.DataLoader(train_patches_queue, batch_size=batch_size, num_workers=0)
val_loader = torch.utils.data.DataLoader(val_patches_queue, batch_size=batch_size, num_workers=0)
class Segmenter(pl.LightningModule):
    def __init__(self):
        super().__init__()
        
        self.model = UNet()
        
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=1e-4)
        self.loss_fn = torch.nn.CrossEntropyLoss()
    
    def forward(self, data):
        pred = self.model(data)
        return pred
    
    def training_step(self, batch, batch_idx):
        
        # You can obtain the raw volume arrays by accessing the data attribute of the subject
        img = batch["MR"]["data"]
        mask = batch["Label"]["data"][:,0]  # Remove single channel as CrossEntropyLoss expects NxHxW
        mask = mask.long()
        
        pred = self(img)
        loss = self.loss_fn(pred, mask)
        
        # Logs
        self.log("Train Loss", loss)
        if batch_idx % 50 == 0:
            self.log_images(img.cpu(), pred.cpu(), mask.cpu(), "Train")
        return loss
    
        
    def validation_step(self, batch, batch_idx):
        # You can obtain the raw volume arrays by accessing the data attribute of the subject
        img = batch["MR"]["data"]
        mask = batch["Label"]["data"][:,0]  # Remove single channel as CrossEntropyLoss expects NxHxW
        mask = mask.long()
        
        pred = self(img)
        loss = self.loss_fn(pred, mask)
        
        # Logs
        self.log("Val Loss", loss)
        self.log_images(img.cpu(), pred.cpu(), mask.cpu(), "Val")
        
        return loss

    
    def log_images(self, img, pred, mask, name):
        
        results = []
        pred = torch.argmax(pred, 1) # Take the output with the highest value
        axial_slice = 50  # Always plot slice 50 of the 96 slices
        
        fig, axis = plt.subplots(1, 2)
        axis[0].imshow(img[0][0][:,:,axial_slice], cmap="bone")
        mask_ = np.ma.masked_where(mask[0][:,:,axial_slice]==0, mask[0][:,:,axial_slice])
        axis[0].imshow(mask_, alpha=0.6)
        axis[0].set_title("Ground Truth")
        
        axis[1].imshow(img[0][0][:,:,axial_slice], cmap="bone")
        mask_ = np.ma.masked_where(pred[0][:,:,axial_slice]==0, pred[0][:,:,axial_slice])
        axis[1].imshow(mask_, alpha=0.6, cmap="autumn")
        axis[1].set_title("Pred")

        self.logger.experiment.add_figure(f"{name} Prediction vs Label", fig, self.global_step)

            
    
    def configure_optimizers(self):
        #Caution! You always need to return a list here (just pack your optimizer into one :))
        return [self.optimizer]

</code></pre>

username: Hassan_Khanzae
cooked: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/5/2531ac46071385580ab06f924d8e8841c14d1b7e.png" data-download-href="https://discuss.pytorch.org/uploads/default/2531ac46071385580ab06f924d8e8841c14d1b7e" title="problem"><img src="https://discuss.pytorch.org/uploads/default/original/3X/2/5/2531ac46071385580ab06f924d8e8841c14d1b7e.png" alt="problem" data-base62-sha1="5j25AIxVOmzCDGEBiwfMf0wvDBA" width="690" height="390" data-dominant-color="F5F5F5"><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">problem</span><span class="informations">775×439 8.07 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>

username: Hassan_Khanzae
cooked: <p>Here are the min and max tensors labels<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/4/3436dad9aff87677a70414c86abdbd40048449f1.png" data-download-href="https://discuss.pytorch.org/uploads/default/3436dad9aff87677a70414c86abdbd40048449f1" title="max_label"><img src="https://discuss.pytorch.org/uploads/default/original/3X/3/4/3436dad9aff87677a70414c86abdbd40048449f1.png" alt="max_label" data-base62-sha1="7rUlYJNScp9Ad1Zmo2ccJrSfmEx" width="690" height="282" data-dominant-color="F6F6F6"><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">max_label</span><span class="informations">811×332 5.02 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><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/3/b/3bcb663b85c446f7d5f583d3f5ec573fa1ef9430.png" data-download-href="https://discuss.pytorch.org/uploads/default/3bcb663b85c446f7d5f583d3f5ec573fa1ef9430" title="min_label"><img src="https://discuss.pytorch.org/uploads/default/original/3X/3/b/3bcb663b85c446f7d5f583d3f5ec573fa1ef9430.png" alt="min_label" data-base62-sha1="8wXX0TiIWjwQxQ6WXJcK5HzArII" width="690" height="282" data-dominant-color="F5F5F5"><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">min_label</span><span class="informations">841×344 16 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><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/6/b6613008a62889719d454b0b864d9fb8ab85f95f.png" data-download-href="https://discuss.pytorch.org/uploads/default/b6613008a62889719d454b0b864d9fb8ab85f95f" title="np_unique"><img src="https://discuss.pytorch.org/uploads/default/original/3X/b/6/b6613008a62889719d454b0b864d9fb8ab85f95f.png" alt="np_unique" data-base62-sha1="q1p7DraCAcqdAnVDYPNOWpG30ir" width="690" height="259" data-dominant-color="F9F9F9"><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">np_unique</span><span class="informations">803×302 4.32 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>