|
import impact.additional_dependencies |
|
from impact.utils import * |
|
|
|
impact.additional_dependencies.ensure_onnx_package() |
|
|
|
try: |
|
import onnxruntime |
|
|
|
def onnx_inference(image, onnx_model): |
|
|
|
pil = tensor2pil(image) |
|
image = np.ascontiguousarray(pil) |
|
image = image[:, :, ::-1] |
|
image = image.astype(np.float32) |
|
image -= [103.939, 116.779, 123.68] |
|
|
|
|
|
onnx_model = onnxruntime.InferenceSession(onnx_model, providers=["CPUExecutionProvider"]) |
|
outputs = onnx_model.run( |
|
[s_i.name for s_i in onnx_model.get_outputs()], |
|
{onnx_model.get_inputs()[0].name: np.expand_dims(image, axis=0)}, |
|
) |
|
|
|
labels = [op for op in outputs if op.dtype == "int32"][0] |
|
scores = [op for op in outputs if isinstance(op[0][0], np.float32)][0] |
|
boxes = [op for op in outputs if isinstance(op[0][0], np.ndarray)][0] |
|
|
|
|
|
idx = np.where(labels[0] == -1)[0][0] |
|
|
|
labels = labels[0][:idx] |
|
scores = scores[0][:idx] |
|
boxes = boxes[0][:idx].astype(np.uint32) |
|
|
|
return labels, scores, boxes |
|
except Exception as e: |
|
print("[ERROR] ComfyUI-Impact-Pack: 'onnxruntime' package doesn't support 'python 3.11', yet.") |
|
print(f"\t{e}") |
|
|