import argparse
import cv2
import numpy as np
import onnxruntime as ort

class FreeYOLO():
    def __init__(self, model_path, confThreshold=0.4, nmsThreshold=0.85, datatype='coco'):
        so = ort.SessionOptions()
        so.log_severity_level = 3
        self.session = ort.InferenceSession(model_path, so)
        model_inputs = self.session.get_inputs()
        self.input_name = model_inputs[0].name
        self.input_shape = model_inputs[0].shape
        self.input_height = int(self.input_shape[2])
        self.input_width = int(self.input_shape[3])
        self.anchors, self.expand_strides = self.generate_anchors((self.input_height, self.input_width), [8, 16, 32])

        if datatype=='coco':
            self.classes = list(map(lambda x: x.strip(), open('coco.names', 'r').readlines()))
        elif datatype=='face':
            self.classes = ['face']
        else:
            self.classes = ['person']
        self.num_class = len(self.classes)
        self.confThreshold = confThreshold
        self.nmsThreshold = nmsThreshold

    def generate_anchors(self, input_shape, strides):
        """
            fmp_size: (List) [H, W]
        """
        all_anchors = []
        all_expand_strides = []
        for stride in strides:
            # generate grid cells
            fmp_h, fmp_w = input_shape[0] // stride, input_shape[1] // stride
            anchor_x, anchor_y = np.meshgrid(np.arange(fmp_w),
                                             np.arange(fmp_h))
            # [H, W, 2]
            anchor_xy = np.stack([anchor_x, anchor_y], axis=-1)
            shape = anchor_xy.shape[:2]
            # [H, W, 2] -> [HW, 2]
            anchor_xy = (anchor_xy.reshape(-1, 2) + 0.5) * stride
            all_anchors.append(anchor_xy)

            # expanded stride
            strides = np.full((*shape, 1), stride)
            all_expand_strides.append(strides.reshape(-1, 1))

        anchors = np.concatenate(all_anchors, axis=0)
        expand_strides = np.concatenate(all_expand_strides, axis=0)

        return anchors, expand_strides

    def decode_boxes(self, anchors, pred_regs, expand_strides):
        """
            anchors:  (List[Tensor]) [1, M, 2] or [M, 2]
            pred_reg: (List[Tensor]) [B, M, 4] or [B, M, 4]
        """
        # center of bbox
        pred_ctr_xy = anchors[..., :2] + pred_regs[..., :2] * expand_strides
        # size of bbox
        pred_box_wh = np.exp(pred_regs[..., 2:]) * expand_strides

        pred_x1y1 = pred_ctr_xy - 0.5 * pred_box_wh
        # pred_x2y2 = pred_ctr_xy + 0.5 * pred_box_wh
        # pred_box = np.concatenate([pred_x1y1, pred_x2y2], axis=-1)
        pred_box = np.concatenate([pred_x1y1, pred_box_wh], axis=-1)
        return pred_box
    def drawPred(self, frame, classId, conf, left, top, right, bottom):
        # Draw a bounding box.
        cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), thickness=2)

        label = '%.2f' % conf
        label = '%s:%s' % (self.classes[classId], label)

        # Display the label at the top of the bounding box
        labelSize, baseLine = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
        top = max(top, labelSize[1])
        # cv.rectangle(frame, (left, top - round(1.5 * labelSize[1])), (left + round(1.5 * labelSize[0]), top + baseLine), (255,255,255), cv.FILLED)
        cv2.putText(frame, label, (left, top - 10), 0, 0.7, (0, 255, 0), thickness=2)
        return frame

    def detect(self, frame):
        padded_image = np.ones((self.input_height, self.input_width, 3), dtype=np.uint8)*114
        ratio = min(self.input_height / frame.shape[0], self.input_width / frame.shape[1])
        neww, newh = int(frame.shape[1] * ratio), int(frame.shape[0] * ratio)
        temp_image = cv2.resize(frame, (neww, newh), interpolation=cv2.INTER_LINEAR)
        padded_image[:newh, :neww, :] = temp_image

        padded_image = padded_image.transpose(2, 0, 1)
        padded_image = np.expand_dims(padded_image, axis=0).astype(np.float32)

        # Inference
        results = self.session.run(None, {self.input_name: padded_image})

        reg_preds = results[0][0][..., :4]
        obj_preds = results[0][0][..., 4:5]
        cls_preds = results[0][0][..., 5:]
        scores = np.sqrt(obj_preds * cls_preds)

        # scores & class_ids
        class_ids = np.argmax(scores, axis=1)  # [M,]
        scores = np.max(scores, axis=1)

        # bboxes
        bboxes = self.decode_boxes(self.anchors, reg_preds, self.expand_strides)  # [M, 4]
        # thresh
        keep = np.where(scores > self.confThreshold)
        scores = scores[keep]
        class_ids = class_ids[keep]
        bboxes = bboxes[keep]
        bboxes /= ratio

        indices = cv2.dnn.NMSBoxes(bboxes.tolist(), scores.tolist(), self.confThreshold, self.nmsThreshold)
        for i in indices:
            left, top, width, height = bboxes[i, :].astype(np.int32)
            frame = self.drawPred(frame, class_ids[i], scores[i], left, top, left + width, top + height)
        return frame

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--modelpath", type=str, default='weights/coco/yolo_free_nano_192x320.onnx', help="model path")
    parser.add_argument("--imgpath", type=str, default='images/coco/dog.jpg', help="image path")
    parser.add_argument("--confThreshold", default=0.6, type=float, help='class confidence')
    parser.add_argument("--nmsThreshold", default=0.5, type=float, help='iou thresh')
    parser.add_argument("--datatype", default='coco', type=str, choices=['coco', 'face', 'person'], help='data type')
    args = parser.parse_args()

    net = FreeYOLO(args.modelpath, confThreshold=args.confThreshold, nmsThreshold=args.nmsThreshold, datatype=args.datatype)
    srcimg = cv2.imread(args.imgpath)
    srcimg = net.detect(srcimg)

    winName = 'Deep learning object detection in ONNXRuntime'
    cv2.namedWindow(winName, cv2.WINDOW_NORMAL)
    cv2.imshow(winName, srcimg)
    cv2.waitKey(0)
    cv2.destroyAllWindows()