Spaces:
Running
Running
File size: 1,367 Bytes
613c9ab |
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 |
import impact.additional_dependencies
from impact.utils import *
impact.additional_dependencies.ensure_onnx_package()
try:
import onnxruntime
def onnx_inference(image, onnx_model):
# prepare image
pil = tensor2pil(image)
image = np.ascontiguousarray(pil)
image = image[:, :, ::-1] # to BGR image
image = image.astype(np.float32)
image -= [103.939, 116.779, 123.68] # 'caffe' mode image preprocessing
# do detection
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]
# filter-out useless item
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}")
|