Update app.py

app.py

@@ -1,15 +1,18 @@
-import gradio as gr
 import torch
 import torch.nn.functional as F
 import numpy as np
 import os
 import time
+import gradio as gr
 import cv2
 from PIL import Image
 from model.CyueNet_models import MMS
 from utils1.data import transform_image
 
-# 设置主题颜色和样式
+# 设置GPU/CPU
+device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+
+# 自定义CSS
 custom_css = """
     .gradio-container {
         background: linear-gradient(to right, #f6f8fa, #ffffff);
@@ -33,10 +36,77 @@ custom_css = """
     }
 """
 
-# [previous model loading and processing functions remain the same]
+def load_model():
+    """加载预训练的模型"""
+    model = MMS()
+    try:
+        model.load_state_dict(torch.load('models/CyueNet_EORSSD6.pth.54', map_location=device))
+        print("模型加载成功")
+    except RuntimeError as e:
+        print(f"加载状态字典时出现部分不匹配，错误信息: {e}")
+    model.to(device)
+    model.eval()
+    return model
+
+def process_image(image, model, threshold=0.5, testsize=256):
+    """处理图像并返回显著性检测结果"""
+    if image is None:
+        return None, None, None, None, None, "请提供有效的图像"
+    
+    # 保存原始图像
+    original_image = image.copy()
+    
+    # 预处理图像
+    image_pil = Image.fromarray(image).convert('RGB')
+    image_tensor = transform_image(image_pil, testsize)
+    image_tensor = image_tensor.unsqueeze(0)
+    image_tensor = image_tensor.to(device)
+    
+    # 计时
+    time_start = time.time()
+    
+    # 推理
+    with torch.no_grad():
+        x1, res, s1_sig, edg1, edg_s, s2, e2, s2_sig, e2_sig, s3, e3, s3_sig, e3_sig, s4, e4, s4_sig, e4_sig, s5, e5, s5_sig, e5_sig, sk1, sk1_sig, sk2, sk2_sig, sk3, sk3_sig, sk4, sk4_sig, sk5, sk5_sig = model(image_tensor)
+    
+    time_end = time.time()
+    inference_time = time_end - time_start
+    
+    # 处理输出结果
+    res = res.sigmoid().data.cpu().numpy().squeeze()
+    res = (res - res.min()) / (res.max() - res.min() + 1e-8)
+    
+    # 调整大小
+    h, w = original_image.shape[:2]
+    res_resized = cv2.resize(res, (w, h))
+    
+    # 应用阈值
+    res_vis = (res_resized * 255).astype(np.uint8)
+    
+    # 创建热力图
+    heatmap = cv2.applyColorMap(res_vis, cv2.COLORMAP_JET)
+    
+    # 叠加结果
+    alpha = 0.5
+    if len(original_image.shape) == 3 and original_image.shape[2] == 3:
+        original_bgr = cv2.cvtColor(original_image, cv2.COLOR_RGB2BGR)
+    else:
+        original_bgr = cv2.cvtColor(original_image, cv2.COLOR_GRAY2BGR)
+    
+    overlayed = cv2.addWeighted(original_bgr, 1-alpha, heatmap, alpha, 0)
+    
+    # 使用阈值进行二值化
+    _, binary_mask = cv2.threshold(res_vis, int(255 * threshold), 255, cv2.THRESH_BINARY)
+    segmented = cv2.bitwise_and(original_bgr, original_bgr, mask=binary_mask)
+    
+    # 转回RGB格式
+    overlayed_rgb = cv2.cvtColor(overlayed, cv2.COLOR_BGR2RGB)
+    segmented_rgb = cv2.cvtColor(segmented, cv2.COLOR_BGR2RGB)
+    
+    return original_image, res_vis, heatmap, overlayed_rgb, segmented_rgb, f"推理时间: {inference_time:.4f}秒"
 
 def run_demo(input_image, threshold):
-    """增强的Gradio界面主函数"""
+    """Gradio界面的主函数"""
     if input_image is None:
         return [None] * 6 + ["请上传图片"]
     
@@ -46,12 +116,19 @@ def run_demo(input_image, threshold):
     )
     
     # 计算检测区域占比
-    mask_area = np.sum(saliency_map > 127) / (saliency_map.shape[0] * saliency_map.shape[1])
-    area_info = f"检测区域占比: {mask_area:.2%}"
+    if saliency_map is not None:
+        mask_area = np.sum(saliency_map > 127) / (saliency_map.shape[0] * saliency_map.shape[1])
+        area_info = f"检测区域占比: {mask_area:.2%}"
+    else:
+        area_info = "无法计算区域占比"
     
     return original, saliency_map, heatmap, overlayed, segmented, time_info, area_info
 
-# 创建增强的Gradio界面
+# 加载模型
+print("正在加载模型...")
+model = load_model()
+
+# 创建Gradio界面
 with gr.Blocks(title="高级显著性目标检测系统", css=custom_css) as demo:
     gr.Markdown(
         """