Torch.utils make_grid with cmaps?

aksg87: <p>I wanted to combine two grids from make_grid. One for the source images, and another from model predictions.</p>
<p>Is it possible to apply a cmap to the masks?</p>
<p>I pasted a few relevant parts of the code‹ below:</p>
<pre><code class="lang-auto">from torchvision.utils import make_grid
...

def display_volumes(
    img_vol, pred_vol,
):

    def show(img, label=None, alpha=0.5):
        npimg = img.numpy()
        plt.imshow(
            np.transpose(npimg, (1, 2, 0)), interpolation="none"
        )
        if label is not None:
            lbimg = label.numpy()
            plt.imshow(
                np.transpose(lbimg, (1, 2, 0)),
                cmap="jet",  # cmap doesn't appear to do anything!
                alpha=alpha,
                interpolation="none",
            )

    x = torch.from_numpy(img_vol)
    y = torch.from_numpy(pred_vol)
  
    show(make_grid(x), make_grid(y), alpha=0.5)
    plt.show()
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/1/91fd5b44cf52211091c186fbe3d2ca2a0e5d639e.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/91fd5b44cf52211091c186fbe3d2ca2a0e5d639e" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/9/1/91fd5b44cf52211091c186fbe3d2ca2a0e5d639e_2_517x285.jpeg" alt="image" data-base62-sha1="kPu2US32XLZ9M8Z6SZv6Xlat2jA" width="517" height="285" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/9/1/91fd5b44cf52211091c186fbe3d2ca2a0e5d639e_2_517x285.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/9/1/91fd5b44cf52211091c186fbe3d2ca2a0e5d639e_2_775x427.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/9/1/91fd5b44cf52211091c186fbe3d2ca2a0e5d639e.jpeg 2x" data-dominant-color="424242"><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">image</span><span class="informations">916×506 65.6 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>

aksg87: <p>I think I found a decent solution to this. If anyone has another approach please share!</p>
<pre><code class="lang-auto">    from matplotlib import cm
    ...
    x = torch.from_numpy(img_vol)
    y = torch.from_numpy(pred_vol)
    cmap_vol = np.apply_along_axis(cm.viridis, 0, y.numpy()) # converts prediction to cmap!
    cmap_vol = torch.from_numpy(np.squeeze(cmap_vol))

    show(make_grid(x), make_grid(y), alpha=0.5)
    plt.show()

</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/8/58be35e625a0c21ebc0be5d48eb79d7a5a461847.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/58be35e625a0c21ebc0be5d48eb79d7a5a461847" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/5/8/58be35e625a0c21ebc0be5d48eb79d7a5a461847_2_690x375.jpeg" alt="image" data-base62-sha1="cF3xzWskYrafacO8cX8CHKt0XPN" width="690" height="375" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/5/8/58be35e625a0c21ebc0be5d48eb79d7a5a461847_2_690x375.jpeg, https://discuss.pytorch.org/uploads/default/original/3X/5/8/58be35e625a0c21ebc0be5d48eb79d7a5a461847.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/5/8/58be35e625a0c21ebc0be5d48eb79d7a5a461847.jpeg 2x" data-dominant-color="573E58"><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">image</span><span class="informations">900×490 69 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>

ptrblck: <p>Thanks for sharing this approach, as it looks like an elegant way to map the colormap to the array. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

aksg87: <p>Hope someone in the future finds it helpful!</p>