| import os |
| import math |
|
|
| import utility |
| from data import common |
|
|
| import torch |
| import cv2 |
|
|
| from tqdm import tqdm |
|
|
| class VideoTester(): |
| def __init__(self, args, my_model, ckp): |
| self.args = args |
| self.scale = args.scale |
|
|
| self.ckp = ckp |
| self.model = my_model |
|
|
| self.filename, _ = os.path.splitext(os.path.basename(args.dir_demo)) |
|
|
| def test(self): |
| torch.set_grad_enabled(False) |
|
|
| self.ckp.write_log('\nEvaluation on video:') |
| self.model.eval() |
|
|
| timer_test = utility.timer() |
| for idx_scale, scale in enumerate(self.scale): |
| vidcap = cv2.VideoCapture(self.args.dir_demo) |
| total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT)) |
| vidwri = cv2.VideoWriter( |
| self.ckp.get_path('{}_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)) |
| ) |
| ) |
|
|
| tqdm_test = tqdm(range(total_frames), ncols=80) |
| for _ in tqdm_test: |
| success, lr = vidcap.read() |
| if not success: break |
|
|
| lr, = common.set_channel(lr, n_channels=self.args.n_colors) |
| lr, = common.np2Tensor(lr, rgb_range=self.args.rgb_range) |
| lr, = self.prepare(lr.unsqueeze(0)) |
| sr = self.model(lr, idx_scale) |
| sr = utility.quantize(sr, self.args.rgb_range).squeeze(0) |
|
|
| normalized = sr * 255 / self.args.rgb_range |
| ndarr = normalized.byte().permute(1, 2, 0).cpu().numpy() |
| vidwri.write(ndarr) |
|
|
| vidcap.release() |
| vidwri.release() |
|
|
| self.ckp.write_log( |
| 'Total: {:.2f}s\n'.format(timer_test.toc()), refresh=True |
| ) |
| torch.set_grad_enabled(True) |
|
|
| def prepare(self, *args): |
| device = torch.device('cpu' if self.args.cpu else 'cuda') |
| def _prepare(tensor): |
| if self.args.precision == 'half': tensor = tensor.half() |
| return tensor.to(device) |
|
|
| return [_prepare(a) for a in args] |
|
|
|
|
