# 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
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval

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 AllocateOutputMemory(model):
    outputData={}
    for key in model.get_outputs().keys():
        outputData[key] = migraphx.allocate_gpu(s=model.get_outputs()[key])
    return outputData


def xyxy2xywh(x):
    # Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right
    y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
    y[:, 0] = (x[:, 0] + x[:, 2]) / 2  # x center
    y[:, 1] = (x[:, 1] + x[:, 3]) / 2  # y center
    y[:, 2] = x[:, 2] - x[:, 0]  # width
    y[:, 3] = x[:, 3] - x[:, 1]  # height
    return y


def save_coco_json(predn, pred_dict, image_id, class_map):
    # Save one JSON result {"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}
    box = xyxy2xywh(predn[:, :4])  # xywh
    box[:, :2] -= box[:, 2:] / 2  # xy center to top-left corner
    for p, b in zip(predn.tolist(), box.tolist()):
        pred_dict.append({'image_id': image_id,
                          'category_id': class_map[int(p[5])],
                          'bbox': [round(x, 3) for x in b],
                          'score': round(p[4], 5)})

def evaluate(cocoGt_file, cocoDt_file):
    cocoGt = COCO(cocoGt_file)
    cocoDt = cocoGt.loadRes(cocoDt_file)
    cocoEval = COCOeval(cocoGt, cocoDt, 'bbox')
    cocoEval.evaluate()
    cocoEval.accumulate()
    cocoEval.summarize()


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 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,
        ground_truth_json='',
        ):
    this_file_device = device
    resultdir = os.path.join('results', device)
    os.makedirs(resultdir, exist_ok=True)

    # 初始化模型并选择相应的计算设备
    device = select_device(device, batch_size=batch_size)
    if weights.split(".")[-1] == "mxr":
        model = migraphx.load(weights)
        inputName=list(model.get_inputs().keys())[0]
    elif weights.split(".")[-1] == "onnx":
        # 解析推理模型
        max_input = {"images":[24,3,640,640]}
        model = migraphx.parse_onnx(weights, map_input_dims=max_input)
        inputName = model.get_parameter_names()[0]
        migraphx.quantize_fp16(model)
        model.compile(t=migraphx.get_target("gpu"), offlod_copy=False, device_id=0)
    else:
        print("请输出正确的模型路径")
    
    modelData=AllocateOutputMemory(model)

    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 = [], [], [], []

    # 数据预处理
    pred_results = []
    i = 0
    infer_times = []        
    total_infer_times = []
    total_start = time.time()
    all_images = []

    for img, _, _, _ in dataloader:
        img = img.float() / 255.0
        img_np=img.numpy()
        all_images.append(img_np.astype(np.float32))

    # warm up   
    modelData[inputName] = migraphx.to_gpu(migraphx.argument(all_images[0]))
    model.run(modelData)
    
    for batch_i, (_, targets, paths, shapes) in enumerate(tqdm(dataloader, desc=s)):
        img = all_images[batch_i]
        modelData[inputName] = migraphx.to_gpu(migraphx.argument(img))

        targets = targets.to(device)
        nb, _, height, width = img.shape  # batch size, channels, height, width
        if nb != 24:
            break
        start = time.time()
        # Run model
        out = model.run(modelData)
        infer_times.append(time.time() - start)
        total_infer_times.append(time.time() - total_start)
        # save to file
        out=np.array(migraphx.from_gpu(out[0]))
        out.tofile(f'{resultdir}/{i}_0.bin')
        out=torch.from_numpy(out).to("cuda")
        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 idx, pred in enumerate(out):
            try:
                scale_coords((640, 640), pred[:, :4], shapes[idx][0][:])

            except:
                pred = torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0]])
            # append to COCO-JSON dictionary
            path = Path(paths[idx])
            image_id = int(path.stem) if path.stem.isnumeric() else path.stem
            save_coco_json(pred, pred_results, image_id, coco80_to_coco91_class())
        total_start = time.time()

    pred_json_file = f"./results/yolov5m_predictions{this_file_device}.json"

    with open(pred_json_file, 'w') as f:
        json.dump(pred_results, f)
    print(f"saving results to {pred_json_file}")

    # evaluate mAP
    evaluate(ground_truth_json, pred_json_file)
    
    
    print("***************************")
    infer_time = sum(infer_times)
    avg_infer_fps = 24 * len(infer_times) / sum(infer_times)
    print(f"total_infer_time: {infer_time}s")
    print(f'avg_infer_fps: {avg_infer_fps}samples/s')
    load_data_infer_time = sum(total_infer_times)
    load_data_avg_infer_fps = len(total_infer_times) * 24 / sum(total_infer_times)
    print(f'load_data_total_infer_time: {load_data_infer_time}s')
    print(f'load_data_avg_total_Infer_fps: {load_data_avg_infer_fps} samples/s')
    print("******************************")
    
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')
    parser.add_argument('--ground_truth_json', type=str, default="/datasets/coco/instances_val2017.json", help='annotation file path')
    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)