Update app.py

app.py

@@ -72,10 +72,19 @@ for model_name, model_path in MODEL_PATHS.items():
     if os.path.isfile(model_path):
         try:
             print(f"正在加载ONNX模型: {model_path}")
-            # Create ONNX runtime session
-            providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+            # 检查可用的执行提供程序
+            available_providers = ort.get_available_providers()
+            providers = []
+            
+            # 优先使用CUDA，如果可用
+            if 'CUDAExecutionProvider' in available_providers:
+                providers.append('CUDAExecutionProvider')
+            
+            # 总是添加CPU作为后备
+            providers.append('CPUExecutionProvider')
+            
             sessions[model_name] = ort.InferenceSession(model_path, providers=providers)
-            print(f"{model_name} ONNX模型加载成功")
+            print(f"{model_name} ONNX模型加载成功，使用提供程序: {providers}")
         except Exception as e:
             print(f"加载 {model_name} ONNX模型出错: {str(e)}")
             # Continue loading other models if one fails
@@ -87,10 +96,26 @@ if not sessions:
     raise RuntimeError("无法加载任何ONNX模型，请确保至少有一个模型文件是正确的")
 
 # Load feature extractor for image preprocessing
-feature_extractor = SegformerFeatureExtractor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
+try:
+    # 尝试使用新的API (如果transformers版本支持)
+    from transformers import AutoImageProcessor
+    feature_extractor = AutoImageProcessor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
+    print("使用AutoImageProcessor加载特征提取器")
+except (ImportError, AttributeError):
+    # 回退到旧API
+    feature_extractor = SegformerFeatureExtractor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
+    print("使用SegformerFeatureExtractor加载特征提取器")
 
 # Define segmentation classes and corresponding colors
 CLASSES = {
+    0: "Background",
+    1: "Leaf",
+    2: "Black Spot",
+    3: "Black Rot"
+}
+
+# 中文类别名称 (用于UI界面)
+CLASSES_CN = {
     0: "背景",
     1: "叶片",
     2: "黑痣病",
@@ -211,10 +236,7 @@ def segment_image(image, model_choice="Segformer", visualization_type="overlay")
     
     # 左侧显示原图
     ax1.imshow(original_array)
-    if font_prop:
-        ax1.set_title("原始图像", fontproperties=font_prop, fontsize=14, pad=10, fontweight='bold')
-    else:
-        ax1.set_title("原始图像", fontsize=14, pad=10, fontweight='bold')
+    ax1.set_title("Original Image", fontsize=14, pad=10, fontweight='bold')
     ax1.axis('off')
     ax1.set_frame_on(True)  # 显示边框
     ax1.patch.set_edgecolor('lightgray')  # 设置边框颜色
@@ -222,10 +244,7 @@ def segment_image(image, model_choice="Segformer", visualization_type="overlay")
     
     # 右侧显示分割结果
     ax2.imshow(overlay_image)
-    if font_prop:
-        ax2.set_title("分割结果", fontproperties=font_prop, fontsize=14, pad=10, fontweight='bold')
-    else:
-        ax2.set_title("分割结果", fontsize=14, pad=10, fontweight='bold')
+    ax2.set_title("Segmentation Result", fontsize=14, pad=10, fontweight='bold')
     ax2.axis('off')
     ax2.set_frame_on(True)  # 显示边框
     ax2.patch.set_edgecolor('lightgray')  # 设置边框颜色
@@ -235,19 +254,13 @@ def segment_image(image, model_choice="Segformer", visualization_type="overlay")
     if visualization_type == "detailed":
         # 热力图显示
         heatmap = ax3.imshow(segmentation_map, cmap=CUSTOM_CMAP, interpolation='nearest')
-        if font_prop:
-            ax3.set_title("病害热力图", fontproperties=font_prop, fontsize=14, pad=10, fontweight='bold')
-        else:
-            ax3.set_title("病害热力图", fontsize=14, pad=10, fontweight='bold')
+        ax3.set_title("Disease Heatmap", fontsize=14, pad=10, fontweight='bold')
         ax3.axis('off')
         
         # 添加颜色条
         cbar = plt.colorbar(heatmap, ax=ax3, orientation='horizontal', pad=0.05, shrink=0.8)
         cbar.set_ticks([0, 1, 2, 3])
-        if font_prop:
-            cbar.set_ticklabels([CLASSES[i] for i in range(4)], fontproperties=font_prop)
-        else:
-            cbar.set_ticklabels([CLASSES[i] for i in range(4)])
+        cbar.set_ticklabels([CLASSES[i] for i in range(4)])
         
         # 饼图显示各类别占比
         non_zero_classes = {k: v for k, v in class_percentages.items() if v > 0.5}  # 只显示占比大于0.5%的类别
@@ -268,16 +281,10 @@ def segment_image(image, model_choice="Segformer", visualization_type="overlay")
         )
         
         # 设置饼图标题
-        if font_prop:
-            ax4.set_title("类别分布", fontproperties=font_prop, fontsize=14, pad=10, fontweight='bold')
-        else:
-            ax4.set_title("类别分布", fontsize=14, pad=10, fontweight='bold')
+        ax4.set_title("Class Distribution", fontsize=14, pad=10, fontweight='bold')
         
         # 自定义饼图图例
-        if font_prop:
-            ax4.legend(wedges, labels, title="类别", loc="center left", bbox_to_anchor=(1, 0, 0.5, 1), prop=font_prop)
-        else:
-            ax4.legend(wedges, labels, title="类别", loc="center left", bbox_to_anchor=(1, 0, 0.5, 1))
+        ax4.legend(wedges, labels, title="Classes", loc="center left", bbox_to_anchor=(1, 0, 0.5, 1))
     
     # 添加美观的图例（对于非详细视图）
     if visualization_type != "detailed":
@@ -291,16 +298,11 @@ def segment_image(image, model_choice="Segformer", visualization_type="overlay")
                     legend_elements.append(plt.Rectangle((0, 0), 1, 1, color=color, label=legend_label))
         
         # 在图像下方添加图例
-        if font_prop:
-            fig.legend(handles=legend_elements, loc='lower center', bbox_to_anchor=(0.5, 0.02),
-                      fancybox=True, shadow=True, ncol=len(legend_elements), fontsize=10, prop=font_prop)
-        else:
-            fig.legend(handles=legend_elements, loc='lower center', bbox_to_anchor=(0.5, 0.02),
-                      fancybox=True, shadow=True, ncol=len(legend_elements), fontsize=10)
+        fig.legend(handles=legend_elements, loc='lower center', bbox_to_anchor=(0.5, 0.02),
+                  fancybox=True, shadow=True, ncol=len(legend_elements), fontsize=10)
     
     # 添加水印
-    fig.text(0.99, 0.01, '植物叶片病害分割系统', fontsize=8, color='gray', ha='right', va='bottom', alpha=0.7,
-             fontproperties=font_prop if font_prop else None)
+    fig.text(0.99, 0.01, 'Plant Leaf Disease Segmentation System', fontsize=8, color='gray', ha='right', va='bottom', alpha=0.7)
     
     # 调整布局，确保图例不会被裁剪
     plt.tight_layout()
@@ -356,4 +358,4 @@ iface = gr.Interface(
         ".gr-image {border-radius: 8px; border: 1px solid #ddd;}"
 )
 
-iface.launch(share=True)
+iface.launch()