numpy - How to change cmap based on array value in python

Question

Current code:

        # Mask array with mask pixels set to 1 and rest set to 0
        mask = nib.load(os.path.join(mask_path, f))
        # Grey scale MRI image slice
        image = nib.load(os.path.join(images_path, f))
        # Extracting pixel arrays
        mask_data = mask.get_fdata()
        image_data = image.get_fdata()
        # Setting image pixels to 0 where mask pixels are set to 1
        masked_image = np.where(mask_data == 0, image_data,0)
        # Transposing and flipping to fix visual orientation
        masked_image = masked_image.transpose((1,0))
        masked_image = np.flip(masked_image,axis=0)
        # Want something like this
        # colors = np.where(mask_data == 1, 'autumn','gray')
        fig,ax = plt.subplots()
        ax.imshow(masked_image,cmap="gray")

Output image:

Currently my code takes in an image and a mask, I set the pixel values in the image_array to 0 based on the mask. I ideally want this to show up as red:

The only way I can see how to do this is a different colormap for all pixels with a value of 0. I don't know how to write a custom colour map based on pixel value. Is there any way to do this?

question from:https://stackoverflow.com/questions/65924186/how-to-change-cmap-based-on-array-value-in-python

Answer 1

Completely forgot matplotlib draws on top of pre existing data if using the same figure. The below code works to do what I want without needing to create my own colormap.

        plt.clf()
        # Mask array with mask pixels set to 1 and rest set to 0
        mask = nib.load(os.path.join(mask_path, f))
        # Grey scale MRI image slice
        image = nib.load(os.path.join(images_path, f))
        # Extracting pixel arrays
        mask_data = mask.get_fdata()
        image_data = image.get_fdata()
        # Transposing and flipping to fix visual orientation
        image_data = image_data.transpose((1,0))
        image_data = np.flip(image_data,axis=0)
        mask_data = mask_data.transpose((1,0))
        mask_data = np.flip(mask_data,axis=0)
        mask_data[mask_data==0] = None
        plt.imshow(image_data,cmap='gray')
        plt.imshow(mask_data,cmap = 'autumn')

Result: