add legend

app.py

@@ -175,12 +175,23 @@ css = """
 def reset_enhance():
     return "请选择样例图片或者上传", None, None
 
-def reset_seg():
-    return "请选择样例图片或者上传", None, None
+def reset_seg_full():
+    return "请选择样例图片或者上传", None, None, None, False
 
 def reset_diag():
     return "请选择样例图片或者上传", None, [], "", ""
 
+def segment_demo_with_legend(img_name):
+    result = segment_demo(img_name)
+    if result is not None:
+        legend_html = """<div style="margin-top:8px; font-size:15px;">
+        <span style="display:inline-block;width:18px;height:18px;background:#F33;border-radius:3px;vertical-align:-4px;border:1.5px solid #B00;margin-right:9px"></span>
+        红色部分为息肉
+        </div>"""
+        return result, legend_html, True
+    else:
+        return None, None, False
+
 with gr.Blocks(title="AI内窥镜影像辅助分析系统") as demo:
     gr.HTML(f"<style>{css}</style>")
     gr.Markdown("<span style='font-size: 32px; font-weight: bold;'>AI内窥镜影像辅助分析系统</span>")
@@ -210,12 +221,16 @@ with gr.Blocks(title="AI内窥镜影像辅助分析系统") as demo:
                 with gr.Row():
                     seg_input_img = gr.Image(label="输入图片", interactive=False, width=300, height=300)
                     seg_output_img = gr.Image(label="分割结果图片", interactive=False, width=300, height=300)
+                seg_legend = gr.HTML(value=None, visible=False)
                 with gr.Row(elem_classes="button-row"):
                     seg_reset_btn = gr.Button("Reset", elem_classes="gray-btn")
                     seg_btn = gr.Button("点击执行分割", elem_classes="orange-btn")
                 seg_select.change(fn=show_seg_input, inputs=seg_select, outputs=seg_input_img)
-                seg_btn.click(fn=segment_demo, inputs=seg_select, outputs=seg_output_img)
-                seg_reset_btn.click(fn=reset_seg, inputs=None, outputs=[seg_select, seg_input_img, seg_output_img])
+                seg_btn.click(fn=segment_demo_with_legend, inputs=seg_select,
+                              outputs=[seg_output_img, seg_legend, seg_legend.visible])
+                seg_reset_btn.click(fn=reset_seg_full, 
+                                    inputs=None, 
+                                    outputs=[seg_select, seg_input_img, seg_output_img, seg_legend, seg_legend.visible])
 
     # 第二行：疾病诊断（大框包裹）
     with gr.Row():