How to save png images with OpenCV
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?

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(&quot;img1.png&quot;,img1_reshaped) #save the images as .png
cv2.imwrite(&quot;img2.png&quot;,img2_reshaped)

