Update app.py

app.py

@@ -27,18 +27,21 @@ def load_model():
 
 def process_image(image, model, testsize=256):
     """处理图像并返回显著性检测结果"""
+    # 保存原始图像用于后续处理
+    original_image = image.copy()
+    
     # 预处理图像
-    image = Image.fromarray(image).convert('RGB')
-    image = transform_image(image, testsize)
-    image = image.unsqueeze(0)
-    image = image.to(device)
+    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)
+        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
@@ -48,7 +51,6 @@ def process_image(image, model, testsize=256):
     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))
     
@@ -60,13 +62,23 @@ def process_image(image, model, testsize=256):
     
     # 将热力图与原始图像混合
     alpha = 0.5
-    overlayed = cv2.addWeighted(original_image, 1-alpha, heatmap, alpha, 0)
+    # 确保原始图像是BGR格式用于OpenCV操作
+    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, 127, 255, cv2.THRESH_BINARY)
-    segmented = cv2.bitwise_and(original_image, original_image, mask=binary_mask)
+    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, segmented, f"推理时间: {inference_time:.4f}秒"
+    return original_image, res_vis, heatmap, overlayed_rgb, segmented_rgb, f"推理时间: {inference_time:.4f}秒"
 
 def run_demo(input_image):
     """Gradio界面的主函数"""
@@ -112,4 +124,5 @@ with gr.Blocks(title="显著性目标检测Demo") as demo:
     gr.Markdown("3. 系统将显示原始图像、显著性图、热力图、叠加结果和分割结果")
 
 # 启动Gradio应用
-demo.launch()
+if __name__ == "__main__":
+    demo.launch(share=True)