Update app.py

app.py

@@ -106,73 +106,6 @@ def build_transform(is_training: bool, interpolation: str):
     cfg = model.data_config.copy()
     cfg.update(dict(interpolation=interpolation))
     return timm.data.create_transform(**cfg, is_training=is_training)
-# 4. Attention Hook & Visualization  =========================================
-# ---------------------------------------------------------------------------
-
-def get_attention_map(module, inputs, output):
-    """
-    forward_hook —— 捕获 softmax 后、dropout 前的注意力权重
-    inputs[0] : [num_windows*B, num_heads, N, N] (N = win_size²)
-    """
-    global attention_maps
-    if inputs and isinstance(inputs[0], torch.Tensor):
-        # 只保存第一张图片第一帧即可
-        attention_maps.append(inputs[0].detach().cpu())
-
-
-def create_attention_visualization(image_pil: Image.Image,
-                                   attn_map: torch.Tensor,
-                                   patch_size: int = 4) -> Image.Image:
-    """
-    1) 把窗口注意力 → token 级热图 (H_token × W_token)
-    2) 最近邻上采样到原始 patch 网格 (96×96 for 384² 输入, patch_size=4)
-    3) 再上采样到像素级并做 blend
-    """
-    # -----------------------------------------------------------
-    # 1. 计算每个 token 的“重要性” : head + query 维均值
-    #    attn_map: [num_windows, num_heads, N, N]  (batch 已经 =1)
-    attn_map = attn_map.mean(dim=1).mean(dim=2)           # → [num_windows, N]
-    attn_map = attn_map.clamp(min=0)
-
-    num_windows, N = attn_map.shape
-    win_size = int(math.sqrt(N))                          # 12
-    assert win_size * win_size == N, "N 不是 win_size²"
-
-    # -----------------------------------------------------------
-    # 2. 先在 token 分辨率下拼一张 heat_token (H_token × W_token)
-    #    token 分辨率 = win_size × windows_per_dim
-    img_h, img_w = image_pil.size[1], image_pil.size[0]   # PIL (w,h)
-    num_patch_h, num_patch_w = img_h // patch_size, img_w // patch_size  # 96×96
-
-    win_per_row = int(round((num_patch_w / win_size)))    # 8 for Stage1, 1 for Stage4
-    token_side  = win_per_row * win_size                  # 96 or 12
-    heat_token  = torch.zeros(token_side, token_side)
-
-    for idx in range(num_windows):
-        row_w = idx // win_per_row
-        col_w = idx %  win_per_row
-        r0, r1 = row_w * win_size, (row_w + 1) * win_size
-        c0, c1 = col_w * win_size, (col_w + 1) * win_size
-        heat_token[r0:r1, c0:c1] = attn_map[idx].view(win_size, win_size)
-
-    # -----------------------------------------------------------
-    # 3. 归一化 & 上采样到 patch 网格尺寸 (96×96)
-    heat_token = heat_token.unsqueeze(0).unsqueeze(0)     # [1,1,H_t,W_t]
-    heat_patch = F.interpolate(heat_token, size=(num_patch_h, num_patch_w),
-                               mode="nearest")[0, 0]      # [H_patch,W_patch]
-
-    heat_patch -= heat_patch.min()
-    heat_patch /= (heat_patch.max() + 1e-6)
-
-    # -----------------------------------------------------------
-    # 4. 再转为像素热图并 Blend
-    heat_np = heat_patch.numpy()
-    heat_img = Image.fromarray((plt.cm.viridis(heat_np)[:, :, :3] * 255).astype(np.uint8))
-    heat_img = heat_img.resize(image_pil.size, Image.BICUBIC)
-
-    blended = Image.blend(image_pil.convert("RGB"), heat_img, alpha=0.55)
-    return blended
-
 
 # ---------------------------------------------------------------------------
 # 5. 推理 + 可选的注意力可视化
@@ -181,38 +114,20 @@ def predict_and_visualize(image_pil: Image.Image,
                           interpolation: str = "bicubic",
                           show_attention: bool = True):
     if image_pil is None:
-        return None, None
+        return None
 
     load_model(ckpt_name)
 
-    global attention_maps
-    attention_maps = []
-
     transform = build_transform(is_training=False, interpolation=interpolation)
     input_tensor = transform(image_pil).unsqueeze(0).to(device)
     
-    hook_handle = None
-    if show_attention:
-        # --- FIX: Target the attn_drop layer inside the attention module ---
-        target_layer = model.backbone.layers[-1].blocks[-1].attn.attn_drop
-        hook_handle = target_layer.register_forward_hook(get_attention_map)
-
     with torch.no_grad():
         logits = model(input_tensor)
-
-    if hook_handle:
-        hook_handle.remove()
-
+    
     probs = F.softmax(logits, dim=1)[0]
     confidences = {class_names[i]: float(probs[i]) for i in range(NUM_CLASSES)}
 
-    viz_image = None
-    if show_attention and attention_maps:
-        viz_image = create_attention_visualization(image_pil.copy(),
-                                                   attention_maps[0],
-                                                   patch_size=4)   # Swin-Large 默认 patch 4
-
-    return confidences, viz_image
+    return confidences
 
 # ---------------------------------------------------------------------------
 # 6. Gradio UI
@@ -234,18 +149,16 @@ def launch_app():
                     ["bilinear", "bicubic", "nearest"], value="bicubic",
                     label="Resize Interpolation (Preprocessing)"
                 )
-                viz_checkbox = gr.Checkbox(value=True, label="Show Attention Visualization")
                 
                 in_img = gr.Image(type="pil", label="Upload an Image")
             
             with gr.Column(scale=2):
                 out_lbl = gr.Label(num_top_classes=2, label="Predictions")
-                out_viz = gr.Image(type="pil", label="Attention Map Visualization", visible=True)
 
         run_btn.click(
             predict_and_visualize,
-            inputs=[in_img, model_choice, interp_choice, viz_checkbox],
-            outputs=[out_lbl, out_viz]
+            inputs=[in_img, model_choice, interp_choice],
+            outputs=[out_lbl]
         )
         
         # Create a dummy examples directory if it doesn't exist
@@ -258,9 +171,9 @@ def launch_app():
         example_files = [os.path.join(example_dir, f) for f in os.listdir(example_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
         if example_files:
             gr.Examples(
-                examples=[[f, DEFAULT_CKPT, "bicubic", True] for f in example_files],
-                inputs=[in_img, model_choice, interp_choice, viz_checkbox],
-                outputs=[out_lbl, out_viz],
+                examples=[[f, DEFAULT_CKPT, "bicubic"] for f in example_files],
+                inputs=[in_img, model_choice, interp_choice],
+                outputs=[out_lbl],
                 fn=predict_and_visualize,
                 cache_examples=False,
             )