app update

app.py

@@ -1,14 +1,16 @@
-# Gradio YOLOv8 Det v0.1
+# Gradio YOLOv8 Det v0.2
 # 创建人：曾逸夫
-# 创建时间：2023-01-15
-
-import os
+# 创建时间：2023-01-20
 
 import argparse
 import csv
+import os
 import sys
 
 from ultralytics import YOLO
+
+csv.field_size_limit(sys.maxsize)
+
 import gc
 import json
 import random
@@ -34,22 +36,16 @@ SimSun = font_manager.FontProperties(fname=SimSun_path, size=12)
 # 新罗马字体
 TimesNesRoman = font_manager.FontProperties(fname=TimesNesRoman_path, size=12)
 
-from copy import deepcopy
-
+import torch
 import yaml
 from PIL import Image, ImageDraw, ImageFont
 
 from util.fonts_opt import is_fonts
-from util.pdf_opt import pdf_generate
 
 ROOT_PATH = sys.path[0]  # 根目录
 
-# 本地模型路径
-local_model_path = f"{ROOT_PATH}/models"
-local_model_path_02 = f"{ROOT_PATH}/yolov5"
-
 # Gradio YOLOv8 Det版本
-GYD_VERSION = "Gradio YOLOv8 Det v0.1"
+GYD_VERSION = "Gradio YOLOv8 Det v0.2"
 
 # 文件后缀
 suffix_list = [".csv", ".yaml"]
@@ -62,7 +58,7 @@ obj_style = ["小目标", "中目标", "大目标"]
 
 
 def parse_args(known=False):
-    parser = argparse.ArgumentParser(description="Gradio YOLOv8 Det v0.1")
+    parser = argparse.ArgumentParser(description="Gradio YOLOv8 Det v0.2")
     parser.add_argument("--model_type", "-mt", default="online", type=str, help="model type")
     parser.add_argument("--source", "-src", default="upload", type=str, help="image input source")
     parser.add_argument("--source_video", "-src_v", default="upload", type=str, help="video input source")
@@ -194,13 +190,11 @@ def random_color(cls_num, is_light=True):
 
 
 # 检测绘制
-def pil_draw(img_path, score_l, bbox_l, cls_l, cls_index_l, textFont, color_list):
-    img = Image.open(img_path)
+def pil_draw(img, score_l, bbox_l, cls_l, cls_index_l, textFont, color_list):
     img_pil = ImageDraw.Draw(img)
     id = 0
 
     for score, (xmin, ymin, xmax, ymax), label, cls_index in zip(score_l, bbox_l, cls_l, cls_index_l):
-
         img_pil.rectangle([xmin, ymin, xmax, ymax], fill=None, outline=color_list[cls_index], width=2)  # 边界框
         countdown_msg = f"{id}-{label} {score:.2f}"
         text_w, text_h = textFont.getsize(countdown_msg)  # 标签尺寸
@@ -226,6 +220,42 @@ def pil_draw(img_path, score_l, bbox_l, cls_l, cls_index_l, textFont, color_list
     return img
 
 
+# 绘制多边形
+def polygon_drawing(img_mask, canvas, color_seg):
+    # ------- RGB转BGR -------
+    color_seg = list(color_seg)
+    color_seg[0], color_seg[2] = color_seg[2], color_seg[0]
+    color_seg = tuple(color_seg)
+    # 定义多边形的顶点
+    pts = np.array(img_mask, dtype=np.int32)
+
+    # 多边形绘制
+    cv2.drawContours(canvas, [pts], -1, color_seg, thickness=-1)
+
+
+# 输出分割结果
+def seg_output(img_path, seg_mask_list, color_list, cls_list):
+    img = cv2.imread(img_path)
+    w, h = img.shape[1], img.shape[0]
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)
+
+    w, h = img.shape[1], img.shape[0]
+    canvas = np.zeros((h, w, 3), dtype=np.uint8)
+    canvas = cv2.cvtColor(canvas, cv2.COLOR_BGR2BGRA)
+
+    # 获取分割坐标
+    for seg_mask, cls_index in zip(seg_mask_list, cls_list):
+        img_mask = []
+        for i in range(len(seg_mask)):
+            img_mask.append([seg_mask[i][0] * w, seg_mask[i][1] * h])
+
+        polygon_drawing(img_mask, canvas, color_list[int(cls_index)])  # 绘制分割图形
+
+    img_mask_merge = cv2.add(img, canvas)  # 合并图像
+
+    return img_mask_merge
+
+
 # YOLOv5图片检测函数
 def yolo_det_img(img_path, model_name, infer_size, conf, iou):
 
