How can I display the predicted mask from UNet?

kwkim128: <p>Hi everyone. I am trying to plot the original image, mask and predicted mask from the UNet model, however, I am getting weird images as my output. This is a method I used before, which worked perfectly but for some weird reason, it is not working anymore…</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/4/24c94ae01b876daf75a93125d48bf1579445e46d.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/24c94ae01b876daf75a93125d48bf1579445e46d" title="Figure 2023-09-13 150823"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/4/24c94ae01b876daf75a93125d48bf1579445e46d_2_690x359.jpeg" alt="Figure 2023-09-13 150823" data-base62-sha1="5fqshd2BG6UrxdlMTaTTpqi6Ndj" width="690" height="359" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/4/24c94ae01b876daf75a93125d48bf1579445e46d_2_690x359.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/2/4/24c94ae01b876daf75a93125d48bf1579445e46d_2_1035x538.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/2/4/24c94ae01b876daf75a93125d48bf1579445e46d_2_1380x718.jpeg 2x" data-dominant-color="D2CCC8"><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">Figure 2023-09-13 150823</span><span class="informations">1920×999 165 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 shown in the image above, I am not getting an image at all…</p>
<p><strong>Here is how I am calling the dataset</strong>:</p>
<blockquote></blockquote>
<p>import os<br>
import numpy as np<br>
import torch<br>
from torch.utils.data import Dataset<br>
from PIL import Image<br>
from torchvision import transforms</p>
<p>class GetData(Dataset):<br>
def <strong>init</strong>(self, img_path, mask_path, transform=None):<br>
self.img_path = img_path<br>
self.mask_path = mask_path<br>
self.transform = transform</p>
<pre><code>def __len__(self):
    return len(self.img_path)

def __getitem__(self, idx):
    image = np.array(Image.open(self.img_path[idx]).convert('RGB'), 
                     dtype=np.float32)
    mask = np.array(Image.open(self.mask_path[idx]).convert('L'), 
                     dtype=np.float32)
    mask = ((mask/np.max([mask.max(), 1e-8])) &gt; 0.5).astype(np.float32)
    
    
    if self.transform is not None:
        augmentations = self.transform(image=image, mask=mask)
        image = augmentations['image']
        mask = augmentations['mask']

    
    return image, mask
</code></pre>
<p><strong>And here is the testing loop which prints to the images</strong>:</p>
<blockquote>
<p>def testinge_loop (model, loader, device=torch.device(‘cuda’)):</p>
</blockquote>
<pre><code>dice = []
model.eval()
with torch.no_grad():
    for x, y in loader:
        
        # split data to image and mask
        image = x
        mask = y
        
        image = image.to(device)
        mask = mask.to(device)
        
        test_outputs = model(image)
        test_outputs = torch.sigmoid(test_outputs)

        dice_metric = dice_metrics(test_outputs, mask)
        dice.append(dice_metric.cpu().numpy())
        
        
        plt.figure(figsize=(12,6), dpi=1200)
        plt.subplot(1,3,1)
        plt.title('Original Image')
        plt.plot(image[0].permute(1,2,0).detach().cpu()[:,:,0])
        plt.subplot(1,3,2)
        plt.title('Ground Truth')
        plt.plot(mask.permute(1,2,0).detach().cpu()[:,:,0])
        plt.subplot(1,3,3)
        plt.title('Predicted Mask')
        plt.plot(test_outputs[0].permute(1,2,0).detach().cpu()[:,:,0])
        break
        
metric = sum(dice)/len(loader)
return metric 
</code></pre>
<p>I genuinely have no idea what I am doing wrong…I tried to copy what others did but that seems to cause more errors… Thank you in advance.</p>

ptrblck: <p>Try to use <code>plt.imshow</code> to display images instead of <code>plt.plot</code> as the latter will interpret the array as rows of different signals.</p>