How to save png images with OpenCV-CodePudding

I am trying to use OpenCV to convert the results of the model I trained into a png image. My output has 4 channels, and I am not sure how to convert these 4 channels to png.

# Load the model
model = CNNSEG()
model.load_state_dict(torch.load(PATH))
model.eval()

for iteration, sample in enumerate(test_data_loader):
    img = sample
    print(img.shape)

    plt.imshow(img[0,...].squeeze(), cmap='gray') #visualise all images in test set
    plt.pause(1)
    
    # output the results
    img_in = img.unsqueeze(1)
    output = model(img_in) # shape: [2, 4, 96, 96]

As here shows, the shape of output is [2, 4, 96, 96] which are batch size, channels, height and width. So how could I do to convert it to png image?

CodePudding user response：

To write the images you need to convert the output into the right format (assuming outputs are in the 0,1 range):

# Convert outputs from 0-1 to 0, 255
img_in *= 255.0

# Convert floats to bytes
img_in = img_in.astype(np.uint8)

# Transpose the images from channel first (4, 96, 96) to channel last (96, 96, 4)
image1 = img_in[0, :, :, :].transpose(2, 1, 0)
image2 = img_in[1, :, :, :].transpose(2, 1, 0)

Then it's just a matter of saving the images:

cv2.imwrite('./example_path/image1.png', image1)
cv2.imwrite('./example_path/image2.png', image2)

CodePudding user response：

You would want to split your image into two essentially and then save them individually.

import numpy as np
import cv2
img = np.ones((2,4,96,96),dtype=np.uint8) #creating a random image


img1 = img[0,:,:,:] #extracting the two separate images
img2 = img[1,:,:,:]

img1_reshaped = img1.transpose() #reshaping them to the desired form of (w,h,c)
img2_reshaped = img2.transpose()

cv2.imwrite("img1.png",img1_reshaped) #save the images as .png
cv2.imwrite("img2.png",img2_reshaped)