How to create a palette-based PNG with alpha channel?

Question

I have a 2D NumPy array with values from 0 to 6 (corresponding to image segmentation classes), and I would like to turn it into a color-indexed PNG image using Image.putpalette() from Pillow (using version 8.0.1). The color palette should contain an alpha channel. According to the official documentation this should be possible:

The palette sequence must contain either 768 integer values, or 1024 integer values if alpha is included. Each group of values represents the red, green, blue (and alpha if included) values for the corresponding pixel index. Instead of an integer sequence, you can use an 8-bit string.

So far I've succeeded only without the alpha channel:

from PIL import Image
import matplotlib.cm as cm
import numpy as np

# test data
np.random.seed(42)
mask = np.random.randint(low=0, high=6, size=(32, 32))

# color palette - random colors for testing
colormap = np.zeros((256, 3))
colormap[:7, :] = cm.jet(np.linspace(0, 1, 7))[:, :3]

img = Image.fromarray(mask.astype(np.uint8))
img = img.convert("P")
img.putpalette((colormap * 255).astype(np.uint8).flatten())

The above works nicely.

However, if I try to include the alpha channel (which, as far as I understand, I have to append to the end), I get nowhere:

img = Image.fromarray(corrected_mask.astype(np.uint8))
img = img.convert("P")
img.putpalette(np.append((colormap[:, :] * 255).astype(np.uint8).flatten(), np.ones(256, dtype=np.uint8) * 255))

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-65-6a955d432bbc> in <module>
      1 img = Image.fromarray(corrected_mask.astype(np.uint8))
      2 img = img.convert("P")
----> 3 img.putpalette(np.append((colormap[:, :] * 255).astype(np.uint8).flatten(), np.ones(256, dtype=np.uint8) * 255))

~/miniconda3/envs/tf2_py38/lib/python3.8/site-packages/PIL/Image.py in putpalette(self, data, rawmode)
   1697         self.palette = palette
   1698         self.palette.mode = "RGB"
-> 1699         self.load()  # install new palette
   1700 
   1701     def putpixel(self, xy, value):

~/miniconda3/envs/tf2_py38/lib/python3.8/site-packages/PIL/Image.py in load(self)
    816         if self.im and self.palette and self.palette.dirty:
    817             # realize palette
--> 818             self.im.putpalette(*self.palette.getdata())
    819             self.palette.dirty = 0
    820             self.palette.mode = "RGB"

ValueError: invalid palette size

Attempting to set the mode to PA also leads nowhere...

Can somebody point me in the right direction? Thanks!

Answer 1

You need to explicitly set rawmode='RGBA' in the putpalette call, since it defaults to 'RGB':

from PIL import Image
import matplotlib.cm as cm
import numpy as np

# test data
np.random.seed(42)
mask = np.random.randint(low=0, high=6, size=(32, 32))

# color palette - random colors for testing
colormap = np.zeros((256, 4))
colormap[:7, :3] = cm.jet(np.linspace(0, 1, 7))[:, :3]
colormap[:7, 3] = np.linspace(0, 1, 7)

img = Image.fromarray(mask.astype(np.uint8))
img = img.convert('P')
img.putpalette((colormap * 255).astype(np.uint8).flatten(), rawmode='RGBA')

img.save('image.png')

That gives an image like this, which has specific alpha values for each color:

----------------------------------------
System information
----------------------------------------
Platform:      Windows-10-10.0.16299-SP0
Python:        3.9.1
Matplotlib:    3.3.4
NumPy:         1.20.1
Pillow:        8.1.0
----------------------------------------