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

image916×506 65.6 KB