update

README.md

@@ -1,13 +1,15 @@
 ---
-title: LIVE
-emoji: 🐨
-colorFrom: green
-colorTo: green
+title: Gradio_yolov5_det
+emoji: 📊
+colorFrom: pink
+colorTo: indigo
 sdk: gradio
-sdk_version: 2.9.4
+sdk_version: 2.9.1
 app_file: app.py
 pinned: false
-license: apache-2.0
+license: gpl-3.0
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces#reference
+
+🚀 项目主页：https://gitee.com/CV_Lab/gradio_yolov5_det

__init__.py

@@ -0,0 +1,2 @@
+__author__ = "Xu Ma"
+__email__ = "[email protected]"

app.py

@@ -1,119 +1,289 @@
+# Gradio YOLOv5 Det v0.1
+# 创建人：曾逸夫
+# 创建时间：2022-04-03
+# email：[email protected]
+# 项目主页：https://gitee.com/CV_Lab/gradio_yolov5_det
+
+import argparse
+import csv
+import sys
+from pathlib import Path
+
 import gradio as gr
-from transformers import DPTFeatureExtractor, DPTForDepthEstimation
 import torch
-import numpy as np
+import yaml
 from PIL import Image
-import open3d as o3d
-from pathlib import Path
-import os
-
-feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
-model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large")
-
-
-def process_image(image_path):
-    image_path = Path(image_path)
-    image_raw = Image.open(image_path)
-    image = image_raw.resize(
-        (800, int(800 * image_raw.size[1] / image_raw.size[0])),
-        Image.Resampling.LANCZOS)
-
-    # prepare image for the model
-    encoding = feature_extractor(image, return_tensors="pt")
-
-    # forward pass
-    with torch.no_grad():
-        outputs = model(**encoding)
-        predicted_depth = outputs.predicted_depth
-
-    # interpolate to original size
-    prediction = torch.nn.functional.interpolate(
-        predicted_depth.unsqueeze(1),
-        size=image.size[::-1],
-        mode="bicubic",
-        align_corners=False,
-    ).squeeze()
-    output = prediction.cpu().numpy()
-    depth_image = (output * 255 / np.max(output)).astype('uint8')
-    try:
-        gltf_path = create_3d_obj(np.array(image), depth_image, image_path)
-        img = Image.fromarray(depth_image)
-        return [img, gltf_path, gltf_path]
-    except Exception as e:
-        gltf_path = create_3d_obj(
-            np.array(image), depth_image, image_path, depth=8)
-        img = Image.fromarray(depth_image)
-        return [img, gltf_path, gltf_path]
-    except:
-        print("Error reconstructing 3D model")
-        raise Exception("Error reconstructing 3D model")
-
-
-def create_3d_obj(rgb_image, depth_image, image_path, depth=10):
-    depth_o3d = o3d.geometry.Image(depth_image)
-    image_o3d = o3d.geometry.Image(rgb_image)
-    rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(
-        image_o3d, depth_o3d, convert_rgb_to_intensity=False)
-    w = int(depth_image.shape[1])
-    h = int(depth_image.shape[0])
-
-    camera_intrinsic = o3d.camera.PinholeCameraIntrinsic()
-    camera_intrinsic.set_intrinsics(w, h, 500, 500, w/2, h/2)
-
-    pcd = o3d.geometry.PointCloud.create_from_rgbd_image(
-        rgbd_image, camera_intrinsic)
-
-    print('normals')
-    pcd.normals = o3d.utility.Vector3dVector(
-        np.zeros((1, 3)))  # invalidate existing normals
-    pcd.estimate_normals(
-        search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.01, max_nn=30))
-    pcd.orient_normals_towards_camera_location(
-        camera_location=np.array([0., 0., 1000.]))
-    pcd.transform([[1, 0, 0, 0],
-                   [0, -1, 0, 0],
-                   [0, 0, -1, 0],
-                   [0, 0, 0, 1]])
-    pcd.transform([[-1, 0, 0, 0],
-                   [0, 1, 0, 0],
-                   [0, 0, 1, 0],
-                   [0, 0, 0, 1]])
-
-    print('run Poisson surface reconstruction')
-    with o3d.utility.VerbosityContextManager(o3d.utility.VerbosityLevel.Debug) as cm:
-        mesh_raw, densities = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(
-            pcd, depth=depth, width=0, scale=1.1, linear_fit=True)
-
-    voxel_size = max(mesh_raw.get_max_bound() - mesh_raw.get_min_bound()) / 256
-    print(f'voxel_size = {voxel_size:e}')
-    mesh = mesh_raw.simplify_vertex_clustering(
-        voxel_size=voxel_size,
-        contraction=o3d.geometry.SimplificationContraction.Average)
-
-    # vertices_to_remove = densities < np.quantile(densities, 0.001)
-    # mesh.remove_vertices_by_mask(vertices_to_remove)
-    bbox = pcd.get_axis_aligned_bounding_box()
-    mesh_crop = mesh.crop(bbox)
-    gltf_path = f'./{image_path.stem}.gltf'
-    o3d.io.write_triangle_mesh(
-        gltf_path, mesh_crop, write_triangle_uvs=True)
-    return gltf_path
-
-
-title = "LIVE: Towards Layer-wise Image Vectorization (CVPR 2022 Oral)"
-description = "This demo shows the effectiveness of LIVE <a href='' target='_blank'>Paper</a>. Given the input image, LIVE is able to progressively build the SVG output with a layer-wise representation." \
-              "<br>NOTE: for efficiency, we resize input images to 240x240 for Huggingface Space. "
-examples = [["examples/" + img] for img in os.listdir("examples/")]
-
-iface = gr.Interface(fn=process_image,
-                     inputs=[gr.inputs.Image(type="filepath", label="Input Image")
-                     ],
-                     outputs=[gr.outputs.Image(label="predicted depth", type="pil"),
-                              gr.outputs.Image3D(label="3d mesh reconstruction", clear_color=[
-                                                 1.0, 1.0, 1.0, 1.0]),
-                              gr.outputs.File(label="3d gLTF")],
-                     title=title,
-                     description=description,
-                     examples=examples,
-                     allow_flagging="never")
-iface.launch(debug=True, enable_queue=False, cache_examples=True)
+
+ROOT_PATH = sys.path[0]  # 根目录
+
+# 模型路径
+model_path = "ultralytics/yolov5"
+
+
+# 模型名称临时变量
+model_name_tmp = ""
+
+# 设备临时变量
+device_tmp = ""
+
+# 文件后缀
+suffix_list = [".csv", ".yaml"]
+
+
+def parse_args(known=False):
+    parser = argparse.ArgumentParser(description="Gradio YOLOv5 Det v0.1")
+    parser.add_argument(
+        "--model_name", "-mn", default="yolov5s", type=str, help="model name"
+    )
+    parser.add_argument(
+        "--model_cfg",
+        "-mc",
+        default="./model_config/model_name_p5_all.yaml",
+        type=str,
+        help="model config",
+    )
+    parser.add_argument(
+        "--cls_name",
+        "-cls",
+        default="./cls_name/cls_name.yaml",
+        type=str,
+        help="cls name",
+    )
+    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(
+        "--label_dnt_show",
+        "-lds",
+        action="store_false",
+        default=True,
+        help="label show",
+    )
+    parser.add_argument(
+        "--device",
+        "-dev",
+        default="cpu",
+        type=str,
+        help="cuda or cpu, hugging face only cpu",
+    )
+    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
+
+
+#  模型加载
+def model_loading(model_name, device):
+
+    # 加载本地模型
+    model = torch.hub.load(model_path, model_name, force_reload=True, device=device)
+
+    return model
+
+
+# 检测信息
+def export_json(results, model, img_size):
+
+    return [
+        [
+            {
+                "id": int(i),
+                "class": int(result[i][5]),
+                "class_name": model.model.names[int(result[i][5])],
+                "normalized_box": {
+                    "x0": round(result[i][:4].tolist()[0], 6),
+                    "y0": round(result[i][:4].tolist()[1], 6),
+                    "x1": round(result[i][:4].tolist()[2], 6),
+                    "y1": round(result[i][:4].tolist()[3], 6),
+                },
+                "confidence": round(float(result[i][4]), 2),
+                "fps": round(1000 / float(results.t[1]), 2),
+                "width": img_size[0],
+                "height": img_size[1],
+            }
+            for i in range(len(result))
+        ]
+        for result in results.xyxyn
+    ]
+
+
+# YOLOv5图片检测函数
+def yolo_det(img, device, model_name, inference_size, conf, iou, label_opt, model_cls):
+
+    global model, model_name_tmp, device_tmp
+
+    if model_name_tmp != model_name:
+        # 模型判断，避免反复加载
+        model_name_tmp = model_name
+        model = model_loading(model_name_tmp, device)
+    elif device_tmp != device:
+        device_tmp = device
+        model = model_loading(model_name_tmp, device)
+
+    # -----------模型调参-----------
+    model.conf = conf  # NMS 置信度阈值
+    model.iou = iou  # NMS IOU阈值
+    model.max_det = 1000  # 最大检测框数
+    model.classes = model_cls  # 模型类别
+
+    results = model(img, size=inference_size)  # 检测
+    results.render(labels=label_opt)  # 渲染
+
+    det_img = Image.fromarray(results.imgs[0])  # 检测图片
+
+    det_json = export_json(results, model, img.size)[0]  # 检测信息
+
+    return det_img, det_json
+
+
+# yaml文件解析
+def yaml_parse(file_path):
+    return yaml.safe_load(open(file_path, "r", 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 main(args):
+    gr.close_all()
+
+    global model
+
+    slider_step = 0.05  # 滑动步长
+
+    nms_conf = args.nms_conf
+    nms_iou = args.nms_iou
+    label_opt = args.label_dnt_show
+    model_name = args.model_name
+    model_cfg = args.model_cfg
+    cls_name = args.cls_name
+    device = args.device
+    inference_size = args.inference_size
+
+    # 模型加载
+    model = model_loading(model_name, device)
+
+    model_names = yaml_csv(model_cfg, "model_names")
+    model_cls_name = yaml_csv(cls_name, "model_cls_name")
+
+    # -------------------输入组件-------------------
+    inputs_img = gr.inputs.Image(type="pil", label="原始图片")
+    device = gr.inputs.Dropdown(
+        choices=["cpu"], default=device, type="value", label="设备"
+    )
+    inputs_model = gr.inputs.Dropdown(
+        choices=model_names, default=model_name, type="value", label="模型"
+    )
+    inputs_size = gr.inputs.Radio(
+        choices=[320, 640], default=inference_size, label="推理尺寸"
+    )
+    input_conf = gr.inputs.Slider(
+        0, 1, step=slider_step, default=nms_conf, label="置信度阈值"
+    )
+    inputs_iou = gr.inputs.Slider(
+        0, 1, step=slider_step, default=nms_iou, label="IoU 阈值"
+    )
+    inputs_label = gr.inputs.Checkbox(default=label_opt, label="标签显示")
+    inputs_clsName = gr.inputs.CheckboxGroup(
+        choices=model_cls_name, default=model_cls_name, type="index", label="类别"
+    )
+
+    # 输入参数
+    inputs = [
+        inputs_img,  # 输入图片
+        device,  # 设备
+        inputs_model,  # 模型
+        inputs_size,  # 推理尺寸
+        input_conf,  # 置信度阈值
+        inputs_iou,  # IoU阈值
+        inputs_label,  # 标签显示
+        inputs_clsName,  # 类别
+    ]
+    # 输出参数
+    outputs = gr.outputs.Image(type="pil", label="检测图片")
+    outputs02 = gr.outputs.JSON(label="检测信息")
+
+    # 标题
+    title = "基于Gradio的YOLOv5通用目标检测系统"
+    # 描述
+    description = "<div align='center'>可自定义目标检测模型、安装简单、使用方便</div>"
+
+    # 示例图片
+    examples = [
+        [
+            "./img_example/bus.jpg",
+            "cpu",
+            "yolov5s",
+            640,
+            0.6,
+            0.5,
+            True,
+            ["人", "公交车"],
+        ],
+        [
+            "./img_example/Millenial-at-work.jpg",
+            "cpu",
+            "yolov5l",
+            320,
+            0.5,
+            0.45,
+            True,
+            ["人", "椅子", "杯子", "笔记本电脑"],
+        ],
+        [
+            "./img_example/zidane.jpg",
+            "cpu",
+            "yolov5m",
+            640,
+            0.25,
+            0.5,
+            False,
+            ["人", "领带"],
+        ],
+    ]
+
+    # 接口
+    gr.Interface(
+        fn=yolo_det,
+        inputs=inputs,
+        outputs=[outputs, outputs02],
+        title=title,
+        description=description,
+        examples=examples,
+        theme="seafoam",
+        # live=True, # 实时变更输出
+        flagging_dir="run"  # 输出目录
+        # ).launch(inbrowser=True, auth=['admin', 'admin'])
+    ).launch(
+        inbrowser=True,  # 自动打开默认浏览器
+        show_tips=True,  # 自动显示gradio最新功能
+        favicon_path="./icon/logo.ico",
+    )
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)

cls_name/cls_name.csv

@@ -0,0 +1,80 @@
+人
+自行车
+汽车
+摩托车
+飞机
+公交车
+火车
+卡车
+船
+红绿灯
+消防栓
+停止标志
+停车收费表
+长凳
+鸟
+猫
+狗
+马
+羊
+牛
+象
+熊
+斑马
+长颈鹿
+背包
+雨伞
+手提包
+领带
+手提箱
+飞盘
+滑雪板
+单板滑雪
+运动球
+风筝
+棒球棒
+棒球手套
+滑板
+冲浪板
+网球拍
+瓶子
+红酒杯
+杯子
+叉子
+刀
+勺
+碗
+香蕉
+苹果
+三明治
+橙子
+西兰花
+胡萝卜
+热狗
+比萨
+甜甜圈
+蛋糕
+椅子
+长椅
+盆栽
+床
+餐桌
+马桶
+电视
+笔记本电脑
+鼠标
+遥控器
+键盘
+手机
+微波炉
+烤箱
+烤面包机
+洗碗槽
+冰箱
+书
+时钟
+花瓶
+剪刀
+泰迪熊
+吹风机
+牙刷

cls_name/cls_name.yaml

@@ -0,0 +1,7 @@
+model_cls_name: ['人', '自行车', '汽车', '摩托车', '飞机', '公交车', '火车', '卡车', '船', '红绿灯', '消防栓', '停止标志',
+                '停车收费表', '长凳', '鸟', '猫', '狗', '马', '羊', '牛', '象', '熊', '斑马', '长颈鹿', '背包', '雨伞', '手提包', '领带',
+                '手提箱', '飞盘', '滑雪板', '单板滑雪', '运动球', '风筝', '棒球棒', '棒球手套', '滑板', '冲浪板', '网球拍', '瓶子', '红酒杯',
+                '杯子', '叉子', '刀', '勺', '碗', '香蕉', '苹果', '三明治', '橙子', '西兰花', '胡萝卜', '热狗', '比萨', '甜甜圈', '蛋糕',
+                '椅子', '长椅', '盆栽', '床', '餐桌', '马桶', '电视', '笔记本电脑', '鼠标', '遥控器', '键盘', '手机', '微波炉', '烤箱',
+                '烤面包机', '洗碗槽', '冰箱', '书', '时钟', '花瓶', '剪刀', '泰迪熊', '吹风机', '牙刷'
+            ]

model_config/model_name_p5_all.csv

@@ -0,0 +1,5 @@
+yolov5n
+yolov5s
+yolov5m
+yolov5l
+yolov5x

model_config/model_name_p5_all.yaml

@@ -0,0 +1 @@
+model_names: ["yolov5n", "yolov5s", "yolov5m", "yolov5l", "yolov5x"]

model_config/model_name_p5_n.csv

@@ -0,0 +1 @@
+yolov5n

model_config/model_name_p5_n.yaml

@@ -0,0 +1 @@
+model_names: ["yolov5n"]

model_config/model_name_p6_all.csv

@@ -0,0 +1,5 @@
+yolov5n6
+yolov5s6
+yolov5m6
+yolov5l6
+yolov5x6

model_config/model_name_p6_all.yaml

@@ -0,0 +1 @@
+model_names: ["yolov5n6", "yolov5s6", "yolov5m6", "yolov5l6", "yolov5x6"]

model_download/yolov5_model_p5_all.sh

@@ -0,0 +1,8 @@
+cd ./yolov5
+
+# 下载YOLOv5模型
+wget -c -t 0 https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5n.pt
+wget -c -t 0 https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5s.pt
+wget -c -t 0 https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5m.pt
+wget -c -t 0 https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5l.pt
+wget -c -t 0 https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5x.pt

model_download/yolov5_model_p5_n.sh

@@ -0,0 +1,4 @@
+cd ./yolov5
+
+# 下载YOLOv5模型
+wget -c -t 0 https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5n.pt

model_download/yolov5_model_p6_all.sh

@@ -0,0 +1,8 @@
+cd ./yolov5
+
+# 下载YOLOv5模型
+wget -c -t 0 https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5n6.pt
+wget -c -t 0 https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5s6.pt
+wget -c -t 0 https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5m6.pt
+wget -c -t 0 https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5l6.pt
+wget -c -t 0 https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5x6.pt

requirements.txt

@@ -1,7 +1,35 @@
-torch
-git+https://github.com/nielsrogge/transformers.git@add_dpt_redesign#egg=transformers
-numpy
-Pillow
-gradio>=2.9.3
-jinja2
-open3d
+# Base ----------------------------------------
+matplotlib>=3.2.2
+numpy>=1.18.5
+opencv-python-headless>=4.5.5.64
+Pillow>=7.1.2
+PyYAML>=5.3.1
+requests>=2.23.0
+scipy>=1.4.1
+torch>=1.7.0
+torchvision>=0.8.1
+tqdm>=4.41.0
+
+# Logging -------------------------------------
+tensorboard>=2.4.1
+# wandb
+
+# Plotting ------------------------------------
+pandas>=1.1.4
+seaborn>=0.11.0
+
+# Export --------------------------------------
+# coremltools>=4.1  # CoreML export
+# onnx>=1.9.0  # ONNX export
+# onnx-simplifier>=0.3.6  # ONNX simplifier
+# scikit-learn==0.19.2  # CoreML quantization
+# tensorflow>=2.4.1  # TFLite export
+# tensorflowjs>=3.9.0  # TF.js export
+# openvino-dev  # OpenVINO export
+
+# Extras --------------------------------------
+# albumentations>=1.0.3
+# Cython  # for pycocotools https://github.com/cocodataset/cocoapi/issues/172
+# pycocotools>=2.0  # COCO mAP
+# roboflow
+thop  # FLOPs computation