Update app.py

app.py

@@ -2,12 +2,22 @@ from ultralytics import YOLO
 import gradio as gr
 import cv2
 import numpy as np
-import torch
 from collections import defaultdict
+import os
 
-model = YOLO('./yolo11x-seg.pt')
+# 初始化模型
+model = None
+
+def load_model():
+    global model
+    if model is None:
+        model = YOLO('yolo11x-seg.pt')
+    return model
 
 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)
@@ -92,69 +102,66 @@ def segment_image(image, conf_threshold, iou_threshold, mask_threshold, line_thi
     
     return gallery_output if gallery_output else None
 
-# 创建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="启用高分辨率分割掩码（可能会降低速度）"
+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.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="启用高分辨率分割掩码（可能会降低速度）"
+                    )
+            
+            with gr.Column(scale=1):
+                output_gallery = gr.Gallery(
+                    label="分割结果",
+                    show_label=True,
+                    columns=2,
+                    rows=2,
+                    height=600,
+                    object_fit="contain"
                 )
         
-        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("开始分割")
-    
-    gr.Examples(
-        examples=["./WX20250102-114134.png"],
-        inputs=input_image,
-    )
+        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
+        )
     
-    submit_btn.click(
-        fn=segment_image,
-        inputs=[
-            input_image,
-            conf_threshold,
-            iou_threshold,
-            mask_threshold,
-            line_thickness,
-            retina_masks
-        ],
-        outputs=output_gallery
-    )
+    return demo
 
-# 启动应用
 if __name__ == "__main__":
-    iface.launch() 
+    demo = create_demo()
+    demo.launch()