gradio_yolov5_det.ipynb · 空林长风/Gradio-YOLOv5-Det

加入 Gitee

与超过 1200万开发者一起发现、参与优秀开源项目，私有仓库也完全免费：）

克隆/下载

gradio_yolov5_det.ipynb 13.30 KB

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "627529e5-9204-4989-855b-fc7193ec17fa",
   "metadata": {},
   "source": [
    "# 🚀 Gradio YOLOv5 Det v0.1 （Jupyter版）\n",
    "## 创建人：曾逸夫\n",
    "## 创建时间：2022-04-08"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1f7190d-4732-44a3-a3b1-cf002dadf007",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install ipywidgets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "489bd085-18d6-4246-8f1d-dcd9f80aa8ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.system(\"pip install gradio==2.9.4\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cdd5e57-6580-415a-b92d-8bd8e104b495",
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import csv\n",
    "import os\n",
    "import sys\n",
    "from pathlib import Path\n",
    "\n",
    "import gradio as gr\n",
    "import torch\n",
    "import yaml\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46ebffa5-d161-4619-a0ad-70b6631c6670",
   "metadata": {},
   "outputs": [],
   "source": [
    "ROOT_PATH = sys.path[0]  # 根目录\n",
    "\n",
    "# 本地模型路径\n",
    "local_model_path = f\"{ROOT_PATH}/yolov5\"\n",
    "\n",
    "\n",
    "# 模型名称临时变量\n",
    "model_name_tmp = \"\"\n",
    "\n",
    "# 设备临时变量\n",
    "device_tmp = \"\"\n",
    "\n",
    "# 文件后缀\n",
    "suffix_list = [\".csv\", \".yaml\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a5c5d0a-7533-44bc-9a44-a01c3a4902cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_args(known=False):\n",
    "    parser = argparse.ArgumentParser(description=\"Gradio YOLOv5 Det v0.1\")\n",
    "    parser.add_argument(\n",
    "        \"--img_dir\", \"-imd\", default=\"img_example\", type=str, help=\"image dir\"\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--model_name\", \"-mn\", default=\"yolov5s\", type=str, help=\"model name\"\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--model_cfg\",\n",
    "        \"-mc\",\n",
    "        default=\"./model_config/model_name_p5_all.yaml\",\n",
    "        type=str,\n",
    "        help=\"model config\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--cls_name\",\n",
    "        \"-cls\",\n",
    "        default=\"./cls_name/cls_name.yaml\",\n",
    "        type=str,\n",
    "        help=\"cls name\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--nms_conf\",\n",
    "        \"-conf\",\n",
    "        default=0.5,\n",
    "        type=float,\n",
    "        help=\"model NMS confidence threshold\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--nms_iou\", \"-iou\", default=0.45, type=float, help=\"model NMS IoU threshold\"\n",
    "    )\n",
    "\n",
    "    parser.add_argument(\n",
    "        \"--label_dnt_show\",\n",
    "        \"-lds\",\n",
    "        action=\"store_false\",\n",
    "        default=True,\n",
    "        help=\"label show\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--device\",\n",
    "        \"-dev\",\n",
    "        default=\"cpu\",\n",
    "        type=str,\n",
    "        help=\"cuda or cpu\",\n",
    "    )\n",
    "    parser.add_argument(\n",
    "        \"--inference_size\", \"-isz\", default=640, type=int, help=\"model inference size\"\n",
    "    )\n",
    "\n",
    "    args = parser.parse_known_args()[0] if known else parser.parse_args(args=[])\n",
    "    return args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66add09b-0a47-45dc-8145-bb1466aca435",
   "metadata": {},
   "outputs": [],
   "source": [
    "#  模型加载\n",
    "def model_loading(model_name, device):\n",
    "\n",
    "    # 加载本地模型\n",
    "    model = torch.hub.load(\n",
    "        local_model_path,\n",
    "        \"custom\",\n",
    "        path=f\"{local_model_path}/{model_name}\",\n",
    "        source=\"local\",\n",
    "        device=device,\n",
    "    )\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6523632-927f-43a1-b656-2094662a4f4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 检测信息\n",
    "def export_json(results, model, img_size):\n",
    "\n",
    "    return [\n",
    "        [\n",
    "            {\n",
    "                \"id\": i,\n",
    "                \"class\": int(result[i][5]),\n",
    "                \"class_name\": model.model.names[int(result[i][5])],\n",
    "                # \"class_name\": model_cls_name_cp[int(result[i][5])],\n",
    "                \"normalized_box\": {\n",
    "                    \"x0\": round(result[i][:4].tolist()[0], 6),\n",
    "                    \"y0\": round(result[i][:4].tolist()[1], 6),\n",
    "                    \"x1\": round(result[i][:4].tolist()[2], 6),\n",
    "                    \"y1\": round(result[i][:4].tolist()[3], 6),\n",
    "                },\n",
    "                \"confidence\": round(float(result[i][4]), 2),\n",
    "                \"fps\": round(1000 / float(results.t[1]), 2),\n",
    "                \"width\": img_size[0],\n",
    "                \"height\": img_size[1],\n",
    "            }\n",
    "            for i in range(len(result))\n",
    "        ]\n",
    "        for result in results.xyxyn\n",
    "    ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b9078a7-c36f-4b03-b1ff-596fdb457681",
   "metadata": {},
   "outputs": [],
   "source": [
    "# YOLOv5图片检测函数\n",
    "def yolo_det(img, device, model_name, inference_size, conf, iou, label_opt, model_cls):\n",
    "\n",
    "    global model, model_name_tmp, device_tmp\n",
    "\n",
    "    if model_name_tmp != model_name:\n",
    "        # 模型判断，避免反复加载\n",
    "        model_name_tmp = model_name\n",
    "        model = model_loading(model_name_tmp, device)\n",
    "    elif device_tmp != device:\n",
    "        device_tmp = device\n",
    "        model = model_loading(model_name_tmp, device)\n",
    "\n",
    "    # -----------模型调参-----------\n",
    "    model.conf = conf  # NMS 置信度阈值\n",
    "    model.iou = iou  # NMS IOU阈值\n",
    "    model.max_det = 1000  # 最大检测框数\n",
    "    model.classes = model_cls  # 模型类别\n",
    "\n",
    "    results = model(img, size=inference_size)  # 检测\n",
    "    results.render(labels=label_opt)  # 渲染\n",
    "\n",
    "    det_img = Image.fromarray(results.imgs[0])  # 检测图片\n",
    "\n",
    "    det_json = export_json(results, model, img.size)[0]  # 检测信息\n",
    "\n",
    "    return det_img, det_json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77be530a-35c0-4a57-8b58-25fc70862c67",
   "metadata": {},
   "outputs": [],
   "source": [
    "# yaml文件解析\n",
    "def yaml_parse(file_path):\n",
    "    return yaml.safe_load(open(file_path, \"r\", encoding=\"utf-8\").read())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cecba5b-f9d1-4031-b3ba-7e8da23217f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# yaml csv 文件解析\n",
    "def yaml_csv(file_path, file_tag):\n",
    "    file_suffix = Path(file_path).suffix\n",
    "    if file_suffix == suffix_list[0]:\n",
    "        # 模型名称\n",
    "        file_names = [i[0] for i in list(csv.reader(open(file_path)))]  # csv版\n",
    "    elif file_suffix == suffix_list[1]:\n",
    "        # 模型名称\n",
    "        file_names = yaml_parse(file_path).get(file_tag)  # yaml版\n",
    "    else:\n",
    "        print(f\"{file_path}格式不正确！程序退出！\")\n",
    "        sys.exit()\n",
    "\n",
    "    return file_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "502b4d1c-d628-4d57-bfa7-1ff94ba835a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def main(args):\n",
    "    gr.close_all()\n",
    "\n",
    "    global model\n",
    "\n",
    "    slider_step = 0.05  # 滑动步长\n",
    "\n",
    "    img_dir = args.img_dir\n",
    "    nms_conf = args.nms_conf\n",
    "    nms_iou = args.nms_iou\n",
    "    label_opt = args.label_dnt_show\n",
    "    model_name = args.model_name\n",
    "    model_cfg = args.model_cfg\n",
    "    cls_name = args.cls_name\n",
    "    device = args.device\n",
    "    inference_size = args.inference_size\n",
    "\n",
    "    # 模型加载\n",
    "    model = model_loading(model_name, device)\n",
    "\n",
    "    model_names = yaml_csv(model_cfg, \"model_names\")\n",
    "    model_cls_name = yaml_csv(cls_name, \"model_cls_name\")\n",
    "\n",
    "    # model_cls_name_cp = model_cls_name.copy()  # 类别名称\n",
    "\n",
    "    # -------------------输入组件-------------------\n",
    "    inputs_img = gr.inputs.Image(type=\"pil\", label=\"原始图片\")\n",
    "    inputs_device = gr.inputs.Dropdown(\n",
    "        choices=[\"0\", \"cpu\"], default=device, type=\"value\", label=\"设备\"\n",
    "    )\n",
    "    inputs_model = gr.inputs.Dropdown(\n",
    "        choices=model_names, default=model_name, type=\"value\", label=\"模型\"\n",
    "    )\n",
    "    inputs_size = gr.inputs.Radio(\n",
    "        choices=[320, 640], default=inference_size, label=\"推理尺寸\"\n",
    "    )\n",
    "    input_conf = gr.inputs.Slider(\n",
    "        0, 1, step=slider_step, default=nms_conf, label=\"置信度阈值\"\n",
    "    )\n",
    "    inputs_iou = gr.inputs.Slider(\n",
    "        0, 1, step=slider_step, default=nms_iou, label=\"IoU 阈值\"\n",
    "    )\n",
    "    inputs_label = gr.inputs.Checkbox(default=label_opt, label=\"标签显示\")\n",
    "    inputs_clsName = gr.inputs.CheckboxGroup(\n",
    "        choices=model_cls_name, default=model_cls_name, type=\"index\", label=\"类别\"\n",
    "    )\n",
    "\n",
    "    # 输入参数\n",
    "    inputs = [\n",
    "        inputs_img,  # 输入图片\n",
    "        inputs_device,  # 设备\n",
    "        inputs_model,  # 模型\n",
    "        inputs_size,  # 推理尺寸\n",
    "        input_conf,  # 置信度阈值\n",
    "        inputs_iou,  # IoU阈值\n",
    "        inputs_label,  # 标签显示\n",
    "        inputs_clsName,  # 类别\n",
    "    ]\n",
    "    # 输出参数\n",
    "    outputs = gr.outputs.Image(type=\"pil\", label=\"检测图片\")\n",
    "    outputs02 = gr.outputs.JSON(label=\"检测信息\")\n",
    "\n",
    "    # 标题\n",
    "    title = \"基于Gradio的YOLOv5通用目标检测系统\"\n",
    "    # 描述\n",
    "    description = \"<div align='center'>可自定义目标检测模型、安装简单、使用方便</div>\"\n",
    "\n",
    "    # 示例图片\n",
    "    examples = [\n",
    "        [\n",
    "            \"./img_example/bus.jpg\",\n",
    "            \"cpu\",\n",
    "            \"yolov5s\",\n",
    "            640,\n",
    "            0.6,\n",
    "            0.5,\n",
    "            True,\n",
    "            [\"人\", \"公交车\"],\n",
    "        ],\n",
    "        [\n",
    "            \"./img_example/Millenial-at-work.jpg\",\n",
    "            \"0\",\n",
    "            \"yolov5l\",\n",
    "            320,\n",
    "            0.5,\n",
    "            0.45,\n",
    "            True,\n",
    "            [\"人\", \"椅子\", \"杯子\", \"笔记本电脑\"],\n",
    "        ],\n",
    "        [\n",
    "            \"./img_example/zidane.jpg\",\n",
    "            \"0\",\n",
    "            \"yolov5m\",\n",
    "            640,\n",
    "            0.25,\n",
    "            0.5,\n",
    "            False,\n",
    "            [\"人\", \"领带\"],\n",
    "        ],\n",
    "    ]\n",
    "\n",
    "    # 接口\n",
    "    gr.Interface(\n",
    "        fn=yolo_det,\n",
    "        inputs=inputs,\n",
    "        outputs=[outputs, outputs02],\n",
    "        title=title,\n",
    "        description=description,\n",
    "        examples=examples,\n",
    "        theme=\"seafoam\",\n",
    "        # live=True, # 实时变更输出\n",
    "        flagging_dir=\"run\"  # 输出目录\n",
    "        # ).launch(inbrowser=True, auth=['admin', 'admin'])\n",
    "    ).launch(\n",
    "        inbrowser=True,  # 自动打开默认浏览器\n",
    "        show_tips=True,  # 自动显示gradio最新功能\n",
    "        # favicon_path=\"./icon/logo.ico\",\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b976790e-77ef-4332-b2a4-1d2d5bcf2cc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    args = parse_args()\n",
    "    main(args)"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "f9d149eed12fac937545f4635efc9e9926f24a5acb73f9a6fbfbd1c202313110"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}

一键复制编辑原始数据按行查看历史

提交于 2022-05-12 09:41 . v0.2.2 jupyter update

马建仓 AI 助手

尝试更多

代码解读

代码找茬

代码优化

Python

https://gitee.com/jiangxiaolan/gradio_yolov5_det.git

git@gitee.com:jiangxiaolan/gradio_yolov5_det.git

jiangxiaolan

gradio_yolov5_det

Gradio-YOLOv5-Det

master

空林长风 / Gradio-YOLOv5-Det

简介

发行版

贡献者

近期动态

空林长风 / Gradio-YOLOv5-Det .gitee-modal { width: 500px !important; }

简介

发行版

贡献者

近期动态

搜索帮助

空林长风 / Gradio-YOLOv5-Det