Spaces:
Runtime error
Runtime error
File size: 3,429 Bytes
4c16f3d | 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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | import cv2
import numpy as np
import onnxruntime
import roop.globals
from roop.utilities import resolve_relative_path
from roop.typing import Frame
class Frame_Colorizer:
plugin_options: dict = None
model_colorizer = None
devicename = None
prev_type = None
processorname = "deoldify"
type = "frame_colorizer"
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.prev_type is not None
and self.prev_type != self.plugin_options["subtype"]
):
self.Release()
self.prev_type = self.plugin_options["subtype"]
if self.model_colorizer is None:
# replace Mac mps with cpu for the moment
self.devicename = self.plugin_options["devicename"].replace("mps", "cpu")
if self.prev_type == "deoldify_artistic":
model_path = resolve_relative_path(
"../models/Frame/deoldify_artistic.onnx"
)
elif self.prev_type == "deoldify_stable":
model_path = resolve_relative_path(
"../models/Frame/deoldify_stable.onnx"
)
onnxruntime.set_default_logger_severity(3)
self.model_colorizer = onnxruntime.InferenceSession(
model_path, None, providers=roop.globals.execution_providers
)
self.model_inputs = self.model_colorizer.get_inputs()
model_outputs = self.model_colorizer.get_outputs()
self.io_binding = self.model_colorizer.io_binding()
self.io_binding.bind_output(model_outputs[0].name, self.devicename)
def Run(self, input_frame: Frame) -> Frame:
temp_frame = cv2.cvtColor(input_frame, cv2.COLOR_BGR2GRAY)
temp_frame = cv2.cvtColor(temp_frame, cv2.COLOR_GRAY2RGB)
temp_frame = cv2.resize(temp_frame, (256, 256))
temp_frame = temp_frame.transpose((2, 0, 1))
temp_frame = np.expand_dims(temp_frame, axis=0).astype(np.float32)
self.io_binding.bind_cpu_input(self.model_inputs[0].name, temp_frame)
self.model_colorizer.run_with_iobinding(self.io_binding)
ort_outs = self.io_binding.copy_outputs_to_cpu()
result = ort_outs[0][0]
del ort_outs
colorized_frame = result.transpose(1, 2, 0)
colorized_frame = cv2.resize(
colorized_frame, (input_frame.shape[1], input_frame.shape[0])
)
temp_blue_channel, _, _ = cv2.split(input_frame)
colorized_frame = cv2.cvtColor(colorized_frame, cv2.COLOR_BGR2RGB).astype(
np.uint8
)
colorized_frame = cv2.cvtColor(colorized_frame, cv2.COLOR_BGR2LAB)
_, color_green_channel, color_red_channel = cv2.split(colorized_frame)
colorized_frame = cv2.merge(
(temp_blue_channel, color_green_channel, color_red_channel)
)
colorized_frame = cv2.cvtColor(colorized_frame, cv2.COLOR_LAB2BGR)
return colorized_frame.astype(np.uint8)
def Release(self):
del self.model_colorizer
self.model_colorizer = None
del self.io_binding
self.io_binding = None
|