| | import cv2 |
| | import torch |
| | import onnx |
| | import onnxruntime |
| | import numpy as np |
| |
|
| | import time |
| |
|
| | |
| | |
| |
|
| |
|
| | class CodeFormerEnhancer: |
| | def __init__(self, model_path="codeformer.onnx", device='cpu'): |
| | model = onnx.load(model_path) |
| | session_options = onnxruntime.SessionOptions() |
| | session_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL |
| | providers = ["CPUExecutionProvider"] |
| | if device == 'cuda': |
| | providers = [("CUDAExecutionProvider", {"cudnn_conv_algo_search": "DEFAULT"}),"CPUExecutionProvider"] |
| | self.session = onnxruntime.InferenceSession(model_path, sess_options=session_options, providers=providers) |
| |
|
| | def enhance(self, img, w=0.9): |
| | img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR) |
| | img = img.astype(np.float32)[:,:,::-1] / 255.0 |
| | img = img.transpose((2, 0, 1)) |
| | nrm_mean = np.array([0.5, 0.5, 0.5]).reshape((-1, 1, 1)) |
| | nrm_std = np.array([0.5, 0.5, 0.5]).reshape((-1, 1, 1)) |
| | img = (img - nrm_mean) / nrm_std |
| |
|
| | img = np.expand_dims(img, axis=0) |
| |
|
| | out = self.session.run(None, {'x':img.astype(np.float32), 'w':np.array([w], dtype=np.double)})[0] |
| | out = (out[0].transpose(1,2,0).clip(-1,1) + 1) * 0.5 |
| | out = (out * 255)[:,:,::-1] |
| |
|
| | return out.astype('uint8') |
| |
|
