UNet Predicts Masks in a Distorted Way

dev1: <p>Hi everyone. I’m experimenting with UNet using <a href="https://www.kaggle.com/c/carvana-image-masking-challenge" rel="noopener nofollow ugc">Carvana Image Masking Dataset</a>. I trained a model and I’m doing some inference with it. However predicted mask is very whitish and distorted. Here are the results:</p>
<p><strong>The Actual Image:</strong><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/2/72716e9d28dce70eb6b909bead2ec8668cce3f93.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/72716e9d28dce70eb6b909bead2ec8668cce3f93" title="The Actual Image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/2/72716e9d28dce70eb6b909bead2ec8668cce3f93_2_516x500.jpeg" alt="The Actual Image" data-base62-sha1="gkpsrrQCvBXHVp8GrABHay76Yrp" width="516" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/2/72716e9d28dce70eb6b909bead2ec8668cce3f93_2_516x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/7/2/72716e9d28dce70eb6b909bead2ec8668cce3f93_2_774x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/7/2/72716e9d28dce70eb6b909bead2ec8668cce3f93_2_1032x1000.jpeg 2x" data-dominant-color="C7C7C6"><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">The Actual Image</span><span class="informations">1432×1386 63 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><strong>Original Mask</strong><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/6/a6d4b8162bb6be8c7e2c17ecb115aef783565751.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/a6d4b8162bb6be8c7e2c17ecb115aef783565751" title="Screenshot 2023-07-20 at 23.13.24"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/6/a6d4b8162bb6be8c7e2c17ecb115aef783565751_2_519x500.jpeg" alt="Screenshot 2023-07-20 at 23.13.24" data-base62-sha1="nNR0xcZH1oUKp1DFsBvV2Xw9d7j" width="519" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/6/a6d4b8162bb6be8c7e2c17ecb115aef783565751_2_519x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/a/6/a6d4b8162bb6be8c7e2c17ecb115aef783565751_2_778x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/a/6/a6d4b8162bb6be8c7e2c17ecb115aef783565751_2_1038x1000.jpeg 2x" data-dominant-color="505050"><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">Screenshot 2023-07-20 at 23.13.24</span><span class="informations">1420×1366 53.8 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><strong>Predicted Mask</strong><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/c/dc198bb1e11075b6ae4637ab9c634f91a7e3cb45.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/dc198bb1e11075b6ae4637ab9c634f91a7e3cb45" title="Predicted Mask"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/c/dc198bb1e11075b6ae4637ab9c634f91a7e3cb45_2_513x500.jpeg" alt="Predicted Mask" data-base62-sha1="vp5M2BJ3SPek1fsfwnLDpY8yGJ7" width="513" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/c/dc198bb1e11075b6ae4637ab9c634f91a7e3cb45_2_513x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/d/c/dc198bb1e11075b6ae4637ab9c634f91a7e3cb45_2_769x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/d/c/dc198bb1e11075b6ae4637ab9c634f91a7e3cb45_2_1026x1000.jpeg 2x" data-dominant-color="E6E6E6"><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">Predicted Mask</span><span class="informations">1400×1364 28.5 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>As you can see the supposed to be mask forms a whitish cluster at the “Predicted Mask” image. Still very distant from what it supposed to be.</p>
<p><strong>Here is my dataset</strong></p>
<pre><code class="lang-auto">import os
from PIL import Image
from torch.utils.data.dataset import Dataset
from torchvision import transforms

class CarvanaDataset(Dataset):
    def __init__(self, root_path):
        self.root_path = root_path
        self.images = sorted([root_path+"/train/"+i for i in os.listdir(root_path+"/train/")])
        self.masks = sorted([root_path+"/train_masks/"+i for i in os.listdir(root_path+"/train_masks/")])
        self.transforms = transforms.Compose([
            transforms.Resize((512, 512)),
            transforms.ToTensor()])

    def __getitem__(self, index):
        img = Image.open(self.images[index])
        mask = Image.open(self.masks[index])

        return self.transforms(img), self.transforms(mask)

    def __len__(self):
        return len(self.images)
</code></pre>
<p><strong>Here is my training loop</strong></p>
<pre><code class="lang-auto">    model = UNet(3, 1).to(device)
    optimizer = optim.AdamW(model.parameters(), lr=LEARNING_RATE)
    criterion = nn.BCEWithLogitsLoss()

    for epoch in tqdm(range(EPOCHS)):
        model.train()
        train_running_loss = 0 
        for idx, img_mask in enumerate(tqdm(train_dataloader)):
            img = img_mask[0].float().to(device)
            mask = img_mask[1].float().to(device)

            y_pred = model(img)
            optimizer.zero_grad()

            loss = criterion(y_pred, mask)
            train_running_loss += loss.item()

            loss.backward()
            optimizer.step()

        train_loss = train_running_loss / idx+1

        model.eval()
        val_running_loss = 0
        with torch.no_grad():
            for idx, img_mask in enumerate(tqdm(val_dataloader)):
                img = img_mask[0].float().to(device)
                mask = img_mask[1].float().to(device)

                y_pred = model(img)
                loss = criterion(y_pred, mask)

                val_running_loss += loss.item()

            val_loss = val_running_loss / idx+1

        print("-"*30)
        print(f"Train Loss EPOCH {epoch+1}: {train_loss:.4f}")
        print(f"Valid Loss EPOCH {epoch+1}: {val_loss:.4f}")
        print("-"*30)

    torch.save(model.state_dict(), MODEL_SAVE_PATH)
</code></pre>
<p><strong>This is how I do inference</strong></p>
<pre><code class="lang-auto">trained_model = UNet(3, 1)
trained_model.load_state_dict(torch.load("./models/unet.pth", map_location=torch.device(device)))
trained_model.eval()

img = train_dataset[1][0].float()
img = img.unsqueeze(0)

pred = trained_model(img)
print(pred.size())

mask = pred.squeeze(0).cpu().detach()
mask = mask.permute(1, 2, 0)

plt.imshow(mask, cmap="gray")
plt.show()
</code></pre>
<p>I thought there may be something wrong about how I do the image transformation. Specifically with how I transform images with <code>Compose([transforms.Resize((512, 512)), transforms.ToTensor()])</code>.</p>
<p>But I’ve seen other sources do the same. I can’t spot what is wrong. I would appreciate some help.</p>

dhruvbird: <p>You’re plotting the raw logits that your network returns. When running inference you need to run them through sigmoid or alternatively just pick the logit with the highest value and return its index (which I’m assuming to be 0 or 1).</p>