TypeError: Invalid shape (360, 480, 4) for image data

username: gb_pytorch
cooked: <p>I am having trouble visualizing the predictions of a multi-class (4 classes) instance segmentation.</p>
<p>I can visualize the masks before the modeling like this but how could I do the same with the predictions?</p>
<pre><code class="lang-auto">CLASSES = ['sky',  'building', 'pole', 'road']
</code></pre>
<pre><code class="lang-auto">dataset = Dataset(x_train_dir, y_train_dir, classes= ['sky', 'building', 'pole', 'road'])
</code></pre>
<pre><code class="lang-auto">image, mask = dataset[4] # get some sample
visualize(
    image=image, 
    #cars_mask=mask.squeeze(),
    sky_mask = mask[...,0].squeeze(),
    cars_mask = mask[...,1] .squeeze(),
    pole_mask = mask[...,2].squeeze(),
    road_mask = mask[...,3] .squeeze(),
)
</code></pre>
<p>Example output:<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/f/4/f48a232577342db5e0b3492f172c52010254eea1.png" alt="Screen Shot 2020-03-06 at 4.28.48 PM" data-base62-sha1="yTisyqzObcF5vMvVcPElpddIC1X" width="688" height="166"></p>
<p>Visualizing predictions</p>
<pre><code class="lang-auto">for i in range(5):
    n = np.random.choice(len(test_dataset))
    
    image_vis = test_dataset_vis[n][0].astype('uint8')
    image, gt_mask = test_dataset[n]
    
    gt_mask = gt_mask.squeeze()
    
    x_tensor = torch.from_numpy(image).to(DEVICE).unsqueeze(0)
    pr_mask = model.predict(x_tensor)
    pr_mask = (pr_mask.squeeze().cpu().numpy().round())
        
    visualize(
        image=image_vis, 
        ground_truth_mask=gt_mask, 
        predicted_mask=pr_mask
    )
</code></pre>

username: Oli
cooked: <p>You could create a <a href="https://pillow.readthedocs.io/en/stable/handbook/concepts.html#palette" rel="nofollow noopener">palette image</a>.</p>
<p>I think you can argmax predictions -&gt; numpy() -&gt; PIL Image.fromarray() -&gt; set the image to pallette mode <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: gb_pytorch
cooked: <p>The following code seems to be working  but ground truth and prediction mask images are condensed <div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/e/6e0e203011f322a2fd7394efd78c7700f5809003.png" data-download-href="https://discuss.pytorch.org/uploads/default/6e0e203011f322a2fd7394efd78c7700f5809003" title="Screen Shot 2020-03-06 at 4.26.05 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/e/6e0e203011f322a2fd7394efd78c7700f5809003_2_690x146.png" alt="Screen Shot 2020-03-06 at 4.26.05 PM" data-base62-sha1="fHAMJmkQzFBAJSrOSgqeH9Ejpl1" width="690" height="146" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/e/6e0e203011f322a2fd7394efd78c7700f5809003_2_690x146.png, https://discuss.pytorch.org/uploads/default/optimized/3X/6/e/6e0e203011f322a2fd7394efd78c7700f5809003_2_1035x219.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/6/e/6e0e203011f322a2fd7394efd78c7700f5809003_2_1380x292.png 2x" data-dominant-color="EDEDEE"><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">Screen Shot 2020-03-06 at 4.26.05 PM</span><span class="informations">1684×358 102 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> . So weird.</p>
<pre><code class="lang-auto">for i in range(2):
    n = np.random.choice(len(test_dataset))
    image_vis = test_dataset_vis[n][0].astype('uint8')
    
    image, gt_mask = test_dataset[n]
    
    gt_mask1 = gt_mask[...,0].squeeze()
    gt_mask2 = gt_mask[...,1].squeeze()
    gt_mask3 = gt_mask[...,2].squeeze()
    gt_mask4 = gt_mask[...,3].squeeze()

    x_tensor = torch.from_numpy(image).to(DEVICE).unsqueeze(0)
    pr_mask = model.predict(x_tensor)
    
    pr_mask1 = (pr_mask[...,0].squeeze().cpu().numpy().round())
    pr_mask2 = (pr_mask[...,1].squeeze().cpu().numpy().round())
    pr_mask3 = (pr_mask[...,2].squeeze().cpu().numpy().round())
    pr_mask4 = (pr_mask[...,3].squeeze().cpu().numpy().round())

        
    visualize(
        image=image_vis,
        gt_mask1 = gt_mask1,
        gt_mask2 = gt_mask2,
        gt_mask3 = gt_mask3,
        gt_mask4 = gt_mask4,
        pr_mask1 = pr_mask1,
        pr_mask2 = pr_mask2,
        pr_mask3 = pr_mask3,
        pr_mask4 = pr_mask4
    )
</code></pre>
<pre><code class="lang-auto"># helper function for data visualization
def visualize(**images):
    """PLot images in one row."""
    n = len(images)
    plt.figure(figsize=(16, 5))
    for i, (name, image) in enumerate(images.items()):
        plt.subplot(1, n, i + 1)
        plt.xticks([])
        plt.yticks([])
        plt.title(' '.join(name.split('_')).title())
        plt.imshow(image)
    plt.show()
</code></pre>

username: copperwiring
cooked: <p>I had a similar problem a while ago. The parts that are condensed are actually because of some incorrect dimension/ coordinate when you are plotting.</p>
<p>You will need to recheck which channel should have which coordinate.</p>