/
vsd.py
61 lines (44 loc) · 2.06 KB
/
vsd.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
import torch
import numpy as np
import PIL.Image as pil_image
import os
from models import FSRCNN, FSRCNN_VSD
from utils import convert_ycbcr_to_rgb, preprocess, calc_psnr
from torchsummary import summary
if __name__ == '__main__':
torch.backends.quantized.engine = 'qnnpack'
args_scale = 4
args_img = 'data/butterfly_GT.bmp'
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model = FSRCNN_VSD(scale_factor=4).to(device)
print(summary(model, (1, 25, 25)))
model.eval()
model.qconfig = torch.quantization.get_default_qconfig('qnnpack')
torch.quantization.prepare(model, inplace=True)
torch.quantization.convert(model, inplace=True)
model.load_state_dict(torch.load(
'./vsd/Qbest.pth', map_location=lambda storage, loc: storage))
image = pil_image.open(args_img).resize((500,500)).convert('RGB')
image_width = (image.width // args_scale) * args_scale
image_height = (image.height // args_scale) * args_scale
hr = image.resize((image_width, image_height), resample=pil_image.BICUBIC)
lr = hr.resize((hr.width // args_scale, hr.height //
args_scale), resample=pil_image.BICUBIC)
bicubic = lr.resize((lr.width * args_scale, lr.height *
args_scale), resample=pil_image.BICUBIC)
bicubic.save(args_img.replace('.', '_bicubic_x{}.'.format(args_scale)))
lr, _ = preprocess(lr, device)
hr, _ = preprocess(hr, device)
bicubic, ycbcr = preprocess(bicubic, device)
with torch.no_grad():
preds = model(lr).clamp(0.0, 1.0)
psnr = calc_psnr(hr, preds)
print('HR PSNR: {:.2f}'.format(psnr))
psnr = calc_psnr(hr, bicubic)
print('Bic PSNR: {:.2f}'.format(psnr))
preds = preds.mul(255.0).cpu().numpy().squeeze(0).squeeze(0)
output = np.array([preds, ycbcr[..., 1], ycbcr[..., 2]]
).transpose([1, 2, 0])
output = np.clip(convert_ycbcr_to_rgb(output), 0.0, 255.0).astype(np.uint8)
output = pil_image.fromarray(output)
output.save(args_img.replace('.', '_fsrcnn_x{}.'.format(args_scale)))