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