Search code examples
pytorchpython-imaging-libraryfast-ai

Pytorch tensor to PIL image with alpha channel

Using fastai v1, I have a model that transforms an image. When I plot the resulting image with matplotlib, the background is white;

ax.imshow(image2np(img.data), alpha=None)
plt.show()

matplot lib output with white background

Without alpha=None, the background is black.

Thing is, I need to output a PIL image. But when doing so, the background becomes black.

from torchvision import transforms
im = transforms.ToPILImage()(img.data).convert('RGBA')

enter image description here

I guess I have to add an alpha channel but I'm not sure how.

This is a print of the tensor;

tensor([[[1.0904, 1.1132, 1.1009,  ..., 1.1189, 1.0949, 1.0662],
         [1.1277, 1.1046, 1.1044,  ..., 1.1015, 1.1039, 1.0865],
         [1.1359, 1.1245, 1.1096,  ..., 1.1164, 1.1179, 1.1051],
         ...,
         [1.0870, 1.0913, 1.0904,  ..., 1.0954, 1.1055, 1.0953],
         [1.0706, 1.0862, 1.0896,  ..., 1.0978, 1.0918, 1.0478],
         [0.9694, 1.0501, 1.0482,  ..., 1.0810, 1.0433, 0.9563]],

        [[1.0870, 1.1245, 1.1248,  ..., 1.1244, 1.1221, 1.0880],
         [1.1188, 1.1277, 1.1278,  ..., 1.1276, 1.1277, 1.1197],
         [1.1196, 1.1276, 1.1273,  ..., 1.1274, 1.1277, 1.1247],
         ...,
         [1.0942, 1.1244, 1.1238,  ..., 1.1266, 1.1265, 1.1188],
         [1.0687, 1.1223, 1.1236,  ..., 1.1274, 1.1269, 1.1173],
         [0.9189, 1.0851, 1.0955,  ..., 1.1191, 1.1116, 1.0639]],

        [[1.0562, 1.0785, 1.0787,  ..., 1.0785, 1.0754, 1.0554],
         [1.0758, 1.0806, 1.0807,  ..., 1.0807, 1.0806, 1.0755],
         [1.0758, 1.0806, 1.0803,  ..., 1.0804, 1.0807, 1.0780],
         ...,
         [1.0583, 1.0773, 1.0772,  ..., 1.0795, 1.0800, 1.0728],
         [1.0375, 1.0735, 1.0752,  ..., 1.0798, 1.0791, 1.0688],
         [0.9594, 1.0424, 1.0479,  ..., 1.0686, 1.0633, 1.0203]]])

Output of tensor.shape:

torch.Size([3, 920, 920])
Solution

  • I can't remember the code as its been a while, but it was something simple you could find in the docs.

    Now for the black background: the real problem came from the PIL display function having an issue with showing the alpha channel. Saving to disk confirmed that the image was created properly.