|
from typing import Any, List, Callable |
|
import cv2 |
|
import numpy as np |
|
import onnxruntime |
|
import roop.globals |
|
|
|
from roop.typing import Face, Frame, FaceSet |
|
from roop.utilities import resolve_relative_path |
|
|
|
class Enhance_CodeFormer(): |
|
model_codeformer = None |
|
|
|
plugin_options:dict = None |
|
|
|
processorname = 'codeformer' |
|
type = 'enhance' |
|
|
|
|
|
def Initialize(self, plugin_options:dict): |
|
if self.plugin_options is not None: |
|
if self.plugin_options["devicename"] != plugin_options["devicename"]: |
|
self.Release() |
|
|
|
self.plugin_options = plugin_options |
|
if self.model_codeformer is None: |
|
|
|
self.devicename = self.plugin_options["devicename"].replace('mps', 'cpu') |
|
model_path = resolve_relative_path('../models/CodeFormer/CodeFormerv0.1.onnx') |
|
self.model_codeformer = onnxruntime.InferenceSession(model_path, None, providers=roop.globals.execution_providers) |
|
self.model_inputs = self.model_codeformer.get_inputs() |
|
model_outputs = self.model_codeformer.get_outputs() |
|
self.io_binding = self.model_codeformer.io_binding() |
|
self.io_binding.bind_cpu_input(self.model_inputs[1].name, np.array([0.5])) |
|
self.io_binding.bind_output(model_outputs[0].name, self.devicename) |
|
|
|
|
|
def Run(self, source_faceset: FaceSet, target_face: Face, temp_frame: Frame) -> Frame: |
|
input_size = temp_frame.shape[1] |
|
|
|
temp_frame = cv2.resize(temp_frame, (512, 512), cv2.INTER_CUBIC) |
|
temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_BGR2RGB) |
|
temp_frame = temp_frame.astype('float32') / 255.0 |
|
temp_frame = (temp_frame - 0.5) / 0.5 |
|
temp_frame = np.expand_dims(temp_frame, axis=0).transpose(0, 3, 1, 2) |
|
|
|
self.io_binding.bind_cpu_input(self.model_inputs[0].name, temp_frame.astype(np.float32)) |
|
self.model_codeformer.run_with_iobinding(self.io_binding) |
|
ort_outs = self.io_binding.copy_outputs_to_cpu() |
|
result = ort_outs[0][0] |
|
del ort_outs |
|
|
|
|
|
result = result.transpose((1, 2, 0)) |
|
|
|
un_min = -1.0 |
|
un_max = 1.0 |
|
result = np.clip(result, un_min, un_max) |
|
result = (result - un_min) / (un_max - un_min) |
|
|
|
result = cv2.cvtColor(result, cv2.COLOR_RGB2BGR) |
|
result = (result * 255.0).round() |
|
scale_factor = int(result.shape[1] / input_size) |
|
return result.astype(np.uint8), scale_factor |
|
|
|
|
|
def Release(self): |
|
del self.model_codeformer |
|
self.model_codeformer = None |
|
del self.io_binding |
|
self.io_binding = None |
|
|
|
|
