feat: add new SAGA attack with attention capture

app.py

@@ -6,7 +6,7 @@ from typing import Optional, List, Tuple
 from utils.model_loader import load_model_and_labels
 from utils.preprocessing import get_default_transform, preprocess_image
 from utils.inference import predict_topk
-from utils.attacks import PGDIterations, FGSM
+from utils.attacks import PGDIterations, FGSM, SAGA
 from utils.visualization import extract_attention_maps, attention_rollout, create_attention_overlay, extract_attention_for_iterations, create_iteration_attention_overlays
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -224,6 +224,8 @@ def run_attack(
         # Configurar ataque baseado no tipo selecionado
         if attack_type == "FGSM":
             attack = FGSM(model, eps=eps)
+        elif attack_type == "SAGA":
+            attack = SAGA(model, eps=eps, steps=steps)
         else:  # PGD
             attack = PGDIterations(model, eps=eps, alpha=alpha, steps=steps)
         
@@ -233,13 +235,19 @@ def run_attack(
         # Executar ataque
         adv_tensor, iteration_images = attack(img_tensor, original_label)
         
-        # Extrair atenção para todas as iterações (incluindo original)
-        attention_masks = extract_attention_for_iterations(
-            model,
-            attack.iteration_tensors,
-            discard_ratio=discard_ratio,
-            head_fusion=head_fusion
-        )
+        # Extrair atenção para todas as iterações
+        # SAGA já calcula atenção internamente, podemos reutilizar!
+        if attack_type == "SAGA" and hasattr(attack, 'attention_masks_cache') and len(attack.attention_masks_cache) > 0:
+            # Usar atenção já calculada pelo SAGA
+            attention_masks = attack.attention_masks_cache
+        else:
+            # Para FGSM e PGD, calcular atenção normalmente
+            attention_masks = extract_attention_for_iterations(
+                model,
+                attack.iteration_tensors,
+                discard_ratio=discard_ratio,
+                head_fusion=head_fusion
+            )
         
         # Criar overlays de atenção
         attention_overlays = create_iteration_attention_overlays(
@@ -267,7 +275,10 @@ def run_attack(
         if attack_type == "PGD":
             result += f"- Alpha (α): {alpha:.4f}\n"
             result += f"- Steps: {steps}\n"
-        else:
+        elif attack_type == "SAGA":
+            result += f"- Steps: {steps}\n"
+            result += f"- Gradiente ponderado por atenção (ViT-specific)\n"
+        else:  # FGSM
             result += f"- Single-step (sem iterações)\n"
         
         result += f"\n**Predição Original:**\n"
@@ -364,10 +375,10 @@ def create_app():
                         gr.Markdown("#### ⚔️ Configuração do Ataque")
                         
                         attack_type = gr.Dropdown(
-                            choices=["PGD", "FGSM"],
+                            choices=["PGD", "FGSM", "SAGA"],
                             value="PGD",
                             label="Tipo de Ataque",
-                            info="PGD: iterativo (múltiplos steps) | FGSM: single-step (mais rápido)"
+                            info="PGD: iterativo | FGSM: single-step | SAGA: gradient × attention (ViT-specific)"
                         )
                         
                         eps_input = gr.Slider(
@@ -400,7 +411,7 @@ def create_app():
                         def update_attack_params(attack_type):
                             if attack_type == "FGSM":
                                 return gr.update(visible=False)
-                            else:  # PGD
+                            else:  # PGD ou SAGA
                                 return gr.update(visible=True)
                         
                         attack_type.change(

utils/attacks.py

@@ -191,4 +191,130 @@ class PGDIterations(torchattacks.PGD):
 
         # Retornar imagem normalizada para o modelo
         adv_images = (adv_images_denorm - mean) / std
+        return adv_images, self.iteration_images
+
+
+class SAGA(torch.nn.Module):
+    """
+    SAGA: Self-Attention Gradient Attack
+    
+    Ataque adversarial específico para Vision Transformers que multiplica
+    o gradiente FGSM pelo mapa de atenção do modelo, focando perturbações
+    nas regiões que o modelo considera importantes.
+    
+    Baseado em: https://github.com/MetaMain/ViTRobust
+    Paper: "On the Robustness of Vision Transformers to Adversarial Examples" (ICCV 2021)
+    """
+    def __init__(self, model, eps=0.03, steps=10):
+        super().__init__()
+        self.model = model
+        self.eps = eps
+        self.steps = steps
+        self.device = next(model.parameters()).device
+        self.iteration_images: List[Image.Image] = []
+        self.iteration_tensors: List[torch.Tensor] = []
+        self.attention_masks_cache: List[np.ndarray] = []  # Cache das máscaras de atenção
+    
+    def get_attention_map(self, images: torch.Tensor, save_for_viz: bool = False) -> tuple:
+        """
+        Extrai mapa de atenção do ViT usando attention rollout.
+        Retorna:
+        - mask_tensor: [B, C, H, W] para uso no ataque
+        - mask_np: [H, W] numpy array para visualização (se save_for_viz=True)
+        """
+        from utils.visualization import extract_attention_maps, attention_rollout
+        import cv2
+        
+        batch_size = images.shape[0]
+        img_size = images.shape[2]
+        
+        # Extrair attention maps
+        attentions = extract_attention_maps(self.model, images)
+        
+        # Aplicar attention rollout
+        mask = attention_rollout(attentions, discard_ratio=0.9, head_fusion='max')
+        
+        # Salvar para visualização se necessário
+        if save_for_viz:
+            self.attention_masks_cache.append(mask.copy())
+        
+        # Redimensionar para tamanho da imagem (14x14 -> 224x224)
+        mask_resized = cv2.resize(mask, (img_size, img_size))
+        
+        # Expandir para 3 canais e batch: [H, W] -> [B, C, H, W]
+        mask_tensor = torch.from_numpy(mask_resized).float().to(self.device)
+        mask_tensor = mask_tensor.unsqueeze(0).unsqueeze(0)  # [1, 1, H, W]
+        mask_tensor = mask_tensor.repeat(batch_size, 3, 1, 1)  # [B, 3, H, W]
+        
+        return mask_tensor, mask if save_for_viz else None
+    
+    def forward(self, images, labels) -> Tuple[torch.Tensor, List[Image.Image]]:
+        """
+        Executa ataque SAGA e retorna:
+        - adv_images: tensor adversarial final
+        - iteration_images: lista de PIL Images de cada iteração
+        """
+        images = images.clone().detach().to(self.device)
+        labels = labels.clone().detach().to(self.device)
+        
+        loss_fn = torch.nn.CrossEntropyLoss()
+        
+        # Desnormalizar
+        mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).to(self.device)
+        std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(self.device)
+        
+        images_denorm = images * std + mean
+        adv_images_denorm = images_denorm.clone().detach()
+        
+        self.iteration_images = []
+        self.iteration_tensors = []
+        self.attention_masks_cache = []
+        
+        # Salvar imagem original (iteração 0)
+        pil_img_orig = tensor_to_pil(images_denorm[0], denormalize=False)
+        self.iteration_images.append(pil_img_orig)
+        self.iteration_tensors.append(images.clone().detach())
+        
+        # Calcular atenção para imagem original e salvar
+        _, _ = self.get_attention_map(images, save_for_viz=True)
+        
+        # Calcular eps_step
+        eps_step = self.eps / self.steps
+        
+        for step in range(self.steps):
+            # Normalizar para passar pelo modelo
+            adv_images = (adv_images_denorm - mean) / std
+            adv_images.requires_grad = True
+            
+            # Forward pass
+            outputs = self.model(adv_images)
+            
+            # Calcular loss
+            cost = loss_fn(outputs, labels)
+            
+            # Calcular gradiente
+            self.model.zero_grad()
+            cost.backward()
+            grad = adv_images.grad.data
+            
+            # SAGA: Multiplicar gradiente pelo mapa de atenção (e salvar máscara)
+            attention_map, _ = self.get_attention_map(adv_images.detach(), save_for_viz=True)
+            grad_weighted = grad * attention_map
+            
+            # Aplicar perturbação no espaço desnormalizado
+            adv_images_denorm = adv_images_denorm + eps_step * grad_weighted.sign() * std
+            
+            # Clip para [0, 1]
+            adv_images_denorm = torch.clamp(adv_images_denorm, min=0, max=1).detach()
+            
+            # Normalizar para salvar tensor
+            adv_images_normalized = (adv_images_denorm - mean) / std
+            
+            # Salvar iteração
+            pil_img = tensor_to_pil(adv_images_denorm[0], denormalize=False)
+            self.iteration_images.append(pil_img)
+            self.iteration_tensors.append(adv_images_normalized.clone().detach())
+        
+        # Retornar normalizado
+        adv_images = (adv_images_denorm - mean) / std
         return adv_images, self.iteration_images