# Gradio YOLOv5 Det Blocks 01_03
# 创建人：曾逸夫
# 创建时间：2022-06-04
# 功能描述：对比检测，同一图片 相同模型 不同配置

from util.gradio_version_opt import gr_v_opt

gr_v_opt()

import argparse
import csv
import gc
import json
import sys
from collections import Counter
from pathlib import Path

import cv2
import gradio as gr
import numpy as np
import pandas as pd
import plotly.express as px
import torch
import yaml
from PIL import Image, ImageDraw, ImageFont

ROOT_PATH = sys.path[0]  # 根目录

# yolov5路径
yolov5_path = "ultralytics/yolov5"

# 本地模型路径
local_model_path = f"{ROOT_PATH}/models"

# Gradio YOLOv5 Det版本
GYD_VERSION = "Gradio YOLOv5 Det block 01_03"

# 模型名称临时变量
model_name_tmp = ""

# 设备临时变量
device_tmp = ""

# 文件后缀
suffix_list = [".csv", ".yaml"]

# 字体大小
FONTSIZE = 25

# 目标尺寸
obj_style = ["小目标", "中目标", "大目标"]


def parse_args(known=False):
    parser = argparse.ArgumentParser(description="Gradio YOLOv5 Det block 01_03")
    parser.add_argument(
        "--model_cfg",
        "-mc5",
        default="./model_config/model_name_p5_p6_all.yaml",
        type=str,
        help="model config",
    )
    parser.add_argument(
        "--nms_conf",
        "-conf",
        default=0.5,
        type=float,
        help="model NMS confidence threshold",
    )
    parser.add_argument("--nms_iou", "-iou", default=0.45, type=float, help="model NMS IoU threshold")
    parser.add_argument("--inference_size", "-isz", default=640, type=int, help="model inference size")

    args = parser.parse_known_args()[0] if known else parser.parse_args()
    return args


# yaml文件解析
def yaml_parse(file_path):
    return yaml.safe_load(open(file_path, encoding="utf-8").read())


# yaml csv 文件解析
def yaml_csv(file_path, file_tag):
    file_suffix = Path(file_path).suffix
    if file_suffix == suffix_list[0]:
        # 模型名称
        file_names = [i[0] for i in list(csv.reader(open(file_path)))]  # csv版
    elif file_suffix == suffix_list[1]:
        # 模型名称
        file_names = yaml_parse(file_path).get(file_tag)  # yaml版
    else:
        print(f"{file_path}格式不正确！程序退出！")
        sys.exit()

    return file_names


#  模型加载
def model_loading(model_name):

    # 加载本地模型
    try:
        torch.hub._validate_not_a_forked_repo = lambda a, b, c: True
        model = torch.hub.load(
            yolov5_path,
            "custom",
            path=f"{local_model_path}/{model_name}",
            device="cuda:0",
            force_reload=False,
            _verbose=True,
        )
    except Exception as e:
        print("模型加载失败！")
        print(e)
        return False
    else:
        print(f"🚀 欢迎使用{GYD_VERSION}，{model_name}加载成功！")

    return model


# YOLOv5图片检测函数
def yolo_det(img, model_name, infer_size, conf, iou):

    global model, model_name_tmp

    if model_name_tmp != model_name:
        # 模型判断，避免反复加载
        model_name_tmp = model_name
        print(f"正在加载模型{model_name_tmp}......")
        model = model_loading(model_name_tmp)
    else:
        print(f"正在加载模型{model_name_tmp}......")
        model = model_loading(model_name_tmp)

    # -----------模型调参-----------
    model.conf = conf  # NMS 置信度阈值
    model.iou = iou  # NMS IOU阈值
    model.max_det = 1000  # 最大检测框数

    results = model(img, size=infer_size)  # 检测
    results.render()  # 渲染

    det_img = Image.fromarray(results.ims[0])  # 检测图片

    return det_img


def main(args):
    gr.close_all()

    slider_step = 0.05  # 滑动步长

    nms_conf = args.nms_conf
    nms_iou = args.nms_iou
    model_cfg = args.model_cfg
    inference_size = args.inference_size

    # 模型加载
    model_names = yaml_csv(model_cfg, "model_names")

    with gr.Blocks() as gyd:

        with gr.Group():
            with gr.Row():
                inputs_img = gr.Image(image_mode="RGB", source="upload", type="pil", label="原始图片")
            with gr.Row():
                inputs_model = gr.inputs.Dropdown(choices=model_names, default="yolov5s", type="value", label="P5模型")
            with gr.Row():
                with gr.Column():
                    inputs_size_p5 = gr.Radio(choices=[320, 640, 1280], value=inference_size, label="推理尺寸")
                    input_conf_p5 = gr.inputs.Slider(0, 1, step=slider_step, default=nms_conf, label="置信度阈值")
                    inputs_iou_p5 = gr.inputs.Slider(0, 1, step=slider_step, default=nms_iou, label="IoU 阈值")
                    det_btn_p5 = gr.Button(value='Detect 01')

                with gr.Column():
                    inputs_size_p6 = gr.Radio(choices=[320, 640, 1280], value=inference_size, label="推理尺寸")
                    input_conf_p6 = gr.inputs.Slider(0, 1, step=slider_step, default=nms_conf, label="置信度阈值")
                    inputs_iou_p6 = gr.inputs.Slider(0, 1, step=slider_step, default=nms_iou, label="IoU 阈值")
                    det_btn_p6 = gr.Button(value='Detect 02')

        with gr.Group():
            with gr.Row():
                with gr.Column():
                    outputs_img_p5 = gr.Image(type="pil", label="检测图片")
                with gr.Column():
                    outputs_img_p6 = gr.Image(type="pil", label="检测图片")

        det_btn_p5.click(fn=yolo_det,
                         inputs=[inputs_img, inputs_model, inputs_size_p5, input_conf_p5, inputs_iou_p5],
                         outputs=[
                             outputs_img_p5,])

        det_btn_p6.click(fn=yolo_det,
                         inputs=[inputs_img, inputs_model, inputs_size_p6, input_conf_p6, inputs_iou_p6],
                         outputs=[
                             outputs_img_p6,])

    gyd.launch(inbrowser=True)


if __name__ == '__main__':
    args = parse_args()
    main(args)