# Gradio YOLOv5 Det Blocks 08_02
# 创建人：曾逸夫
# 创建时间：2022-07-11
# 功能描述：目标尺寸，多选

from util.gradio_version_opt import gr_v_opt

gr_v_opt()

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

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

from util.fonts_opt import is_fonts

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 08_02"

# 模型名称临时变量
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 08_02")
    parser.add_argument(
        "--model_cfg_p5",
        "-mc5",
        default="./model_config/model_name_p5_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 color_set(cls_num):
    color_list = []
    for i in range(cls_num):
        color = tuple(np.random.choice(range(256), size=3))
        # color = ["#"+''.join([random.choice('0123456789ABCDEF') for j in range(6)])]
        color_list.append(color)

    return color_list


# plt转pil
def fig2img(fig):
    buf = io.BytesIO()
    fig.savefig(buf, bbox_inches="tight", dpi=100)
    buf.seek(0)
    img = Image.open(buf)
    return img


# matplotlib绘制
def plt_draw(img, img_size, score_l, bbox_l, cls_l, color_list):
    img_dpi = 100
    plt.figure(figsize=[img_size[0] / img_dpi, img_size[1] / img_dpi], dpi=img_dpi, frameon=False)
    plt.imshow(img)
    ax = plt.gca()
    for score, (xmin, ymin, xmax, ymax), label, color in zip(score_l, bbox_l, cls_l, color_list):
        _color = [color[i] / 255 for i in range(len(color))]
        ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, fill=False, color=_color, linewidth=2))
        ax.text(xmin + 5,
                ymin - 10,
                f"{label}: {score:0.2f}",
                fontsize=15,
                color=(1, 1, 1),
                bbox=dict(facecolor=_color, alpha=0.5))
    plt.axis("off")
    return fig2img(plt.gcf())


# 绘制（pillow版）
def pil_draw(img, obj_Size, textFont, color_list):

    img_pil = ImageDraw.Draw(img)
    id = 0

    for i in obj_Size:
        img_pil.rectangle(i[1], fill=None, outline=color_list[i[3]], width=2)  # 边界框
        countdown_msg = f"{id}-{i[2]} {i[0]:.2f}"
        text_w, text_h = textFont.getsize(countdown_msg)  # 标签尺寸

        # 标签背景
        img_pil.rectangle(
            (i[1][0], i[1][1], i[1][0] + text_w, i[1][1] + text_h),
            fill=color_list[i[3]],
            outline=color_list[i[3]],
        )

        # 标签
        img_pil.multiline_text(
            (i[1][0], i[1][1]),
            countdown_msg,
            fill=(255, 255, 255),
            font=textFont,
            align="center",
        )

        id += 1

    return img


#  模型加载
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, obj_size):

    global model, model_name_tmp

    s_objSize = []
    m_objSize = []
    l_objSize = []
    all_objSize = []

    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  # 最大检测框数
    img_size = img.size  # 帧尺寸

    results = model(img, size=infer_size)  # 检测

    model_cls_name = yaml_csv("./cls_name/cls_name_en.yaml", "model_cls_name")  # 类别名称

    color_list = color_set(len(model_cls_name))  # 设置颜色
    textFont = ImageFont.truetype("./fonts/TimesNewRoman.ttf", FONTSIZE)

    for result in results.xyxyn:
        for i in range(len(result)):
            # id = int(i)  # 实例ID
            obj_cls_index = int(result[i][5])  # 类别索引

            obj_cls = model_cls_name[obj_cls_index]  # 类别

            # ------------边框坐标------------
            x0 = float(result[i][:4].tolist()[0])
            y0 = float(result[i][:4].tolist()[1])
            x1 = float(result[i][:4].tolist()[2])
            y1 = float(result[i][:4].tolist()[3])

            # ------------边框实际坐标------------
            x0 = int(img_size[0] * x0)
            y0 = int(img_size[1] * y0)
            x1 = int(img_size[0] * x1)
            y1 = int(img_size[1] * y1)
            bbox_area = (x1 - x0) * (y1 - y0)

            conf = float(result[i][4])  # 置信度

            all_objSize.append([conf, (x0, y0, x1, y1), obj_cls, obj_cls_index])

            if 0 < bbox_area <= 32 ** 2:
                s_objSize.append([conf, (x0, y0, x1, y1), obj_cls, obj_cls_index])
            elif 32 ** 2 < bbox_area < 96 ** 2:
                m_objSize.append([conf, (x0, y0, x1, y1), obj_cls, obj_cls_index])
            elif bbox_area >= 96 ** 2:
                l_objSize.append([conf, (x0, y0, x1, y1), obj_cls, obj_cls_index])

    det_sImg, det_mImg, det_lImg = None, None, None
    img_s, img_m, img_l = img.copy(), img.copy(), img.copy()

    if "小目标" in obj_size:
        det_sImg = pil_draw(
            img_s,
            s_objSize,
            textFont,
            color_list,
        )
    if "中目标" in obj_size:
        det_mImg = pil_draw(
            img_m,
            m_objSize,
            textFont,
            color_list,
        )
    if "大目标" in obj_size:
        det_lImg = pil_draw(
            img_l,
            l_objSize,
            textFont,
            color_list,
        )

    return det_sImg, det_mImg, det_lImg


