Update demo.py

demo.py

@@ -1,24 +1,17 @@
 import torch
-import models
+from models.pcsr import PCSR
+
 from torchvision import transforms
 from utils import *
 from PIL import Image
 import numpy as np
 
-img_path = 'myimage.png' # only support .png
-scale = 4 # only support x4
+##############################################
+img_path = 'sample/comic.png' # your LR input image path, only support .png
+##############################################
 
-'''
-k: hyperparameter to traverse PSNR-FLOPs trade-off. smaller k → larger FLOPs & PSNR. range is about [-1,2].
-adaptive: whether to use automatic decision of k
-no_refinement: whether not to use pixel-wise refinement (postprocessing for reducing artifacts)
-parser.add_argument('--opacity', type=float, default=0.65, help='opacity for colored visualization')
-parser.add_argument('--pixel_batch_size', type=int, default=300000)
-'''
-
-resume_path = 'carn-pcsr-phase1.pth'
-sv_file = torch.load(resume_path)
-model = models.make(sv_file['model'], load_sd=True).cuda()
+scale = 4 # only support x4
+model = PCSR.from_pretrained("3587jjh/pcsr_carn").cuda()
 model.eval()
 
 rgb_mean = torch.tensor([0.4488, 0.4371, 0.4040], device='cuda').view(1,3,1,1)
@@ -35,6 +28,11 @@ with torch.no_grad():
     cell[:,:,1] *= 2/W
     inp_lr = (lr - rgb_mean) / rgb_std
 
+    '''
+    k: hyperparameter to traverse PSNR-FLOPs trade-off. smaller k → larger FLOPs & PSNR. range is about [-1,2].
+    adaptive: whether to use automatic decision of k
+    refinement: whether to use pixel-wise refinement (postprocessing for reducing artifacts)
+    '''
     pred, flag = model(inp_lr, coord=coord, cell=cell, scale=scale, k=0, 
         pixel_batch_size=300000, adaptive_cluster=True, refinement=True)
     flops = get_model_flops(model, inp_lr, coord=coord, cell=cell, scale=scale, k=0, 
@@ -47,7 +45,7 @@ with torch.no_grad():
 pred = pred.transpose(1,2).view(-1,3,H,W)
 pred = pred * rgb_std + rgb_mean
 pred = tensor2numpy(pred)
-Image.fromarray(pred).save(f'output.png')
+Image.fromarray(pred).save(f'sample/output.png')
 
 flag = flag.view(-1,1,H,W).repeat(1,3,1,1).squeeze(0).detach().cpu()
 H,W = pred.shape[:2]
@@ -55,4 +53,4 @@ vis_img = np.zeros_like(pred)
 vis_img[flag[0] == 0] = np.array([0,255,0])
 vis_img[flag[0] == 1] = np.array([255,0,0])
 vis_img = vis_img*0.35 + pred*0.65
-Image.fromarray(vis_img.astype('uint8')).save('output_vis.png')
+Image.fromarray(vis_img.astype('uint8')).save('sample/output_vis.png')