# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
"""
Validate a trained YOLOv5 model accuracy on a custom dataset

Usage:
    $ python path/to/val.py --data coco128.yaml --weights yolov5s.pt --img 640
"""

import argparse
import json
import os
import sys
from pathlib import Path
from threading import Thread
import cv2
import numpy as np
import torch
from torchvision.transforms import Resize
import migraphx
from tqdm import tqdm
import glob
import time

FILE = Path(__file__).resolve()
ROOT = FILE.parents[0]  # YOLOv5 root directory
if str(ROOT) not in sys.path:
    sys.path.append(str(ROOT))  # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd()))  # relative

from models.experimental import attempt_load
from utils.datasets import create_dataloader
from utils.general import coco80_to_coco91_class, check_dataset, check_img_size, check_requirements, \
    check_suffix, check_yaml, box_iou, non_max_suppression, scale_coords, xyxy2xywh, xywh2xyxy, \
    increment_path, colorstr, print_args
from utils.metrics import ap_per_class, ConfusionMatrix
from utils.plots import output_to_target, plot_images, plot_val_study
from utils.torch_utils import select_device, time_sync


def process_batch(detections, labels, iouv):
    """
    Return correct predictions matrix. Both sets of boxes are in (x1, y1, x2, y2) format.
    Arguments:
        detections (Array[N, 6]), x1, y1, x2, y2, conf, class
        labels (Array[M, 5]), class, x1, y1, x2, y2
    Returns:
        correct (Array[N, 10]), for 10 IoU levels
    """
    correct = torch.zeros(detections.shape[0], iouv.shape[0], dtype=torch.bool, device=iouv.device)
    iou = box_iou(labels[:, 1:], detections[:, :4])
    x = torch.where((iou >= iouv[0]) & (labels[:, 0:1] == detections[:, 5]))  # IoU above threshold and classes match
    if x[0].shape[0]:
        matches = torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1).cpu().numpy()  # [label, detection, iou]
        if x[0].shape[0] > 1:
            matches = matches[matches[:, 2].argsort()[::-1]]
            matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
            # matches = matches[matches[:, 2].argsort()[::-1]]
            matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
        matches = torch.Tensor(matches).to(iouv.device)
        correct[matches[:, 1].long()] = matches[:, 2:3] >= iouv
    return correct

def migraphx_yolov(model, data_tensor):
    # 将输入的tensor数据转换为numpy
    data_numpy=data_tensor.detach().cpu().numpy()
    device = torch.device("cuda")

    # 注意：这里需要执行赋值操作，否则会造成migraphx中输入数据步长不对
    img_data = np.zeros(data_numpy.shape).astype("float32")
    for i in range(data_numpy.shape[0]):
        img_data[i, :, :, :] = data_numpy[i, :, :, :]

    # 执行推理
    result = model.run({"images": img_data})

    # 将结果转换为tensor
    result0=torch.from_numpy(np.array(result[0], copy=False)).to(device)

    return result0

