Test a model trained with images with a video

Dear senior programmers,
I have built and trained a deep learning model for single image dehazing using images. Subsequently, I would like to test the trained model on some shorts videos. I have been able to first convert the video to images, then the obtained images were processed using the trained model. Finally, I converted the processed images into a video. However, performing the task this way is time consuming given that it includes many steps. Hence, I have managed to write the following funstion for testing the model with the video without having to convert it to images. The model is as follows.

def dehaze_image(image_path):
  
  #data_hazy = Image.open(image_path)
  data_hazy = image_path
  data_hazy = (np.asarray(data_hazy)/255.0)
   
  data_hazy = torch.from_numpy(data_hazy).float()
  
  data_hazy = data_hazy.permute(2,0,1)
  data_hazy = data_hazy.cuda().unsqueeze(0)
  
  dehaze_net = net.dehaze_net().cuda()
  dehaze_net.load_state_dict(torch.load('snapshots4/dehazer.pth'))
  
  clean_image = dehaze_net(data_hazy)
  #torchvision.utils.save_image(torch.cat((data_hazy, clean_image),0), "resultsM22/" + image_path.split("/")[-1])
  #torchvision.utils.save_image(clean_image, "results3/" + image_path.split("/")[-1])
  return clean_image

if __name__ == '__main__':

  start = timeit.default_timer()
  

  cap=cv2.VideoCapture('/home/HDD2/14.mp4')
  fourcc = cv2.VideoWriter_fourcc(*'XVID')
  frame_width=640
  frame_height=480
  save = cv2.VideoWriter('/home/HDD2/output14.mp4',fourcc, 20.0, (frame_width,frame_height))
    
    
while(cap.isOpened()):
    ret, frame=cap.read()
    if(ret== True):
            
        frame=cv2.resize(frame, (frame_width, frame_height))
        out = dehaze_image(frame)
        print(out.shape)
        cv2.imshow("Before",frame) 
        cv2.imshow("After", np.array(out.cuda().detach().cpu()))
        out = cv2.flip(out,0)
        save.write(out)
        key=cv2.waitKey(1)
        if key==27:
            break
    else:
        break

cap.release()
save.release()
cv2.destroyAllWindows()

stop = timeit.default_timer()
  
print('Time: ', stop - start)  

This function can run. However, the reconstructed video can not open. It displays file unsupported or corrupted. Please, could you help me sort this out?

Thank you for your time and patience

Could you try to remove the PyTorch code and just create the video output using random numpy arrays with the same shape and dtype as in your current code?
I think this issue might be more related to OpenCV than PyTorch and you might need to e.g. permute the arrays before storing them.

Thank you for your reply sir. When I remove the Pytorch and just create video output, It works well. However, when I process the frame with the trained pytorch model, I can not open the saved video. I have corsschecked it several time. Please, any other suggestions?

Could you post an executable code snippet, which reproduces this error, using random input tensors, please?

Hello sir,
I am sorry for the late reply. I have been trying to check whether the problem might be related to OpenCV. Unfortunately, I still could not solve it. The code does run. However, the saved video after processing using the Pytorch model is empty. I still can not figure out what might be the cause.