Update app.py

app.py

@@ -151,36 +151,48 @@ demo.launch()
 
 def generate_heatmap(last_conv_output, pred_class):
     """
-    生成熱力圖。
+    生成局部化的热力图。
     """
-    # 獲取目標類別的輸出
-    target_output = last_conv_output[0, pred_class]
-    
-    # 執行全局平均池化以獲得類別特定的特徵圖
-    heatmap = torch.mean(target_output, dim=0)
-    
-    # 標準化熱力圖
+
+    class_conv_output = last_conv_output[0, pred_class]
+
+
+    heatmap = torch.mean(class_conv_output, dim=0)
+
+
     heatmap = np.maximum(heatmap.detach().cpu().numpy(), 0)
     heatmap /= np.max(heatmap)
     
+
+    heatmap[heatmap < 0.5] = 0  # 只保留高信心区域
+
     return heatmap
 
+
 def overlay_heatmap(image, heatmap, intensity=0.5, colormap=cv2.COLORMAP_JET):
     """
-    將熱力圖疊加到原始圖片上。
+    调整热力图叠加到原始图像上的逻辑。
     """
-    # 將熱力圖調整為與原始圖片相同的大小
-    heatmap = cv2.resize(heatmap, (image.shape[1], image.shape[0]))  # 使用 NumPy 數組的尺寸
+    # 调整热力图大小
+    heatmap = cv2.resize(heatmap, (image.shape[1], image.shape[0]))
     heatmap = np.uint8(255 * heatmap)
     
-    # 將熱力圖應用到原始圖片上
+    # 应用颜色映射
     heatmap = cv2.applyColorMap(heatmap, colormap)
-    superimposed_img = heatmap * intensity + image  # 直接使用 NumPy 數組
+    
+    # 创建一个只有热力图区域的蒙版
+    mask = heatmap > 0
+
+    # 创建叠加图像
+    superimposed_img = image.copy()
+    superimposed_img[mask] = superimposed_img[mask] * (1 - intensity) + heatmap[mask] * intensity
+    
     superimposed_img = np.clip(superimposed_img, 0, 255).astype(np.uint8)
     superimposed_img = Image.fromarray(superimposed_img)
 
     return superimposed_img
 
+
 def predict(image):
     """
     預測並生成熱力圖。