def prepare_input(image):
    input_img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    input_img = cv2.resize(input_img, (640, 640))
    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 run(data,
        weights=None,       # model.pt path(s)
        batch_size=1,       # batch size
        imgsz=640,          # inference size (pixels)
        conf_thres=0.001,    # confidence threshold
        iou_thres=0.65,      # NMS IoU threshold
        device='',          # cuda device, i.e. 0 or 0,1,2,3 or cpu
        single_cls=False,   # treat as single-class dataset
        save_hybrid=False,  # save label+prediction hybrid results to *.txt
        dataloader=None,
        plots=False,
        ):

    resultdir = os.path.join('results', device)
    os.makedirs(resultdir, exist_ok=True)

    # 初始化模型并选择相应的计算设备
    device = select_device(device, batch_size=batch_size)

    if os.path.isfile("/workspace/gaoruiqi_inference/cmcc-code/yolov5-6.0/yolov5m.mxr"):
        model = migraphx.load("/workspace/gaoruiqi_inference/cmcc-code/yolov5-6.0/yolov5m.mxr")
    else:
        # 解析推理模型
        max_input = {"images":[24,3,1024,1024]}
        model = migraphx.parse_onnx(weights, map_input_dims=max_input)

        # 获取模型输入/输出节点信息
        inputs = model.get_inputs()
        outputs = model.get_outputs()

        # 获取模型的输入name
        inputName = model.get_parameter_names()[0]

        # 获取模型的输入尺寸
        inputShape = inputs[inputName].lens()
        inputHeight = int(inputShape[2])
        inputWidth = int(inputShape[3])

        migraphx.quantize_fp16(model)
        
        # 模型编译
        model.compile(t=migraphx.get_target("gpu") ,device_id=0)

    gs = 32
    imgsz = 640

    # Data
    data = check_dataset(data)  # check

    # Configure
    is_coco = isinstance(data.get('val'), str) and data['val'].endswith('coco/val2017.txt')  # COCO dataset
    nc = int(data['nc'])  # number of classes
    iouv = torch.linspace(0.5, 0.95, 10).to(device)  # iou vector for mAP@0.5:0.95
    niou = iouv.numel()

    # Dataloader
    task = 'val'  # path to val images
    dataloader = create_dataloader(data[task], imgsz, batch_size, gs, single_cls, pad=0.5, rect=False,
                                    prefix=colorstr(f'{task}: '))[0]

    seen = 0
    confusion_matrix = ConfusionMatrix(nc=nc)
    class_map = coco80_to_coco91_class() 
    s = ('%20s' + '%11s' * 6) % ('Class', 'Images', 'Labels', 'P', 'R', 'mAP@.5', 'mAP@.5:.95')
    dt, p, r, f1, mp, mr, map50, map = [0.0, 0.0, 0.0], 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
    jdict, stats, ap, ap_class = [], [], [], []

    # 数据预处理
    i = 0
    for batch_i, (img, targets, paths, shapes) in enumerate(tqdm(dataloader, desc=s)):
        img = img.to(device, non_blocking=True)
        img = img.half() 
        img /= 255.0  # 0 - 255 to 0.0 - 1.0

        targets = targets.to(device)
        nb, _, height, width = img.shape  # batch size, channels, height, width

        # Run model
        out = migraphx_yolov(model, img)

        # save to file
        out.cpu().numpy().tofile(f'{resultdir}/{i}_0.bin')
        i += 1

        # Run NMS
        targets[:, 2:] *= torch.Tensor([width, height, width, height]).to(device)  # to pixels
        lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else []  # for autolabelling
        out = non_max_suppression(out, conf_thres, iou_thres, labels=lb, multi_label=True, agnostic=single_cls)

        # Statistics per image
        for si, pred in enumerate(out):
            labels = targets[targets[:, 0] == si, 1:]
            nl = len(labels)
            tcls = labels[:, 0].tolist() if nl else []          # target class
            path, shape = Path(paths[si]), shapes[si][0]
            seen += 1

            if len(pred) == 0:
                if nl:
                    stats.append((torch.zeros(0, niou, dtype=torch.bool), torch.Tensor(), torch.Tensor(), tcls))
                continue

            # Predictions
            if single_cls:
                pred[:, 5] = 0
            predn = pred.clone()
            scale_coords(img[si].shape[1:], predn[:, :4], shape, shapes[si][1])  # native-space pred

            # Evaluate
            if nl:
                tbox = xywh2xyxy(labels[:, 1:5])                             # target boxes
                scale_coords(img[si].shape[1:], tbox, shape, shapes[si][1])  # native-space labels
                labelsn = torch.cat((labels[:, 0:1], tbox), 1)               # native-space labels
                correct = process_batch(predn, labelsn, iouv)
                if plots:
                    confusion_matrix.process_batch(predn, labelsn)
            else:
                correct = torch.zeros(pred.shape[0], niou, dtype=torch.bool)
            stats.append((correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls))  # (correct, conf, pcls, tcls)

    # 计算统计数据
    stats = [np.concatenate(x, 0) for x in zip(*stats)]  # to numpy
    if len(stats) and stats[0].any():
        p, r, ap, f1, ap_class = ap_per_class(*stats, plot=False)
        ap50, ap = ap[:, 0], ap.mean(1)  # AP@0.5, AP@0.5:0.95
        mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
        nt = np.bincount(stats[3].astype(np.int64), minlength=nc)  # number of targets per class
    else:
        nt = torch.zeros(1)

    # Print results
    print('map50:', map50)
    print('map50-95:', map)

def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
    parser.add_argument('--weights', type=str, default='', help='model.onnx path(s)')
    parser.add_argument('--batch-size', type=int, default=32, help='batch size')
    parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='inference size (pixels)')
    parser.add_argument('--conf-thres', type=float, default=0.001, help='confidence threshold')
    parser.add_argument('--iou-thres', type=float, default=0.6, help='NMS IoU threshold')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    opt = parser.parse_args()
    opt.data = check_yaml(opt.data)  # check YAML
    print_args(FILE.stem, opt)
    return opt

def main(opt):
    # 检测requirements文件中需要的包是否安装好了
    check_requirements(exclude=('tensorboard', 'thop'))
    run(**vars(opt))

if __name__ == "__main__":
    opt = parse_opt()
    main(opt)