Update app.py

app.py

@@ -14,6 +14,9 @@ from pathlib import Path
 
 import cv2
 import gradio as gr
+from gradio_imageslider import ImageSlider
+import tempfile
+import uuid
 import numpy as np
 from matplotlib import font_manager
 from ultralytics import YOLO
@@ -63,9 +66,28 @@ EXAMPLES_DET = [
     ["./img_examples/bus.jpg", "yolov8s", "cpu", 640, 0.6, 0.5, 100, "所有尺寸"],
     ["./img_examples/giraffe.jpg", "yolov8l", "cpu", 320, 0.5, 0.45, 100, "所有尺寸"],
     ["./img_examples/zidane.jpg", "yolov8m", "cpu", 640, 0.6, 0.5, 100, "所有尺寸"],
-    ["./img_examples/Millenial-at-work.jpg", "yolov8x", "cpu", 1280, 0.5, 0.5, 100, "所有尺寸"],
+    [
+        "./img_examples/Millenial-at-work.jpg",
+        "yolov8x",
+        "cpu",
+        1280,
+        0.5,
+        0.5,
+        100,
+        "所有尺寸",
+    ],
     ["./img_examples/bus.jpg", "yolov8s-seg", "cpu", 640, 0.6, 0.5, 100, "所有尺寸"],
-    ["./img_examples/Millenial-at-work.jpg", "yolov8x-seg", "cpu", 1280, 0.5, 0.5, 100, "所有尺寸"],]
+    [
+        "./img_examples/Millenial-at-work.jpg",
+        "yolov8x-seg",
+        "cpu",
+        1280,
+        0.5,
+        0.5,
+        100,
+        "所有尺寸",
+    ],
+]
 
 EXAMPLES_CLAS = [
     ["./img_examples/img_clas/ILSVRC2012_val_00000008.JPEG", "yolov8s-cls"],
@@ -73,7 +95,8 @@ EXAMPLES_CLAS = [
     ["./img_examples/img_clas/ILSVRC2012_val_00000023.JPEG", "yolov8m-cls"],
     ["./img_examples/img_clas/ILSVRC2012_val_00000067.JPEG", "yolov8m-cls"],
     ["./img_examples/img_clas/ILSVRC2012_val_00000077.JPEG", "yolov8m-cls"],
-    ["./img_examples/img_clas/ILSVRC2012_val_00000247.JPEG", "yolov8m-cls"],]
+    ["./img_examples/img_clas/ILSVRC2012_val_00000247.JPEG", "yolov8m-cls"],
+]
 
 GYD_CSS = """#disp_image {
         text-align: center; /* Horizontally center the content */
@@ -82,7 +105,9 @@ GYD_CSS = """#disp_image {
 
 def parse_args(known=False):
     parser = argparse.ArgumentParser(description=GYD_VERSION)
-    parser.add_argument("--model_name", "-mn", default="yolov8s", type=str, help="model name")
+    parser.add_argument(
+        "--model_name", "-mn", default="yolov8s", type=str, help="model name"
+    )
     parser.add_argument(
         "--model_cfg",
         "-mc",
@@ -111,10 +136,18 @@ def parse_args(known=False):
         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")
-    parser.add_argument("--max_detnum", "-mdn", default=50, type=float, help="model max det num")
-    parser.add_argument("--slider_step", "-ss", default=0.05, type=float, help="slider step")
+    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"
+    )
+    parser.add_argument(
+        "--max_detnum", "-mdn", default=50, type=float, help="model max det num"
+    )
+    parser.add_argument(
+        "--slider_step", "-ss", default=0.05, type=float, help="slider step"
+    )
     parser.add_argument(
         "--is_login",
         "-isl",
@@ -122,12 +155,14 @@ def parse_args(known=False):
         default=False,
         help="is login",
     )
-    parser.add_argument('--usr_pwd',
-                        "-up",
-                        nargs='+',
-                        type=str,
-                        default=["admin", "admin"],
-                        help="user & password for login")
+    parser.add_argument(
+        "--usr_pwd",
+        "-up",
+        nargs="+",
+        type=str,
+        default=["admin", "admin"],
+        help="user & password for login",
+    )
     parser.add_argument(
         "--is_share",
         "-is",
@@ -135,7 +170,9 @@ def parse_args(known=False):
         default=False,
         help="is login",
     )
