Update app.py

app.py

@@ -2,22 +2,12 @@ from ultralytics import YOLO
 import gradio as gr
 import cv2
 import numpy as np
+import torch
 from collections import defaultdict
-import os
 
-# 初始化模型
-model = None
-
-def load_model():
-    global model
-    if model is None:
-        model = YOLO('yolo11x-seg.pt')
-    return model
+model = YOLO('./yolo11x-seg.pt')
 
 def segment_image(image, conf_threshold, iou_threshold, mask_threshold, line_thickness, use_retina_masks):
-    # 加载模型
-    model = load_model()
-    
     # 确保图像是BGR格式
     if len(image.shape) == 2:
         image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
@@ -102,66 +92,69 @@ def segment_image(image, conf_threshold, iou_threshold, mask_threshold, line_thi
     
     return gallery_output if gallery_output else None
 
-def create_demo():
-    with gr.Blocks() as demo:
-        gr.Markdown("# YOLO 图像分割")
-        gr.Markdown("上传一张图片，模型将对图片进行实例分割。每个检测到的类别将单独显示。")
-        
-        with gr.Row():
-            with gr.Column(scale=1):
-                input_image = gr.Image()
-                with gr.Row():
-                    conf_threshold = gr.Slider(
-                        minimum=0.1, maximum=1.0, value=0.25, step=0.05,
-                        label="置信度阈值", info="检测置信度的最小值"
-                    )
-                    iou_threshold = gr.Slider(
-                        minimum=0.1, maximum=1.0, value=0.45, step=0.05,
-                        label="IOU阈值", info="非极大值抑制的IOU阈值"
-                    )
-                with gr.Row():
-                    mask_threshold = gr.Slider(
-                        minimum=0.1, maximum=1.0, value=0.5, step=0.05,
-                        label="掩码阈值", info="分割掩码的阈值"
-                    )
-                    line_thickness = gr.Slider(
-                        minimum=1, maximum=5, value=2, step=1,
-                        label="线条粗细", info="边界框和文本的粗细"
-                    )
-                with gr.Row():
-                    retina_masks = gr.Checkbox(
-                        label="高分辨率掩码",
-                        value=True,
-                        info="启用高分辨率分割掩码（可能会降低速度）"
-                    )
-            
-            with gr.Column(scale=1):
-                output_gallery = gr.Gallery(
-                    label="分割结果",
-                    show_label=True,
-                    columns=2,
-                    rows=2,
-                    height=600,
-                    object_fit="contain"
+# 创建Gradio界面
+with gr.Blocks() as iface:
+    gr.Markdown("# YOLO 图像分割")
+    gr.Markdown("上传一张图片，模型将对图片进行实例分割。每个检测到的类别将单独显示。")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_image = gr.Image()
+            with gr.Row():
+                conf_threshold = gr.Slider(
+                    minimum=0.1, maximum=1.0, value=0.25, step=0.05,
+                    label="置信度阈值", info="检测置信度的最小值"
+                )
+                iou_threshold = gr.Slider(
+                    minimum=0.1, maximum=1.0, value=0.7, step=0.05,
+                    label="IOU阈值", info="非极大值抑制的IOU阈值"
+                )
+            with gr.Row():
+                mask_threshold = gr.Slider(
+                    minimum=0.1, maximum=1.0, value=0.5, step=0.05,
+                    label="掩码阈值", info="分割掩码的阈值"
+                )
+                line_thickness = gr.Slider(
+                    minimum=1, maximum=5, value=2, step=1,
+                    label="线条粗细", info="边界框和文本的粗细"
+                )
+            with gr.Row():
+                retina_masks = gr.Checkbox(
+                    label="高分辨率掩码",
+                    value=True,
+                    info="启用高分辨率分割掩码（可能会降低速度）"
                 )
         
-        submit_btn = gr.Button("开始分割")
-        
-        submit_btn.click(
-            fn=segment_image,
-            inputs=[
-                input_image,
-                conf_threshold,
-                iou_threshold,
-                mask_threshold,
-                line_thickness,
-                retina_masks
-            ],
-            outputs=output_gallery
-        )
+        with gr.Column(scale=1):
+            output_gallery = gr.Gallery(
+                label="分割结果",
+                show_label=True,
+                columns=2,
+                rows=2,
+                height=600,
+                object_fit="contain"
+            )
+    
+    submit_btn = gr.Button("开始分割")
     
-    return demo
+    gr.Examples(
+        examples=["./WX20250102-114134.png"],
+        inputs=input_image,
+    )
+    
+    submit_btn.click(
+        fn=segment_image,
+        inputs=[
+            input_image,
+            conf_threshold,
+            iou_threshold,
+            mask_threshold,
+            line_thickness,
+            retina_masks
+        ],
+        outputs=output_gallery
+    )
 
+# 启动应用
 if __name__ == "__main__":
-    demo = create_demo()
-    demo.launch() 
+    iface.launch()