import numpy as np
from mindx.sdk import base
from mindx.sdk.base import Tensor, Model, Size, log, ImageProcessor, post
import cv2
"""
本样例使用以下模型：
https://gitee.com/ascend/ModelZoo-TensorFlow/tree/master/ACL_TensorFlow/built-in/cv/YOLOv3_for_ACL/
其中ATC命令需添加输出节点指定：
--out_nodes="yolov3/yolov3_head/Conv_6/BiasAdd:0;yolov3/yolov3_head/Conv_14/BiasAdd:0;yolov3/yolov3_head/Conv_22/BiasAdd:0"
"""
device_id = 0  # 芯片ID
image_path = "./test.jpg"  # 输入图片

b_usedvpp = True  # 使用dvpp图像处理器时启用，使用opencv为False
b_useaipp = True   # ATC转换时使用--insert_op_conf参数为模型添加了AIPP时启用
yolo_resizelen = 416  # 模型输入大小，YOLOv3长宽均为416


def main():
    # ****0 初始化
    base.mx_init()  # 全局资源初始化
    imageTensorList = []

    # ****1 模型前处理
    if b_usedvpp:
        print("using ImageProcessor for preprocess.")
        # 创造图像处理器对象，使用该方法处理后数据在device侧
        imageProcessor0 = ImageProcessor(device_id)

        if b_useaipp:
            # 本分支为使用aipp的YOLOV3模型(ATC转换时使用--insert_op_conf参数),限定输入图像为YUV420SP
            model_path = "./model/yolov3_tf_bs1_fp16.om"
            yolov3 = Model(model_path, device_id)  # 创造模型对象

            decodedImg = imageProcessor0.decode(image_path, base.nv12) # HNWC
            size_cof = Size(yolo_resizelen, yolo_resizelen)
            resizeImg = imageProcessor0.resize(decodedImg, size_cof)

            # 推理需要转换为tensor的List(数据已在device侧无需转移)
            imageTensorList = [resizeImg.to_tensor()]
        else:
            # 本分支为未使用aipp的YOLOV3模型(输入同模型原始输入:BGR格式float32)
            model_path = "./model/yolov3_bs1.om"
            yolov3 = Model(model_path, device_id)  # 创造模型对象

            decodedImg = imageProcessor0.decode(image_path, base.bgr) # HNWC
            size_cof = Size(yolo_resizelen, yolo_resizelen)
            resizeImg = imageProcessor0.resize(decodedImg, size_cof)

            resizeImg.to_host()  # 需要在host侧才能自行处理数据
            image_resize = np.array(resizeImg.to_tensor())  # NHWC，取出为numpy数组

            img_ndarray = image_resize.astype(np.float32)/255  # int8->float32
            img_mxtensor = Tensor(img_ndarray) # 转换为Tensor对象
            img_mxtensor.to_device(device_id) # 推理前需部署到device侧
            # 推理需要转换为tensor的List
            imageTensorList = [img_mxtensor]

    else:
        print("using opencv for preprocess.")
        # 本分支为未使用aipp的YOLOV3模型(输入为BGR格式float32)
        model_path = "./model/yolov3_bs1.om"
        yolov3 = Model(model_path, device_id)  # 创造模型对象

        image_cv2 = cv2.imread(image_path)  # HWC
        image_nd_fp32 = image_cv2.astype(np.float32)/255  # int8->float32
        size_cof = (yolo_resizelen, yolo_resizelen)
        resizeImg = cv2.resize(image_nd_fp32, size_cof,
                               interpolation=cv2.INTER_LINEAR)
        np_image_addbatch = np.expand_dims(resizeImg, axis=0)  # NHWC

        imageTensor = Tensor(np_image_addbatch)  # 推理前需要转换为tensor，使用Tensor类来构建。
        imageTensor.to_device(device_id)  # !需要转移至device侧，该函数单独执行
        imageTensorList = [imageTensor]  # 构造tensor的List类型
        """
        ！！！如使用了transpose,slice,append,reshape等改变数据内存形状的操作后，需要使用numpy.ascontiguousarray对内存进行重新排序成连续的
        如使用非图像数据，也是转为numpy.ndarray数据类型再进行Tensor转换，使用{tensor_data} = Tensor({numpy_data})方式
        外部文件读入的numpy输入（例如np.fromfile）需要reshpe为对应的shape
        for i in range(input.shape[0]):
            input_tensor = Tensor(inputs[i, :].reshape(1,-1)) # 每个batch的内容转换
            input_tensor.to_device(device_id)
            input_tensors.append(input_tensor)
        """

    # ****2 模型推理
    outputs = yolov3.infer(imageTensorList) 

    post_tensor_inputs = [] # tensor结果数组
    nps = [] # 原始numpy结果数组
    for i in range(len(outputs)):
        outputs[i].to_host()
        n = np.array(outputs[i])
        nps.append(n) # 使用自定义后处理时直接用此numpy数组即可!!!
        tensor = Tensor(n)  # 使用内置后处理类型时需要转换为Tenosr
        post_tensor_inputs.append(tensor)

    # ****3 SDK内置模型后处理/自行编写后处理时建议直接使用numpy结果数组
    config_path = "./model/yolov3_tf_bs1_fp16.cfg"  # 模型配置文件的路径
    label_path = "./model/yolov3.names"  # 分类标签文件的路径

    yolov3_post = post.Yolov3PostProcess(
        config_path=config_path, label_path=label_path)  # 构造对应的后处理对象

    resizeInfo = base.ResizedImageInfo()
    resizeInfo.heightResize = yolo_resizelen
    resizeInfo.widthResize = yolo_resizelen
    if b_usedvpp:
        resizeInfo.heightOriginal = decodedImg.original_height
        resizeInfo.widthOriginal = decodedImg.original_width
    else:
        resizeInfo.heightOriginal = image_cv2.shape[0]
        resizeInfo.widthOriginal = image_cv2.shape[1]

    results = yolov3_post.process(post_tensor_inputs, resizeInfo)
    """
    如果使用多batch进行后处理则需要concat连接tensor，且后处理输出要对应。此处为本样例的示例：
        inputs.append(base.batch_concat([tensor] * 2))
    results = yolov3_post.process(inputs, [resizeInfo] * 2)
    """

    # ****4 结果打印
    print("\nresults:")
    for i in range(len(results)):
        for j in range(len(results[i])):
            print("bbox:", results[i][j].x0, ",", results[i]
                  [j].y0, ",", results[i][j].x1, ",", results[i][j].y1)
            print("confidence:", results[i][j].confidence)
            print("classId:", results[i][j].classId)
            print("className:", results[i][j].className)
            print("******")
            
    base.mx_deinit()

try:
    main()
except Exception as e:
    print(e)