[SOLVED] Turn Grayscale mask-image into RGB while changing one grayscale into specific color

Turn Grayscale mask-image into RGB while changing one grayscale into specific color

While using torch and matplotlib I was able to change the grayscale color #50 e.g. to red. However, only by saving the image inbetween. How can it be solved without saving the image? I am sorry for the simple question, but I am new with torch, torchvision and matplotlib.

Here is my code, first to save the image:

import torch
import matplotlib.pyplot as plt
from torchvision.io import read_image 

mask = read_image("A_mask.png")

mask_np = mask.numpy()
print(mask_np.shape)    #(1, 438, 567)
mask_np[mask_np == 4] = 50
output_image = torch.tensor(mask_np)

from torchvision import transforms
im = transforms.ToPILImage()(output_image).convert("RGB")
im.save("A-Mask-Colored-By-Hand.png")

Then to load it:

from torchvision import transforms

m = read_image("A-Mask-Colored-By-Hand.png")
#print(type(m))  # <class 'torch.Tensor'>

Here is my question, how to convert "im" into expected format, instead reading it. This code block is what I have tried and which is not working: So commenting out this block ... will work but is not optimal by saving and loading again.

from PIL import Image
#tensor_image = im.permute(1,2,0) # #
tensor_image = im
pil_image = Image.fromarray((tensor_image.numpy() * 255).astype("uint8")) 
# here it stops, throwing an AttributeError: 'Image' object has no attribute 'numpy'

This code block works:

m = pil_image
print(m.shape)    #(3, 438, 567)
m = m.numpy()
print(m.shape)    #(3, 438, 567)

cgrey = 50 # grey #50 > into red #255
mask = (m[0]==cgrey) & (m[1]==cgrey) & (m[2]==cgrey)
m[0][mask] = 255
m[1][mask] = 0
m[2][mask] = 0

output_image = torch.tensor(m)

from torchvision import transforms
im = transforms.ToPILImage()(output_image).convert("RGB")
print(im) # tensor
im.save("A-Mask-Colored-RGB.png")

Thanks in advance!

Solution

Meanwhile I resolved the issue with this function:

import numpy as np 
def pil_to_tensor(images): 
    images = np.array(images) 
    images = torch.from_numpy(images.transpose(2, 0, 1)) 
    return images

‍ ‍ ‍ ‍ ‍ ‍ And the two code blocks including saving the image can be deleted. Instead calling the function:

pil_image = pil_to_tensor(im)