Spaces:
Build error
Build error
import os | |
import torch | |
import numpy as np | |
from torchvision import transforms | |
from PIL import Image | |
import time | |
import torchvision | |
import argparse | |
from models.SCET import SCET | |
def inference_img(img_path,Net,device): | |
low_image = Image.open(img_path).convert('RGB') | |
enhance_transforms = transforms.Compose([ | |
transforms.ToTensor() | |
]) | |
with torch.no_grad(): | |
low_image = enhance_transforms(low_image) | |
low_image = low_image.unsqueeze(0) | |
start = time.time() | |
restored2 = Net(low_image.to(device)) | |
end = time.time() | |
return restored2,end-start | |
if __name__ == '__main__': | |
parser=argparse.ArgumentParser() | |
parser.add_argument('--test_path',type=str,required=True,help='Path to test') | |
parser.add_argument('--save_path',type=str,required=True,help='Path to save') | |
parser.add_argument('--pk_path',type=str,default='model_zoo/SRx4.pth',help='Path of the checkpoint') | |
parser.add_argument('--scale',type=int,default=4,help='scale factor') | |
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
opt = parser.parse_args() | |
if not os.path.isdir(opt.save_path): | |
os.mkdir(opt.save_path) | |
if opt.scale == 3: | |
Net = SCET(63, 128, opt.scale).eval() | |
else: | |
Net = SCET(64, 128, opt.scale).eval() | |
Net.load_state_dict(torch.load(opt.pk_path)) | |
Net=Net.to(device) | |
image=opt.test_path | |
print(image) | |
restored2,time_num=inference_img(image,Net,device) | |
torchvision.utils.save_image(restored2,opt.save_path+'output.png') | |