# -*- coding: utf-8 -*-
import cv2
import numpy as np
import os
import argparse
import time
import migraphx


class YOLOv3:
    def __init__(self, path, obj_thres=0.5, conf_thres=0.25, iou_thres=0.5):
        self.objectThreshold = obj_thres
        self.confThreshold = conf_thres
        self.nmsThreshold = iou_thres

        # 获取模型检测的类别信息
        self.classNames = list(map(lambda x: x.strip(), open('../Resource/Models/coco.names', 'r').readlines()))

        # 解析推理模型
        self.model = migraphx.parse_onnx(path)

        # 获取模型输入/输出节点信息
        print("inputs:")
        inputs = self.model.get_inputs()
        for key,value in inputs.items():
            print("{}:{}".format(key,value))
        
        print("outputs:")
        outputs = self.model.get_outputs()
        for key,value in outputs.items():
            print("{}:{}".format(key,value))

        # 获取模型的输入name
        self.inputName = "images"

        # 获取模型的输入尺寸
        inputShape = inputShape=inputs[self.inputName].lens()
        self.inputHeight = int(inputShape[2])
        self.inputWidth = int(inputShape[3])
        print("inputName:{0} \ninputShape:{1}".format(self.inputName, inputShape))

    def detect(self, image):
        # 输入图片预处理
        input_img = self.prepare_input(image)

        # 模型编译
        self.model.compile(t=migraphx.get_target("gpu"), device_id=0)  # device_id: 设置GPU设备，默认为0号设备
        print("Success to compile")
        # 执行推理
        print("Start to inference")
        start = time.time()
        result = self.model.run({self.model.get_parameter_names()[0]: input_img})
        print('net forward time: {:.4f}'.format(time.time() - start))
        # 模型输出结果后处理
        boxes, scores, class_ids = self.process_output(result)

        return boxes, scores, class_ids

    def prepare_input(self, image):
        self.img_height, self.img_width = image.shape[:2]
        input_img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        input_img = cv2.resize(input_img, (self.inputWidth, self.inputHeight))
        input_img = input_img.transpose(2, 0, 1)
        input_img = np.expand_dims(input_img, 0)
        input_img = np.ascontiguousarray(input_img)
        input_img = input_img.astype(np.float32)
        input_img = input_img / 255

        return input_img

    def process_output(self, output):
        predictions = np.squeeze(output[0])

        # 筛选包含物体的anchor
        obj_conf = predictions[:, 4]
        predictions = predictions[obj_conf > self.objectThreshold]
        obj_conf = obj_conf[obj_conf > self.objectThreshold]

        # 筛选大于置信度阈值的anchor
        predictions[:, 5:] *= obj_conf[:, np.newaxis]
        scores = np.max(predictions[:, 5:], axis=1)
        valid_scores = scores > self.confThreshold
        predictions = predictions[valid_scores]
        scores = scores[valid_scores]

        # 获取最高置信度分数对应的类别ID
        class_ids = np.argmax(predictions[:, 5:], axis=1)

        # 获取每个物体对应的anchor
        boxes = self.extract_boxes(predictions)

        # 执行非极大值抑制消除冗余anchor
        indices = cv2.dnn.NMSBoxes(boxes.tolist(), scores.tolist(), self.confThreshold, self.nmsThreshold).flatten()

        return boxes[indices], scores[indices], class_ids[indices]

    def extract_boxes(self, predictions):
        # 获取anchor的坐标信息
        boxes = predictions[:, :4]

        # 将anchor的坐标信息映射到输入image
        boxes = self.rescale_boxes(boxes)

        # 格式转换
        boxes_ = np.copy(boxes)
        boxes_[..., 0] = boxes[..., 0] - boxes[..., 2] * 0.5
        boxes_[..., 1] = boxes[..., 1] - boxes[..., 3] * 0.5
        return boxes_

    def rescale_boxes(self, boxes):
        # 对anchor尺寸进行变换
        input_shape = np.array([self.inputWidth, self.inputHeight, self.inputWidth, self.inputHeight])
        boxes = np.divide(boxes, input_shape, dtype=np.float32)
        boxes *= np.array([self.img_width, self.img_height, self.img_width, self.img_height])
        return boxes

    def draw_detections(self, image, boxes, scores, class_ids):
        for box, score, class_id in zip(boxes, scores, class_ids):
            cx, cy, w, h = box.astype(int)

            # 绘制检测物体框
            cv2.rectangle(image, (cx, cy), (cx + w, cy + h), (0, 255, 255), thickness=2)
            label = self.classNames[class_id]
            label = f'{label} {score:.2f}'
            labelSize, baseLine = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
            cv2.putText(image, label, (cx, cy - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), thickness=2)
        return image


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--imgpath', type=str, default='../Resource/Images/dog.jpg', help="image path")
    parser.add_argument('--modelpath', type=str, default='../Resource/Models/yolov3-tiny.onnx', help="onnx filepath")
    parser.add_argument('--objectThreshold', default=0.4, type=float, help='class confidence')
    parser.add_argument('--confThreshold', default=0.2, type=float, help='class confidence')
    parser.add_argument('--nmsThreshold', default=0.4, type=float, help='nms iou thresh')
    args = parser.parse_args()

    yolov3_detector = YOLOv3(args.modelpath, obj_thres=args.objectThreshold, conf_thres=args.confThreshold,
                             iou_thres=args.nmsThreshold)
    srcimg = cv2.imread(args.imgpath, 1)

    boxes, scores, class_ids = yolov3_detector.detect(srcimg)

    dstimg = yolov3_detector.draw_detections(srcimg, boxes, scores, class_ids)

    # 保存检测结果
    cv2.imwrite("./Result.jpg", dstimg)
    print("Success to save result")