69 lines
2.4 KiB
Python
69 lines
2.4 KiB
Python
import argparse
|
|
import cv2
|
|
import glob
|
|
import numpy as np
|
|
import os
|
|
import torch
|
|
from basicsr.archs.rrdbnet_arch import RRDBNet
|
|
from torch.nn import functional as F
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--model_path', type=str, default='experiments/pretrained_models/RealESRGAN_x4plus.pth')
|
|
parser.add_argument('--scale', type=int, default=4)
|
|
parser.add_argument('--input', type=str, default='inputs', help='input image or folder')
|
|
args = parser.parse_args()
|
|
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
# set up model
|
|
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=args.scale)
|
|
loadnet = torch.load(args.model_path)
|
|
model.load_state_dict(loadnet['params_ema'], strict=True)
|
|
model.eval()
|
|
model = model.to(device)
|
|
|
|
os.makedirs('results/', exist_ok=True)
|
|
for idx, path in enumerate(sorted(glob.glob(os.path.join(args.input, '*')))):
|
|
imgname = os.path.splitext(os.path.basename(path))[0]
|
|
print('Testing', idx, imgname)
|
|
# read image
|
|
img = cv2.imread(path, cv2.IMREAD_COLOR).astype(np.float32) / 255.
|
|
img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float()
|
|
img = img.unsqueeze(0).to(device)
|
|
|
|
if args.scale == 2:
|
|
mod_scale = 2
|
|
elif args.scale == 1:
|
|
mod_scale = 4
|
|
else:
|
|
mod_scale = None
|
|
if mod_scale is not None:
|
|
h_pad, w_pad = 0, 0
|
|
_, _, h, w = img.size()
|
|
if (h % mod_scale != 0):
|
|
h_pad = (mod_scale - h % mod_scale)
|
|
if (w % mod_scale != 0):
|
|
w_pad = (mod_scale - w % mod_scale)
|
|
img = F.pad(img, (0, w_pad, 0, h_pad), 'reflect')
|
|
|
|
try:
|
|
# inference
|
|
with torch.no_grad():
|
|
output = model(img)
|
|
# remove extra pad
|
|
if mod_scale is not None:
|
|
_, _, h, w = output.size()
|
|
output = output[:, :, 0:h - h_pad, 0:w - w_pad]
|
|
# save image
|
|
output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy()
|
|
output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))
|
|
output = (output * 255.0).round().astype(np.uint8)
|
|
cv2.imwrite(f'results/{imgname}_RealESRGAN.png', output)
|
|
except Exception as error:
|
|
print('Error', error)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|