Delete visualize_attention.py with huggingface_hub

visualize_attention.py

@@ -1,178 +0,0 @@
-import os
-import torch
-import torch.nn.functional as F
-import decord
-import numpy as np
-import matplotlib.pyplot as plt
-import cv2
-from transformers import AutoModel, AutoConfig
-import torchvision.transforms.v2 as T
-import warnings
-warnings.filterwarnings("ignore")
-
-os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
-
-# ================= 配置 =================
-import glob
-video_files = glob.glob("/root/hri30/train/*/*.avi")
-if len(video_files) > 0:
-    idx = min(50, len(video_files)-1)
-    VIDEO_PATH = video_files[idx]
-else:
-    VIDEO_PATH = "" 
-
-CKPT_PATH = "/root/autodl-tmp/checkpoints_final/final_sota_best.pth"
-MODEL_ID = "OpenGVLab/VideoMAEv2-giant"
-CACHE_DIR = "/root/autodl-tmp/hf_cache"
-
-NUM_FRAMES = 16
-IMG_SIZE = 224
-
-# ================= 模型定义 (智能 Hook) =================
-class DualHeadMAE(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-        v_config = AutoConfig.from_pretrained(MODEL_ID, trust_remote_code=True, cache_dir=CACHE_DIR)
-        v_config.use_cache = False
-        self.visual = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True, config=v_config, cache_dir=CACHE_DIR, torch_dtype=torch.float32)
-        self.attention_map = None
-        self._register_hooks()
-
-    def _register_hooks(self):
-        def hook_fn(module, input, output):
-            self.attention_map = output.detach()
-
-        target_module = None
-        # 优先找 attn_drop
-        for name, module in self.visual.named_modules():
-            if "attn_drop" in name:
-                target_module = module
-        
-        if target_module is not None:
-            target_module.register_forward_hook(hook_fn)
-            print("✅ Hooked Attention Layer")
-
-    def forward(self, x):
-        _ = self.visual(x)
-        return self.attention_map
-
-# ================= 图像处理 =================
-def get_attention_map(model, video_tensor):
-    model.eval()
-    with torch.no_grad():
-        _ = model(video_tensor)
-    
-    att_mat = model.attention_map
-    if att_mat is None: return None
-
-    # [B, Heads, N, N] -> Mean Heads -> [B, N, N]
-    if att_mat.dim() == 4:
-        att_mat = torch.mean(att_mat, dim=1)
-    
-    # 获取 [CLS] 的 attention
-    # 假设第0个是CLS
-    # 如果 N=2048 (无CLS?) 或者 N=2049 (有CLS)
-    seq_len = att_mat.shape[-1]
-    
-    # 尝试取第0行
-    cls_attn = att_mat[:, 0, :] # [B, N]
-    
-    # 如果包含自己，去掉自己
-    # 这里我们做一个简单的处理：直接用全部
-    # 归一化
-    cls_attn = (cls_attn - cls_attn.min()) / (cls_attn.max() - cls_attn.min())
-    return cls_attn
-
-def visualize(video_path, save_path="attention_vis.png"):
-    if not os.path.exists(video_path): return
-    print(f"🎥 Video: {video_path}")
-    
-    # 读取
-    vr = decord.VideoReader(video_path)
-    idx = torch.linspace(0, len(vr)-1, NUM_FRAMES).long()
-    batch = vr.get_batch(idx).asnumpy()
-    
-    # 预处理
-    buffer = torch.from_numpy(batch).permute(0, 3, 1, 2).float()
-    transform = T.Compose([T.Resize((IMG_SIZE, IMG_SIZE), antialias=True)])
-    buffer = transform(buffer)
-    mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
-    std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
-    norm_buffer = (buffer / 255.0 - mean) / std
-    input_tensor = norm_buffer.permute(1, 0, 2, 3).unsqueeze(0).cuda()
-    
-    # 推理
-    model = DualHeadMAE().cuda()
-    try:
-        sd = torch.load(CKPT_PATH)
-        # 只加载 visual
-        new_sd = {}
-        for k, v in sd.items():
-            if "visual" in k: new_sd[k.replace("visual.", "visual.")] = v 
-            elif "backbone" in k: new_sd[k.replace("backbone.", "visual.")] = v
-        model.load_state_dict(new_sd, strict=False)
-        print("✅ Weights Loaded")
-    except:
-        print("⚠️ Random Weights")
-
-    model.eval()
-    attn_score = get_attention_map(model, input_tensor) # [1, N]
-    
-    # 🔥🔥🔥 暴力 Reshape 修复 🔥🔥🔥
-    num_tokens = attn_score.shape[1]
-    print(f"Tokens: {num_tokens}")
-    
-    # 目标：变成 [T, H, W]
-    # 我们知道 T=8 (16/2)
-    # 剩下的 spatial_tokens = num_tokens / 8
-    
-    # 假设有 CLS，先去掉一个看看能不能整除
-    if num_tokens % 8 != 0:
-        attn_score = attn_score[:, 1:] # 丢掉第一个
-        num_tokens -= 1
-    
-    spatial = num_tokens // 8
-    h = int(np.sqrt(spatial))
-    w = h
-    
-    print(f"Reshaping to [8, {h}, {w}]")
-    
-    try:
-        attn_score = attn_score.reshape(8, h, w)
-    except:
-        # 实在不行，硬插值
-        print("⚠️ Shape mismatch, forcing interpolation...")
-        attn_score = F.interpolate(attn_score.unsqueeze(0), size=8*14*14, mode='linear').reshape(8, 14, 14)
-
-    # 插值回视频尺寸
-    attn_score = F.interpolate(attn_score.unsqueeze(0).unsqueeze(0), size=(16, 224, 224), mode='trilinear').squeeze()
-    attn_score = attn_score.cpu().numpy()
-    
-    # 绘图
-    frame_indices = [2, 6, 10, 14]
-    fig, axes = plt.subplots(2, 4, figsize=(16, 8))
-    orig_imgs = F.interpolate(torch.from_numpy(batch).permute(0,3,1,2).float(), size=(224,224)).permute(0,2,3,1).numpy().astype(np.uint8)
-
-    for i, frame_idx in enumerate(frame_indices):
-        img = orig_imgs[frame_idx]
-        heatmap = attn_score[frame_idx]
-        heatmap = (heatmap - heatmap.min()) / (heatmap.max() - heatmap.min() + 1e-8)
-        heatmap = np.uint8(255 * heatmap)
-        heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
-        overlay = cv2.addWeighted(img, 0.6, heatmap, 0.4, 0)
-        
-        axes[0, i].imshow(img)
-        axes[0, i].axis('off')
-        axes[0, i].set_title(f"Frame {frame_idx}")
-        
-        axes[1, i].imshow(cv2.cvtColor(overlay, cv2.COLOR_BGR2RGB))
-        axes[1, i].axis('off')
-        axes[1, i].set_title(f"Attention")
-        
-    plt.tight_layout()
-    plt.savefig(save_path)
-    print(f"✅ Saved: {save_path}")
-
-if __name__ == "__main__":
-    if VIDEO_PATH: visualize(VIDEO_PATH)