Torch.utils make_grid with cmaps?

vision
1

I wanted to combine two grids from make_grid. One for the source images, and another from model predictions.

Is it possible to apply a cmap to the masks?

I pasted a few relevant parts of the code‹ below:

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()

image
image916×506 65.6 KB

1 Like
2

I think I found a decent solution to this. If anyone has another approach please share!

    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()

image
image900×490 69 KB

2 Likes
3

Thanks for sharing this approach, as it looks like an elegant way to map the colormap to the array. :slight_smile:

1 Like
4

Hope someone in the future finds it helpful!

1 Like