Problem in plotting the masks

Adwaitt: <p>Hello everyone <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<p>I am working on the segmentation problem and images are grayscale but I stacked them along the depth which made each channel to be a 4 channel image and I also did the one hot encoding for my mask there were 4 labels in total one of them was the background (I also want to know how to exclude that altogether). My average IoU score was 0.86 keeping the threshold to be 0.9 but when I finally want to generate the mask this is the output I am getting.</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/9/d/9d678351cb56abe9f48ea9c4f9b226ae1a72e750.png" alt="image" data-base62-sha1="mssJMi2dJFCqmWaPY49xiCDSPXW" width="163" height="269"></p>
<p>Any idea what might be wrong?</p>
<p>Here’s my code:</p>
<pre><code class="lang-auto">%matplotlib inline
import os
import warnings
warnings.filterwarnings("ignore", category=RuntimeWarning)
warnings.filterwarnings("ignore", category=UserWarning)
from IPython.display import FileLink

import random

import pandas as pd
import numpy as np

import nibabel as nib

import nibabel as nib
import tarfile

import tensorflow as tf

import torch
from torch.nn.functional import one_hot
from torch import nn, optim
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import Dataset, DataLoader

import segmentation_models_pytorch as smp

import albumentations as A
from albumentations.pytorch import ToTensorV2

import cv2

from tqdm.auto import tqdm

import matplotlib.pyplot as plt
import seaborn as sns

if torch.cuda.is_available():
    device = "cuda:0"
    torch.cuda.empty_cache()
else:
    device = "cpu"

SEED = 42
    
np.random.seed(SEED)
torch.manual_seed(SEED)

torch.cuda.empty_cache()
device = torch.device(device)

class datagen(Dataset):
    def __init__(self, IDs, tfms):
        super(datagen, self).__init__()
        self.IDs = IDs
        self.tfms = tfms

    def __len__(self):
        return len(self.IDs)
    
    def __getitem__(self, idx):
        ID = self.IDs[idx]
        
        x1 = nib.load(f'./data/{ID}/{ID}_flair.nii.gz').get_fdata()
        x2 = nib.load(f'./data/{ID}/{ID}_t1.nii.gz').get_fdata()
        x3 = nib.load(f'./data/{ID}/{ID}_t1ce.nii.gz').get_fdata()
        x4 = nib.load(f'./data/{ID}/{ID}_t2.nii.gz').get_fdata()
        
        y = nib.load(f'./data/{ID}/{ID}_seg.nii.gz').get_fdata()
        y[y == 4] = 3
        
        images = []
        masks = []
        
        for i in range(len(x1)):
            image = np.dstack((x1[i], x2[i], x3[i], x4[i]))
            if np.count_nonzero(image) == 0:
                image = tfms(image = image)['image']
            else:
                image = (image - np.min(image))/(np.max(image) - np.min(image))
                image = tfms(image = image)['image']
            images.append(image)
            
            mask = y[i]
            mask = tf.keras.utils.to_categorical(mask, num_classes = 4)
            mask = tfms(image = mask)['image']
            masks.append(mask)
        
        return (images, masks)

class Network(nn.Module):
    def __init__(self):
        super(Network, self).__init__()
        self.model = smp.Unet(encoder_name = 'efficientnet-b7',
                              encoder_weights = 'imagenet',
                              in_channels = 4,
                              classes = 4,
                              activation = 'sigmoid')
        for param in self.model.encoder.parameters():
            param.requires_grad = False
    def forward(self, x):
        x = self.model(x)
        return x

ids = os.listdir('./data')
random.shuffle(ids)
train_ids = ids[: 1002]
val_ids = ids[1002: ]

tfms = A.PadIfNeeded(min_height = 256, min_width = 160, always_apply = True)

train_set = datagen(train_ids, tfms)
train = DataLoader(train_set, batch_size = 1, shuffle = True)

