Spaces:
Runtime error
Runtime error
File size: 1,542 Bytes
1fe5b59 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
import cv2
import torch
import onnx
import onnxruntime
import numpy as np
import time
# codeformer converted to onnx
# using https://github.com/redthing1/CodeFormer
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")
|