| import numpy as np |
| import cv2 |
| import onnxruntime |
| import roop.globals |
|
|
| from roop.typing import Frame |
| from roop.utilities import resolve_relative_path, conditional_thread_semaphore |
|
|
|
|
|
|
| class Mask_XSeg(): |
| plugin_options:dict = None |
|
|
| model_xseg = None |
|
|
| processorname = 'mask_xseg' |
| type = 'mask' |
|
|
|
|
| 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_xseg is None: |
| model_path = resolve_relative_path('../models/xseg.onnx') |
| onnxruntime.set_default_logger_severity(3) |
| self.model_xseg = onnxruntime.InferenceSession(model_path, None, providers=roop.globals.execution_providers) |
| self.model_inputs = self.model_xseg.get_inputs() |
| self.model_outputs = self.model_xseg.get_outputs() |
|
|
| |
| self.devicename = self.plugin_options["devicename"].replace('mps', 'cpu') |
|
|
|
|
| def Run(self, img1, keywords:str) -> Frame: |
| temp_frame = cv2.resize(img1, (256, 256), cv2.INTER_CUBIC) |
| temp_frame = temp_frame.astype('float32') / 255.0 |
| temp_frame = temp_frame[None, ...] |
| io_binding = self.model_xseg.io_binding() |
| io_binding.bind_cpu_input(self.model_inputs[0].name, temp_frame) |
| io_binding.bind_output(self.model_outputs[0].name, self.devicename) |
| self.model_xseg.run_with_iobinding(io_binding) |
| ort_outs = io_binding.copy_outputs_to_cpu() |
| result = ort_outs[0][0] |
| result = np.clip(result, 0, 1.0) |
| result[result < 0.1] = 0 |
| |
| result = 1.0 - result |
| return result |
|
|
|
|
| def Release(self): |
| del self.model_xseg |
| self.model_xseg = None |
|
|
|
|
|
|