-    parser.add_argument("--server_port", "-sp", default=7860, type=int, help="server port")
+    parser.add_argument(
+        "--server_port", "-sp", default=7860, type=int, help="server port"
+    )
 
     args = parser.parse_known_args()[0] if known else parser.parse_args()
     return args
@@ -167,6 +204,7 @@ def check_online():
     # 参考：https://github.com/ultralytics/yolov5/blob/master/utils/general.py
     # Check internet connectivity
     import socket
+
     try:
         socket.create_connection(("1.1.1.1", 443), 5)  # check host accessibility
         return True
@@ -203,9 +241,12 @@ 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)  # 边界框
+    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)  # 标签尺寸 pillow 9.5.0
         # left, top, left + width, top + height
@@ -214,7 +255,12 @@ def pil_draw(img, score_l, bbox_l, cls_l, cls_index_l, textFont, color_list):
         # 标签背景
         img_pil.rectangle(
             # (xmin, ymin, xmin + text_w, ymin + text_h), # pillow 9.5.0
-            (xmin, ymin, xmin + text_xmax - text_xmin, ymin + text_ymax - text_ymin),  # pillow 10.0.0
+            (
+                xmin,
+                ymin,
+                xmin + text_xmax - text_xmin,
+                ymin + text_ymax - text_ymin,
+            ),  # pillow 10.0.0
             fill=color_list[cls_index],
             outline=color_list[cls_index],
         )
@@ -268,10 +314,19 @@ def seg_output(img_path, seg_mask_list, color_list, cls_list):
 
 
 # 目标检测和图像分割模型加载
-def model_det_loading(img_path, device_opt, conf, iou, infer_size, max_det, yolo_model="yolov8n.pt"):
+def model_det_loading(
+    img_path, device_opt, conf, iou, infer_size, max_det, yolo_model="yolov8n.pt"
+):
     model = YOLO(yolo_model)
 
-    results = model(source=img_path, device=device_opt, imgsz=infer_size, conf=conf, iou=iou, max_det=max_det)
+    results = model(
+        source=img_path,
+        device=device_opt,
+        imgsz=infer_size,
+        conf=conf,
+        iou=iou,
+        max_det=max_det,
+    )
     results = list(results)[0]
     return results
 
@@ -286,8 +341,9 @@ def model_cls_loading(img_path, yolo_model="yolov8s-cls.pt"):
 
 
 # YOLOv8图片检测函数
-def yolo_det_img(img_path, model_name, device_opt, infer_size, conf, iou, max_det, obj_size):
-
+def yolo_det_img(
+    img_path, model_name, device_opt, infer_size, conf, iou, max_det, obj_size
+):
     global model, model_name_tmp, device_tmp
 
     s_obj, m_obj, l_obj = 0, 0, 0
