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}")