@@ -242,12 +272,22 @@ def yolo_det_img(img_path, model_name, infer_size, conf, iou):
 
     # 模型加载
     predict_results = model_loading(img_path, conf, iou, infer_size, yolo_model=f"{model_name}.pt")
+    # 检测参数
     xyxy_list = predict_results.boxes.xyxy.cpu().numpy().tolist()
     conf_list = predict_results.boxes.conf.cpu().numpy().tolist()
     cls_list = predict_results.boxes.cls.cpu().numpy().tolist()
 
+    # 颜色列表
     color_list = random_color(len(model_cls_name_cp), True)
 
+    # 图像分割
+    if (model_name[-3:] == "seg"):
+        masks_list = predict_results.masks.segments
+        img_mask_merge = seg_output(img_path, masks_list, color_list, cls_list)
+        img = Image.fromarray(cv2.cvtColor(img_mask_merge, cv2.COLOR_BGRA2RGBA))
+    else:
+        img = Image.open(img_path)
+
     # 判断检测对象是否为空
     if (xyxy_list != []):
 
@@ -266,7 +306,6 @@ def yolo_det_img(img_path, model_name, infer_size, conf, iou):
             textFont = ImageFont.truetype(str(f"{ROOT_PATH}/fonts/malgun.ttf"), size=FONTSIZE)
 
         for i in range(len(xyxy_list)):
-            # id = int(i)  # 实例ID
             obj_cls_index = int(cls_list[i])  # 类别索引
             cls_index_det_stat.append(obj_cls_index)
 
@@ -284,16 +323,13 @@ def yolo_det_img(img_path, model_name, infer_size, conf, iou):
             conf = float(conf_list[i])  # 置信度
             score_det_stat.append(conf)
 
-            # fps = f"{(1000 / float(results.t[1])):.2f}"  # FPS
-
             # ---------- 加入目标尺寸 ----------
             w_obj = x1 - x0
             h_obj = y1 - y0
             area_obj = w_obj * h_obj
             area_obj_all.append(area_obj)
 
-        det_img = pil_draw(img_path, score_det_stat, bbox_det_stat, cls_det_stat, cls_index_det_stat, textFont,
-                           color_list)
+        det_img = pil_draw(img, score_det_stat, bbox_det_stat, cls_det_stat, cls_index_det_stat, textFont, color_list)
 
         # -------------- 目标尺寸计算 --------------
         for i in range(len(area_obj_all)):
@@ -372,7 +408,7 @@ def main(args):
     title = "Gradio YOLOv8 Det"
 
     # 描述
-    description = "Author: 曾逸夫（Zeng Yifu）, Github: https://github.com/Zengyf-CVer, thanks to [Gradio](https://github.com/gradio-app/gradio) & [YOLOv8](https://github.com/ultralytics/ultralytics)"
+    description = "<div align='center'>Object detection and image segmentation system based on YOLOv8</div><div align='center'>Author: 曾逸夫（Zeng Yifu）, Github: https://github.com/Zengyf-CVer, thanks to [Gradio](https://github.com/gradio-app/gradio) & [YOLOv8](https://github.com/ultralytics/ultralytics)</div>"
 
     # 示例图片
     examples_imgs = [
@@ -399,6 +435,18 @@ def main(args):
             "yolov8x",
             1280,
             0.5,
+            0.5,],
+        [
+            "./img_examples/bus.jpg",
+            "yolov8s-seg",
+            640,
+            0.6,
+            0.5,],
+        [
+            "./img_examples/Millenial-at-work.jpg",
+            "yolov8x-seg",
+            1280,
+            0.5,
             0.5,],]
 
     # 接口