Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- import sys
- import os.path
- import glob
- import cv2
- import numpy as np
- import torch
- import architecture as arch
- model_path = sys.argv[1] # models/RRDB_ESRGAN_x4.pth OR models/RRDB_PSNR_x4.pth
- device = torch.device('cuda') # if you want to run on CPU, change 'cuda' -> cpu
- # device = torch.device('cpu')
- test_img_folder = 'LR_V/*'
- model = arch.RRDB_Net(3, 3, 64, 23, gc=32, upscale=4, norm_type=None, act_type='leakyrelu', \
- mode='CNA', res_scale=1, upsample_mode='upconv')
- model.load_state_dict(torch.load(model_path), strict=True)
- model.eval()
- for k, v in model.named_parameters():
- v.requires_grad = False
- model = model.to(device)
- print('Model path {:s}. \nTesting...'.format(model_path))
- idx = 0
- for path in glob.glob(test_img_folder):
- idx += 1
- base = os.path.splitext(os.path.basename(path))[0]
- print(idx, base)
- cap = cv2.VideoCapture(path)
- width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) * 4
- height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) * 4
- fps = int(cap.get(cv2.CAP_PROP_FPS))
- out = cv2.VideoWriter('results/{:s}.avi'.format(base), cv2.VideoWriter_fourcc('M','J','P','G'), fps, (width,height))
- while True:
- ret, img = cap.read();
- if img is None:
- break
- img = img * 1.0 / 255
- img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float()
- img_LR = img.unsqueeze(0)
- img_LR = img_LR.to(device)
- output = model(img_LR).data.squeeze().float().cpu().clamp_(0, 1).numpy()
- output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))
- output = (output * 255.0).round()
- output = np.uint8(output)
- out.write(output)
- cap.release()
- out.release()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement