import torch
import numpy as np
import cv2
import os
def enhance(img, net):
img = img / 255.
img = np.transpose(img, (2, 0, 1))
img = np.expand_dims(img, axis=0)
img = torch.from_numpy(img.astype(np.float32))
result = net(img)
result = np.squeeze(result.detach().numpy(), axis=0)
result = np.transpose(result, (1, 2, 0))
result = result * 255
result = result.astype(np.int)
result = result.astype(np.uint8)
return result
def main(video_path, model_path, target_file):
net = torch.load(model_path, map_location='cpu')
video = cv2.VideoCapture(video_path)
fps = video.get(cv2.CAP_PROP_FPS)
width, height = video.get(cv2.CAP_PROP_FRAME_WIDTH), video.get(cv2.CAP_PROP_FRAME_HEIGHT)
size = (int(width * 4), int(height * 4))
videoWrite = cv2.VideoWriter(target_file, cv2.VideoWriter_fourcc('m', 'p', '4', 'v'), fps, size)
success, frame = video.read()
cv2.imwrite('data/onepiece.jpg', frame)
exit()
while success:
print(frame.shape, frame.dtype)
result = enhance(frame, net)
print(result.shape, result.dtype)
videoWrite.write(result)
success, frame = video.read()
video.release()
if __name__ == '__main__':
video_file = 'data/onepiece_demo.mp4'
save_path = 'onepiece_demo.mp4'
pt_path = 'RealESRGAN_x4plus.pt'
main(video_file, pt_path, save_path)
