Upload 3 files

README.md

@@ -1,12 +1,26 @@
 ---
-title: Saliency Detection Demo
-emoji: 😻
-colorFrom: purple
-colorTo: pink
+title: 显著性目标检测Demo
+emoji: 🔍
+colorFrom: blue
+colorTo: red
 sdk: gradio
-sdk_version: 5.27.0
+sdk_version: 3.50.2
 app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# 显著性目标检测Demo
+
+这个应用使用CyueNet模型进行显著性目标检测。上传一张图片，系统将自动检测并高亮显示图像中的显著性区域。
+
+## 功能
+
+- 上传图像进行显著性检测
+- 显示原始图像、显著性图、热力图、叠加结果和分割结果
+- 实时处理，显示推理时间
+
+## 使用方法
+
+1. 点击"输入图像"区域上传一张图片
+2. 点击"开始检测"按钮进行显著性目标检测
+3. 查看检测结果

app.py

@@ -0,0 +1,115 @@
+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')
+
+def load_model():
+    """加载预训练的模型"""
+    model = MMS()
+    try:
+        # 使用相对路径，模型文件将存储在HuggingFace Spaces上
+        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, testsize=256):
+    """处理图像并返回显著性检测结果"""
+    # 预处理图像
+    image = Image.fromarray(image).convert('RGB')
+    image = transform_image(image, testsize)
+    image = image.unsqueeze(0)
+    image = image.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)
+    
+    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)
+    
+    # 将输出调整为原始图像大小
+    original_image = np.array(Image.fromarray(image.cpu().squeeze().permute(1, 2, 0).numpy()))
+    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
+    overlayed = cv2.addWeighted(original_image, 1-alpha, heatmap, alpha, 0)
+    
+    # 二值化结果用于分割
+    _, binary_mask = cv2.threshold(res_vis, 127, 255, cv2.THRESH_BINARY)
+    segmented = cv2.bitwise_and(original_image, original_image, mask=binary_mask)
+    
+    return original_image, res_vis, heatmap, overlayed, segmented, f"推理时间: {inference_time:.4f}秒"
+
+def run_demo(input_image):
+    """Gradio界面的主函数"""
+    if input_image is None:
+        return [None] * 5 + ["请上传图片"]
+    
+    # 处理图像
+    original, saliency_map, heatmap, overlayed, segmented, time_info = process_image(input_image, model)
+    
+    return original, saliency_map, heatmap, overlayed, segmented, time_info
+
+# 加载模型
+print("正在加载模型...")
+model = load_model()
+
+# 创建Gradio界面
+with gr.Blocks(title="显著性目标检测Demo") as demo:
+    gr.Markdown("# 显著性目标检测Demo")
+    gr.Markdown("上传一张图片，系统将自动检测显著性区域")
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(label="输入图像", type="numpy")
+            submit_btn = gr.Button("开始检测")
+        
+        with gr.Column():
+            original_output = gr.Image(label="原始图像")
+            saliency_output = gr.Image(label="显著性图")
+            heatmap_output = gr.Image(label="热力图")
+            overlayed_output = gr.Image(label="叠加结果")
+            segmented_output = gr.Image(label="分割结果")
+            time_info = gr.Textbox(label="处理信息")
+    
+    submit_btn.click(
+        fn=run_demo,
+        inputs=input_image,
+        outputs=[original_output, saliency_output, heatmap_output, overlayed_output, segmented_output, time_info]
+    )
+    
+    gr.Markdown("## 使用说明")
+    gr.Markdown("1. 点击'输入图像'区域上传一张图片")
+    gr.Markdown("2. 点击'开始检测'按钮进行显著性目标检测")
+    gr.Markdown("3. 系统将显示原始图像、显著性图、热力图、叠加结果和分割结果")
+
+# 启动Gradio应用
+demo.launch()

requirements.txt

@@ -0,0 +1,6 @@
+torch
+torchvision
+numpy
+opencv-python
+pillow
+gradio