# import onnxruntime
import axengine as axe

CLASS_NAMES = [
    "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
    "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
    "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
    "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
    "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
    "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
    "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
    "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
    "hair drier", "toothbrush"]


class axmodel_inferencer:

    def __init__(self, model_path) -> None:
        # self.onnx_model_sess = onnxruntime.InferenceSession(model_path)
        self.onnx_model_sess = axe.InferenceSession(model_path)
        self.output_names = []
        self.input_names = []
        print(model_path)
        for i in range(len(self.onnx_model_sess.get_inputs())):
            self.input_names.append(self.onnx_model_sess.get_inputs()[i].name)
            print("    input:", i,
                  self.onnx_model_sess.get_inputs()[i].name,
                  self.onnx_model_sess.get_inputs()[i].shape)

        for i in range(len(self.onnx_model_sess.get_outputs())):
            self.output_names.append(
                self.onnx_model_sess.get_outputs()[i].name)
            print("    output:", i,
                  self.onnx_model_sess.get_outputs()[i].name,
                  self.onnx_model_sess.get_outputs()[i].shape)
        print("")

    def get_input_count(self):
        return len(self.input_names)

    def get_input_shape(self, idx: int):
        return self.onnx_model_sess.get_inputs()[idx].shape

    def get_input_names(self):
        return self.input_names

    def get_output_count(self):
        return len(self.output_names)

    def get_output_shape(self, idx: int):
        return self.onnx_model_sess.get_outputs()[idx].shape

    def get_output_names(self):
        return self.output_names

    def inference(self, tensor):
        return self.onnx_model_sess.run(
            self.output_names, input_feed={self.input_names[0]: tensor})

    def inference_multi_input(self, tensors: list):
        inputs = dict()
        for idx, tensor in enumerate(tensors):
            inputs[self.input_names[idx]] = tensor
        return self.onnx_model_sess.run(input_feed=inputs)
    
    def numpy_sigmoid(self,x):
        """
        用NumPy实现的sigmoid函数
        
        参数:
        x (np.ndarray): 输入数组
        
        返回:
        np.ndarray: 经过sigmoid处理后的数组
        """
        return 1 / (1 + np.exp(-x))
    



if __name__ == "__main__":
    axmodel_model_path = "rtdetr_msda.axmodel"
    test_model = axmodel_inferencer(axmodel_model_path)

        # import onnxruntime as ort 
    from PIL import Image, ImageDraw
    # from torchvision.transforms import ToTensor
    import numpy as np
    # import torch

    # # print(onnx.helper.printable_graph(mm.graph))


    image = Image.open('ssd_horse.jpg').convert('RGB')
    im = image.resize((640, 640))
    im_data = np.array([im])
    print(im_data.shape)

    pred_logits,pred_boxes = test_model.inference(im_data)

    pred_logits = np.array(pred_logits)
    pred_boxes = np.array(pred_boxes)
    print(pred_boxes.shape,pred_logits.shape)

    
    # pred_logits = 1/(1+np.exp(-pred_logits))

    pred_logits = test_model.numpy_sigmoid(pred_logits)


    # print(pred["pred_logits"].shape,pred["pred_boxes"].shape)
    # argmax = torch.argmax(pred_logits,2).reshape(-1)
    argmax = np.argmax(pred_logits, axis=2).reshape(-1)
    print(argmax.shape)

    # pred_logits = pred["pred_logits"]
    # pred_boxes = pred["pred_boxes"]
    draw = ImageDraw.Draw(image)

    for i,idx in enumerate(argmax):
        score = pred_logits[0,i,idx]
        if score > 0.6:
            print(score,idx)
            bbox = pred_boxes[0,i]
            print(bbox)
            cx,cy,w,h = bbox
            x0 = (cx-0.5*w)*image.width
            y0 = (cy-0.5*h)*image.height
            x1 = (cx+0.5*w)*image.width
            y1 = (cy+0.5*h)*image.height
            draw.rectangle([x0,y0,x1,y1],outline="red")
            draw.text([x0,y0],CLASS_NAMES[idx]+" %.2f"%score)
    image.save("output.jpg")