val_set = datagen(val_ids, tfms)
val = DataLoader(val_set, batch_size = 1, shuffle = True)

torch.cuda.empty_cache()
model = Network().to(device)
criterion = smp.utils.losses.DiceLoss()
metrics = [smp.utils.metrics.IoU(threshold=0.95)]
optimizer = torch.optim.Adam(model.parameters(), lr = 0.005)
epochs = 1

for e in range(epochs):
    train_losses = []
    train_iou = []
    val_losses = []
    val_iou = []
    for image, mask in tqdm(train):
        for i in range(len(image)):
            optimizer.zero_grad()
            x, y = torch.tensor(image[i], dtype = torch.float32).reshape(1, 4, 256, 160).to(device), torch.tensor(mask[i], dtype = torch.float32).reshape(1, 4, 256, 160).to(device)
            y_pred = model.forward(x)
            loss = criterion(y_pred, y)
            train_losses.append(loss.item())
            iou = metrics[0].forward(y_pred, y)
            train_iou.append(iou)
            loss.backward()
            optimizer.step()
            print(f'Loss: {loss.item():.4f}, IoU: {iou:.4f}', end = '\r')
    with torch.no_grad():
        for image, mask in val:
            for i in range(len(image)):
                x, y = torch.tensor(image[i], dtype = torch.float32).reshape(1, 4, 256, 160).to(device), torch.tensor(mask[i], dtype = torch.float32).reshape(1, 4, 256, 160).to(device)
                y_pred = model.forward(x)
                loss = criterion(y_pred, y)
                val_losses.append(loss.item())
                iou = metrics[0].forward(y_pred, y)
                val_iou.append(iou)
                print(f'Loss: {loss.item():.4f}, IoU: {iou:.4f}', end = '\r')
                
    print(f'Epochs: {e + 1}/{3}, Training Loss: {sum(train_losses)/len(train_losses)}, Training IoU: {sum(train_iou)/len(iou)}, Validation Loss: {sum(val_losses)/len(val_losses)}, Validation IoU: {sum(val_iou)/len(val_iou)}')
    torch.save(model.state_dict, f'./model_{1}.pth')

with torch.no_grad():
    y_pred = model.forward(torch.tensor(images[135].reshape(1, 4, 256, 160), dtype = torch.float32).cuda())
    y = masks[135].reshape(1, 4, 256, 160)
    b = criterion(y_pred, y.cuda())
</code></pre>
<p>Here are the original masks:<br>
(background, not necessary)<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/e/b/eb5e5535b01fee32814cde0e140cf4b3e600778f.png" alt="image" data-base62-sha1="xAajoeoUeANTux8Y8FX8cK8BHCn" width="172" height="273"></p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/8/0/805abe68b45ba382347f268c9644168c85f24f2c.png" alt="image" data-base62-sha1="ijtx6eGmQvb02QX7ENI71eRIVhW" width="168" height="266"></p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/7/1/719a88e525f66c23b12426ad2fd5bc6508044187.png" alt="image" data-base62-sha1="gcZ2Gyk8SadtyeQ1ecDxizs7kgL" width="163" height="257"></p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/9/c/9cb49bf10fd9c0a513d178e4907a25d95e4d6e17.png" alt="image" data-base62-sha1="mmhrcndiKqKGC8SRaAa7EiqQylV" width="147" height="241"></p>
<p>The images were 240, 155 initially but I padded them to 256, 160 because it was required by the frameworok to have a image of the dimensions divisible by 32</p>

ptrblck: <p>I don’t know what memory layout <code>nib</code> uses, but I would guess it’s channels-last (same as OpenCV)?<br>
If so, you are interleaving the images most likely via:</p>
<pre><code class="lang-python">images[135].reshape(1, 4, 256, 160)
</code></pre>
<p>and should change it to a <code>tensor.permute</code> or <code>numpy_array.transpose</code> if you want to change the dimension order.</p>

Adwaitt: <p>Thank you Peter, this seems to work.</p>