Rename app.py to appxxxx.py

app.py → appxxxx.py

@@ -149,50 +149,64 @@ demo.launch()
 
 
 
-def generate_heatmap(last_conv_output, pred_class):
+import torch.nn.functional as F
+
+def generate_heatmap(model, last_conv_layer, input_tensor, pred_class):
     """
-    生成局部化的热力图。
+    使用 Grad-CAM 生成更精确的热力图。
     """
-
-    class_conv_output = last_conv_output[0, pred_class]
-
-
-    heatmap = torch.mean(class_conv_output, dim=0)
-
-
+    # 获取最后卷积层的特征图和模型输出
+    model.eval()
+    features = []
+    def hook_function(module, input, output):
+        features.append(output)
+
+    hook = last_conv_layer.register_forward_hook(hook_function)
+    model_output = model(input_tensor)
+    hook.remove()
+
+    # 获取预测类别的得分
+    class_score = model_output[:, pred_class]
+
+    # 计算梯度
+    model.zero_grad()
+    class_score.backward(retain_graph=True)
+
+    # 获取梯度和特征图
+    gradients = features[0].grad  # 获取最后卷积层的梯度
+    pooled_gradients = torch.mean(gradients, dim=[0, 2, 3])  # 全局平均池化
+
+    # 获取特征图，并应用梯度加权
+    features = features[0]
+    for i in range(features.shape[1]):
+        features[:, i, :, :] *= pooled_gradients[i]
+
+    # 创建热力图
+    heatmap = torch.mean(features, dim=1).squeeze()
     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):
+
+def overlay_heatmap(image, heatmap, intensity=0.5, threshold=0.5):
     """
-    调整热力图叠加到原始图像上的逻辑。
+    优化热力图叠加逻辑。
     """
-    # 调整热力图大小
     heatmap = cv2.resize(heatmap, (image.shape[1], image.shape[0]))
     heatmap = np.uint8(255 * heatmap)
+    heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
     
-    # 应用颜色映射
-    heatmap = cv2.applyColorMap(heatmap, colormap)
-    
-    # 创建一个只有热力图区域的蒙版
-    mask = heatmap > 0
-
-    # 创建叠加图像
+    # 只在高信心区域应用热力图
+    mask = heatmap > np.max(heatmap) * threshold
     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):
     """
     預測並生成熱力圖。