# 检测图片放大
def det_img_expandShow(img):
    return img


def main(args):
    gr.close_all()

    slider_step = 0.05  # 滑动步长

    nms_conf = args.nms_conf
    nms_iou = args.nms_iou
    model_cfg_p5 = args.model_cfg_p5
    inference_size = args.inference_size

    is_fonts(f"{ROOT_PATH}/fonts")  # 检查字体文件

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

    with gr.Blocks() as gyd:

        with gr.Box():
            with gr.Row():
                gr.Markdown("### P5检测")

            with gr.Row():
                with gr.Column():
                    with gr.Row():
                        inputs_img_p5 = gr.Image(image_mode="RGB", source="upload", type="pil", label="原始图片")
                    with gr.Row():
                        inputs_model_p5 = gr.Radio(choices=model_names_p5, value="yolov5s", label="P5模型")
                    with gr.Row():
                        inputs_size_p5 = gr.Radio(choices=[320, 640, 1280], value=inference_size, label="推理尺寸")
                    with gr.Row():
                        input_conf_p5 = gr.inputs.Slider(0, 1, step=slider_step, default=nms_conf, label="置信度阈值")
                    with gr.Row():
                        inputs_iou_p5 = gr.inputs.Slider(0, 1, step=slider_step, default=nms_iou, label="IoU 阈值")
                    with gr.Row():
                        obj_size = gr.CheckboxGroup(choices=["小目标", "中目标", "大目标"],
                                                    value=["小目标", "中目标", "大目标"],
                                                    type="value",
                                                    label="目标尺寸")
                    with gr.Row():
                        det_btn_01 = gr.Button(value='Detect 01', variant="primary")

                with gr.Column():
                    with gr.Row():
                        outputs_sImg_p5 = gr.Image(type="pil", label="小目标")
                        outputs_mImg_p5 = gr.Image(type="pil", label="中目标")
                        outputs_lImg_p5 = gr.Image(type="pil", label="大目标")

                    with gr.Row():
                        det_btn_s = gr.Button(value='Show Small')
                        det_btn_m = gr.Button(value='Show Medium')
                        det_btn_l = gr.Button(value='Show Large')

                    with gr.Row():
                        outputs_img_p5 = gr.Image(image_mode="RGB", source="upload", type="pil", label="检测图片")

        det_btn_01.click(fn=yolo_det,
                         inputs=[
                             inputs_img_p5, inputs_model_p5, inputs_size_p5, input_conf_p5, inputs_iou_p5, obj_size],
                         outputs=[outputs_sImg_p5, outputs_mImg_p5, outputs_lImg_p5])

        det_btn_s.click(fn=det_img_expandShow, inputs=[outputs_sImg_p5], outputs=[outputs_img_p5])
        det_btn_m.click(fn=det_img_expandShow, inputs=[outputs_mImg_p5], outputs=[outputs_img_p5])
        det_btn_l.click(fn=det_img_expandShow, inputs=[outputs_lImg_p5], outputs=[outputs_img_p5])

    gyd.launch(inbrowser=True)


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