|
|
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') |
|
|
|
