| import os | |
| import cv2 | |
| import time | |
| import torch | |
| import argparse | |
| import numpy as np | |
| from tqdm import tqdm | |
| import common | |
| import imgproc | |
| import onnxruntime as ort | |
| torch.manual_seed(1) | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--model", type=str, default="edsr_baseline_x2_1.onnx", help="onnx model path") | |
| parser.add_argument('--scale', nargs='+', type=int, default=[2], help='super resolution scale') | |
| parser.add_argument("--dir_demo", type=str, default='../video/test_1920x1080.mp4', help="demo image directory") | |
| parser.add_argument('--rgb_range', type=int, default=255, help='maximum value of RGB') | |
| parser.add_argument('--arch', type=str, default='espcn', help='model architecture (options: edsr、espcn)') | |
| def quantize(img, rgb_range): | |
| pixel_range = 255 / rgb_range | |
| return np.round(np.clip(img * pixel_range, 0, 255)) / pixel_range | |
| def from_numpy(x): | |
| return x if isinstance(x, np.ndarray) else np.array(x) | |
| class VideoTester(): | |
| def __init__(self, scale, my_model, dir_demo, rgb_range=255, cuda=True, arch='espcn'): | |
| self.scale = scale | |
| self.rgb_range = rgb_range | |
| self.providers = ['CUDAExecutionProvider'] if cuda else ['CPUExecutionProvider'] | |
| self.session = ort.InferenceSession(my_model, providers=self.providers) | |
| self.output_names = [x.name for x in self.session.get_outputs()] | |
| self.input_name = self.session.get_inputs()[0].name | |
| self.dir_demo = dir_demo | |
| self.filename, _ = os.path.splitext(os.path.basename(dir_demo)) | |
| self.arch = arch | |
| def test(self): | |
| torch.set_grad_enabled(False) | |
| if not os.path.exists('experiment'): | |
| os.makedirs('experiment') | |
| for idx_scale, scale in enumerate(self.scale): | |
| vidcap = cv2.VideoCapture(self.dir_demo) | |
| total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| vidwri = cv2.VideoWriter( | |
| os.path.join('experiment', ('{}_x{}.avi'.format(self.filename, scale))), | |
| cv2.VideoWriter_fourcc(*'XVID'), | |
| vidcap.get(cv2.CAP_PROP_FPS), | |
| ( | |
| int(scale * vidcap.get(cv2.CAP_PROP_FRAME_WIDTH)), | |
| int(scale * vidcap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| ) | |
| ) | |
| total_times = 0 | |
| tqdm_test = tqdm(range(total_frames), ncols=80) | |
| if self.arch == 'edsr': | |
| for _ in tqdm_test: | |
| success, lr = vidcap.read() | |
| if not success: break | |
| start_time = time.time() | |
| lr_y_image, = common.set_channel(lr, n_channels=3) | |
| lr_y_image, = common.np_prepare(lr_y_image, rgb_range=self.rgb_range) | |
| sr = self.session.run(self.output_names, {self.input_name: lr_y_image}) | |
| end_time = time.time() | |
| total_times += end_time - start_time | |
| if isinstance(sr, (list, tuple)): | |
| sr = from_numpy(sr[0]) if len(sr) == 1 else [from_numpy(x) for x in sr] | |
| else: | |
| sr = from_numpy(sr) | |
| sr = quantize(sr, self.rgb_range).squeeze(0) | |
| normalized = sr * 255 / self.rgb_range | |
| ndarr = normalized.transpose(1, 2, 0).astype(np.uint8) | |
| vidwri.write(ndarr) | |
| elif self.arch == 'espcn': | |
| for _ in tqdm_test: | |
| success, lr = vidcap.read() | |
| if not success: break | |
| start_time = time.time() | |
| lr_y_image, lr_cb_image, lr_cr_image = imgproc.preprocess_one_frame(lr) | |
| bic_cb_image = cv2.resize(lr_cb_image, | |
| (int(lr_cb_image.shape[1] * scale), | |
| int(lr_cb_image.shape[0] * scale)), | |
| interpolation=cv2.INTER_CUBIC) | |
| bic_cr_image = cv2.resize(lr_cr_image, | |
| (int(lr_cr_image.shape[1] * scale), | |
| int(lr_cr_image.shape[0] * scale)), | |
| interpolation=cv2.INTER_CUBIC) | |
| sr = self.session.run(self.output_names, {self.input_name: lr_y_image}) | |
| end_time = time.time() | |
| total_times += end_time - start_time | |
| if isinstance(sr, (list, tuple)): | |
| sr = from_numpy(sr[0]) if len(sr) == 1 else [from_numpy(x) for x in sr] | |
| else: | |
| sr = from_numpy(sr) | |
| ndarr = imgproc.array_to_image(sr) | |
| sr_y_image = ndarr.astype(np.float32) / 255.0 | |
| sr_ycbcr_image = cv2.merge([sr_y_image[:, :, 0], bic_cb_image, bic_cr_image]) | |
| sr_image = imgproc.ycbcr_to_bgr(sr_ycbcr_image) | |
| sr_image = np.clip(sr_image* 255.0, 0 , 255).astype(np.uint8) | |
| vidwri.write(sr_image) | |
| print('Total time: {:.3f} seconds for {} frames'.format(total_times, total_frames)) | |
| print('Average time: {:.3f} seconds for each frame'.format(total_times / total_frames)) | |
| vidcap.release() | |
| vidwri.release() | |
| torch.set_grad_enabled(True) | |
| def main(): | |
| args = parser.parse_args() | |
| t = VideoTester(args.scale, args.model, args.dir_demo, arch=args.arch) | |
| t.test() | |
| if __name__ == '__main__': | |
| main() | |