@@ -300,13 +356,15 @@ def yolo_det_img(img_path, model_name, device_opt, infer_size, conf, iou, max_de
     cls_index_det_stat = []  # 1
 
     # 模型加载
-    predict_results = model_det_loading(img_path,
-                                        device_opt,
-                                        conf,
-                                        iou,
-                                        infer_size,
-                                        max_det,
-                                        yolo_model=f"{model_name}.pt")
+    predict_results = model_det_loading(
+        img_path,
+        device_opt,
+        conf,
+        iou,
+        infer_size,
+        max_det,
+        yolo_model=f"{model_name}.pt",
+    )
     # 检测参数
     xyxy_list = predict_results.boxes.xyxy.cpu().numpy().tolist()
     conf_list = predict_results.boxes.conf.cpu().numpy().tolist()
@@ -315,34 +373,39 @@ def yolo_det_img(img_path, model_name, device_opt, infer_size, conf, iou, max_de
     # 颜色列表
     color_list = random_color(len(model_cls_name_cp), True)
 
+    img = Image.open(img_path)
+    img_cp = img.copy()
+
     # 图像分割
-    if (model_name[-3:] == "seg"):
+    if model_name[-3:] == "seg":
         # masks_list = predict_results.masks.xyn
         masks_list = predict_results.masks.xy
         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)
+        img = Image.fromarray(cv2.cvtColor(img_mask_merge, cv2.COLOR_BGRA2RGB))
 
     # 判断检测对象是否为空
-    if (xyxy_list != []):
-
+    if xyxy_list != []:
         # ---------------- 加载字体 ----------------
         yaml_index = cls_name.index(".yaml")
-        cls_name_lang = cls_name[yaml_index - 2:yaml_index]
+        cls_name_lang = cls_name[yaml_index - 2 : yaml_index]
 
         if cls_name_lang == "zh":
             # 中文
-            textFont = ImageFont.truetype(str(f"{ROOT_PATH}/fonts/SimSun.ttf"), size=FONTSIZE)
+            textFont = ImageFont.truetype(
+                str(f"{ROOT_PATH}/fonts/SimSun.ttf"), size=FONTSIZE
+            )
         elif cls_name_lang in ["en", "ru", "es", "ar"]:
             # 英文、俄语、西班牙语、阿拉伯语
-            textFont = ImageFont.truetype(str(f"{ROOT_PATH}/fonts/TimesNewRoman.ttf"), size=FONTSIZE)
+            textFont = ImageFont.truetype(
+                str(f"{ROOT_PATH}/fonts/TimesNewRoman.ttf"), size=FONTSIZE
+            )
         elif cls_name_lang == "ko":
             # 韩语
-            textFont = ImageFont.truetype(str(f"{ROOT_PATH}/fonts/malgun.ttf"), size=FONTSIZE)
+            textFont = ImageFont.truetype(
+                str(f"{ROOT_PATH}/fonts/malgun.ttf"), size=FONTSIZE
+            )
 
         for i in range(len(xyxy_list)):
-
             # ------------ 边框坐标 ------------
             x0 = int(xyxy_list[i][0])
             y0 = int(xyxy_list[i][1])
@@ -354,7 +417,7 @@ def yolo_det_img(img_path, model_name, device_opt, infer_size, conf, iou, max_de
             h_obj = y1 - y0
             area_obj = w_obj * h_obj  # 目标尺寸
 
-            if (obj_size == obj_style[0] and area_obj > 0 and area_obj <= 32 ** 2):
+            if obj_size == obj_style[0] and area_obj > 0 and area_obj <= 32**2:
                 obj_cls_index = int(cls_list[i])  # 类别索引
                 cls_index_det_stat.append(obj_cls_index)
 
@@ -367,7 +430,9 @@ def yolo_det_img(img_path, model_name, device_opt, infer_size, conf, iou, max_de
                 score_det_stat.append(conf)
 
                 area_obj_all.append(area_obj)
-            elif (obj_size == obj_style[1] and area_obj > 32 ** 2 and area_obj <= 96 ** 2):
+            elif (
+                obj_size == obj_style[1] and area_obj > 32**2 and area_obj <= 96**2
+            ):
                 obj_cls_index = int(cls_list[i])  # 类别索引
                 cls_index_det_stat.append(obj_cls_index)
 
@@ -380,7 +445,7 @@ def yolo_det_img(img_path, model_name, device_opt, infer_size, conf, iou, max_de
                 score_det_stat.append(conf)
 
                 area_obj_all.append(area_obj)
-            elif (obj_size == obj_style[2] and area_obj > 96 ** 2):
+            elif obj_size == obj_style[2] and area_obj > 96**2:
                 obj_cls_index = int(cls_list[i])  # 类别索引
                 cls_index_det_stat.append(obj_cls_index)
 
@@ -393,7 +458,7 @@ def yolo_det_img(img_path, model_name, device_opt, infer_size, conf, iou, max_de
                 score_det_stat.append(conf)
 
                 area_obj_all.append(area_obj)
-            elif (obj_size == "所有尺寸"):
+            elif obj_size == "所有尺寸":
                 obj_cls_index = int(cls_list[i])  # 类别索引
                 cls_index_det_stat.append(obj_cls_index)
 
@@ -407,20 +472,30 @@ def yolo_det_img(img_path, model_name, device_opt, infer_size, conf, iou, max_de
 
                 area_obj_all.append(area_obj)
 
-        det_img = pil_draw(img, 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)):
-            if (0 < area_obj_all[i] <= 32 ** 2):
+            if 0 < area_obj_all[i] <= 32**2:
                 s_obj = s_obj + 1
-            elif (32 ** 2 < area_obj_all[i] <= 96 ** 2):
+            elif 32**2 < area_obj_all[i] <= 96**2:
                 m_obj = m_obj + 1
-            elif (area_obj_all[i] > 96 ** 2):
+            elif area_obj_all[i] > 96**2:
                 l_obj = l_obj + 1
 
         sml_obj_total = s_obj + m_obj + l_obj
         objSize_dict = {}
-        objSize_dict = {obj_style[i]: [s_obj, m_obj, l_obj][i] / sml_obj_total for i in range(3)}
+        objSize_dict = {
+            obj_style[i]: [s_obj, m_obj, l_obj][i] / sml_obj_total for i in range(3)
+        }
 
         # ------------ 类别统计 ------------
         clsRatio_dict = {}
@@ -429,15 +504,23 @@ def yolo_det_img(img_path, model_name, device_opt, infer_size, conf, iou, max_de
         for k, v in clsDet_dict.items():
             clsRatio_dict[k] = v / clsDet_dict_sum
 
+        images = (det_img, img_cp)
+        images_names = ("det", "raw")
+        images_path = tempfile.mkdtemp()
+        images_paths = []
+        uuid_name = uuid.uuid4()
+        for image, image_name in zip(images, images_names):
+            image.save(images_path + f"/img_{uuid_name}_{image_name}.jpg")
+            images_paths.append(images_path + f"/img_{uuid_name}_{image_name}.jpg")
+
         gr.Info("图片检测成功！")
-        return det_img, objSize_dict, clsRatio_dict
+        return (det_img, img_cp), images_paths, objSize_dict, clsRatio_dict
     else:
         raise gr.Error("图片检测失败！")
 
 
 # YOLOv8图片分类函数
 def yolo_cls_img(img_path, model_name):
-
     # 模型加载
     predict_results = model_cls_loading(img_path, yolo_model=f"{model_name}.pt")
 
@@ -484,8 +567,10 @@ def main(args):
     model_cls_name_cp = model_cls_name.copy()  # 类别名称
     model_cls_imagenet_name_cp = model_cls_imagenet_name.copy()  # 类别名称
 
-    custom_theme = gr.themes.Soft(primary_hue="blue").set(button_secondary_background_fill="*neutral_100",
-                                                          button_secondary_background_fill_hover="*neutral_200")
+    custom_theme = gr.themes.Soft(primary_hue="blue").set(
+        button_secondary_background_fill="*neutral_100",
+        button_secondary_background_fill_hover="*neutral_200",
+    )
 
     custom_css = GYD_CSS
 
@@ -500,58 +585,116 @@ def main(args):
                 with gr.Tabs():
                     with gr.TabItem("目标检测与图像分割"):
                         with gr.Row():
-                            inputs_img = gr.Image(image_mode="RGB", type="filepath", label="原始图片")
+                            inputs_img = gr.Image(
+                                image_mode="RGB", type="filepath", label="原始图片"
+                            )
                         with gr.Row():
-                            device_opt = gr.Radio(choices=["cpu", "0", "1", "2", "3"], value="cpu", label="设备")
+                            device_opt = gr.Radio(
+                                choices=["cpu", "0", "1", "2", "3"],
+                                value="cpu",
+                                label="设备",
+                            )
                         with gr.Row():
-                            inputs_model = gr.Dropdown(choices=model_names, value=model_name, type="value", label="模型")
+                            inputs_model = gr.Dropdown(
+                                choices=model_names,
+                                value=model_name,
+                                type="value",
+                                label="模型",
+                            )
                         with gr.Accordion("高级设置", open=True):
                             with gr.Row():
-                                inputs_size = gr.Slider(320, 1600, step=1, value=inference_size, label="推理尺寸")
-                                max_det = gr.Slider(1, 1000, step=1, value=max_detnum, label="最大检测数")
+                                inputs_size = gr.Slider(
+                                    320,
+                                    1600,
+                                    step=1,
+                                    value=inference_size,
+                                    label="推理尺寸",
+                                )
+                                max_det = gr.Slider(
+                                    1, 1000, step=1, value=max_detnum, label="最大检测数"
+                                )
                             with gr.Row():
-                                input_conf = gr.Slider(0, 1, step=slider_step, value=nms_conf, label="置信度阈值")
-                                inputs_iou = gr.Slider(0, 1, step=slider_step, value=nms_iou, label="IoU 阈值")
+                                input_conf = gr.Slider(
+                                    0,
+                                    1,
+                                    step=slider_step,
+                                    value=nms_conf,
+                                    label="置信度阈值",
+                                )
+                                inputs_iou = gr.Slider(
+                                    0,
+                                    1,
+                                    step=slider_step,
+                                    value=nms_iou,
+                                    label="IoU 阈值",
+                                )
                             with gr.Row():
-                                obj_size = gr.Radio(choices=["所有尺寸", "小目标", "中目标", "大目标"], value="所有尺寸", label="目标尺寸")
+                                obj_size = gr.Radio(
+                                    choices=["所有尺寸", "小目标", "中目标", "大目标"],
+                                    value="所有尺寸",
+                                    label="目标尺寸",
+                                )
                         with gr.Row():
                             gr.ClearButton(inputs_img, value="清除")
-                            det_btn_img = gr.Button(value='检测', variant="primary")
+                            det_btn_img = gr.Button(value="检测", variant="primary")
                         with gr.Row():
                             gr.Examples(
                                 examples=EXAMPLES_DET,
                                 fn=yolo_det_img,
                                 inputs=[
-                                    inputs_img, inputs_model, device_opt, inputs_size, input_conf, inputs_iou, max_det,
-                                    obj_size],
+                                    inputs_img,
+                                    inputs_model,
+                                    device_opt,
+                                    inputs_size,
+                                    input_conf,
+                                    inputs_iou,
+                                    max_det,
+                                    obj_size,
+                                ],
                                 # outputs=[outputs_img, outputs_objSize, outputs_clsSize],
-                                cache_examples=False)
+                                cache_examples=False,
+                            )
 
                     with gr.TabItem("图像分类"):
                         with gr.Row():
-                            inputs_img_cls = gr.Image(image_mode="RGB", type="filepath", label="原始图片")
+                            inputs_img_cls = gr.Image(
+                                image_mode="RGB", type="filepath", label="原始图片"
+                            )
                         with gr.Row():
-                            inputs_model_cls = gr.Dropdown(choices=[
-                                "yolov8n-cls", "yolov8s-cls", "yolov8l-cls", "yolov8m-cls", "yolov8x-cls"],
-                                                           value="yolov8s-cls",
-                                                           type="value",
-                                                           label="模型")
+                            inputs_model_cls = gr.Dropdown(
+                                choices=[
+                                    "yolov8n-cls",
+                                    "yolov8s-cls",
+                                    "yolov8l-cls",
+                                    "yolov8m-cls",
+                                    "yolov8x-cls",
+                                ],
+                                value="yolov8s-cls",
+                                type="value",
+                                label="模型",
+                            )
                         with gr.Row():
                             gr.ClearButton(inputs_img, value="清除")
-                            det_btn_img_cls = gr.Button(value='检测', variant="primary")
+                            det_btn_img_cls = gr.Button(value="检测", variant="primary")
                         with gr.Row():
                             gr.Examples(
                                 examples=EXAMPLES_CLAS,
                                 fn=yolo_cls_img,
                                 inputs=[inputs_img_cls, inputs_model_cls],
                                 # outputs=[outputs_img_cls, outputs_ratio_cls],
-                                cache_examples=False)
+                                cache_examples=False,
+                            )
 
+            # -------- 输出 --------
             with gr.Column(scale=1):
                 with gr.Tabs():
                     with gr.TabItem("目标检测与图像分割"):
+                        # with gr.Row():
+                        #     outputs_img = gr.Image(type="pil", label="检测图片")
+                        with gr.Row():
+                            outputs_img_slider = ImageSlider(position=0.5, label="检测图片")
                         with gr.Row():
-                            outputs_img = gr.Image(type="pil", label="检测图片")
+                            outputs_imgfiles = gr.Files(label="图片下载")
                         with gr.Row():
                             outputs_objSize = gr.Label(label="目标尺寸占比统计")
                         with gr.Row():
@@ -604,15 +747,31 @@ def main(args):
             """
             )
 
-        det_btn_img.click(fn=yolo_det_img,
-                          inputs=[
-                              inputs_img, inputs_model, device_opt, inputs_size, input_conf, inputs_iou, max_det,
-                              obj_size],
-                          outputs=[outputs_img, outputs_objSize, outputs_clsSize])
+        det_btn_img.click(
+            fn=yolo_det_img,
+            inputs=[
+                inputs_img,
+                inputs_model,
+                device_opt,
+                inputs_size,
+                input_conf,
+                inputs_iou,
+                max_det,
+                obj_size,
+            ],
+            outputs=[
+                outputs_img_slider,
+                outputs_imgfiles,
+                outputs_objSize,
+                outputs_clsSize,
+            ],
+        )
 
-        det_btn_img_cls.click(fn=yolo_cls_img,
-                              inputs=[inputs_img_cls, inputs_model_cls],
-                              outputs=[outputs_img_cls, outputs_ratio_cls])
+        det_btn_img_cls.click(
+            fn=yolo_cls_img,
+            inputs=[inputs_img_cls, inputs_model_cls],
+            outputs=[outputs_img_cls, outputs_ratio_cls],
+        )
 
     return gyd