VERSION = "0.0.1"
import onnxruntime as ort
import cv2 
import numpy as np

class Det:
    def __init__(self, model_path="src/rai.onnx", iou_thres:float=0.5):
        self.session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
        self.input_name = self.session.get_inputs()[0].name
        self.output_name = self.session.get_outputs()[0].name
        self.input_shape = self.session.get_inputs()[0].shape
        self.inputH = self.input_shape[2]
        self.inputW = self.input_shape[3]        
        self.iou_thres = iou_thres
        self.imgH = self.inputH
        self.imgW = self.inputW

    def preprocess(self, image):
        self.imgH, self.imgW = image.shape[:2]
        img = cv2.resize(image, (self.inputW, self.inputH))
        img = np.array(img) / 255.0
        img = np.transpose(img, (2, 0, 1))
        img = img.astype(np.float32)
        img = np.expand_dims(img, axis=0)
        return img

    def detect(self, image, conf=0.5):
        input_data = self.preprocess(image)
        outputs = self.session.run([self.output_name], {self.input_name: input_data})
        results = self.postprocess(outputs, conf=conf)
        return results

    def postprocess(self, ioutputs, conf):
        outputs = np.transpose(np.squeeze(ioutputs[0]))
        rows = outputs.shape[0]

        boxes = []
        scores = []

        for i in range(rows):
            classes_scores = outputs[i][4:]
            max_score = np.amax(classes_scores)
            if max_score >= conf:
                print(classes_scores)
                # class_id = np.argmax(classes_scores)
                x, y, w, h = outputs[i][0], outputs[i][1], outputs[i][2], outputs[i][3]
                # __hW = w/2
                # __hH = h/2
                # __left = x - __hW
                # __top = y - __hH
                # __right = x + __hW
                # __bottom = y + __hH

                scores.append(max_score)
                boxes.append([x, y, w, h])
                # boxes.append([__left, __top, __right, __bottom])

        indices = cv2.dnn.NMSBoxes(boxes, scores, conf, self.iou_thres)
        results = []
        for i in np.array(indices).flatten():
            box = boxes[int(i)]
            score = scores[int(i)]

            box = [int(x) for x in box]
            score = float(score)
            
            print(box, score)
            results.append([score, box])
        return results

    def __version__(self):
        print(VERSION)

def is_intel_cpu():
    try:
        import cpuinfo
        brand = cpuinfo.get_cpu_info()['brand_raw']
        return "Intel" in brand
    except Exception:
        return False

if __name__ == '__main__':
    # golet = Goletrai()
    # exit()
    det = Det("src/rai.onnx")

    vPlayer = cv2.VideoCapture(0)
    if not vPlayer.isOpened():
        exit()

    while True:
        ret, image = vPlayer.read()
        if not ret: break

        results = det.detect(image)
        # for result in results:
        #     box
        # print(results)

        cv2.imshow("image", image)
        key = cv2.waitKey(1)
        if key == ord('q'): break

    