import glob import sys sys.path.append('../..') from natsort import natsort import SRFlow.code.options.options as option import torch from SRFlow.code.utils.util import opt_get from SRFlow.code.models.SRFlow_model import SRFlowModel import numpy as np import os import cv2 def fiFindByWildcard(wildcard): return natsort.natsorted(glob.glob(wildcard, recursive=True)) def load_model(conf_path): opt = option.parse(conf_path, is_train=False) opt['gpu_ids'] = None opt = option.dict_to_nonedict(opt) model = SRFlowModel(opt, 0) model_path = opt_get(opt, ['model_path'], None) model.load_network(load_path='models/SRFlow/35000_G.pth', network=model.netG) return model, opt def predict(model, lr): model.feed_data({"LQ": t(lr)}, need_GT=False) model.test() visuals = model.get_current_visuals(need_GT=False) return visuals.get('rlt', visuals.get("SR")) def t(array): return torch.Tensor(np.expand_dims(array.transpose([2, 0, 1]), axis=0).astype(np.float32)) / 255 def rgb(t): return ( np.clip((t[0] if len(t.shape) == 4 else t).detach().cpu().numpy().transpose([1, 2, 0]), 0, 1) * 255).astype( np.uint8) def imread(path): return cv2.imread(path)[:, :, [2, 1, 0]] def imwrite(path, img): os.makedirs(os.path.dirname(path), exist_ok=True) cv2.imwrite(path, img[:, :, [2, 1, 0]]) def imCropCenter(img, size): h, w, c = img.shape h_start = max(h // 2 - size // 2, 0) h_end = min(h_start + size, h) w_start = max(w // 2 - size // 2, 0) w_end = min(w_start + size, w) return img[h_start:h_end, w_start:w_end] def impad(img, top=0, bottom=0, left=0, right=0, color=255): return np.pad(img, [(top, bottom), (left, right), (0, 0)], 'reflect')