attacks.py

"""
attacks.py

提供对检测模型（以 YOLOv8/ultralytics 为主）执行 FGSM 与 PGD 的实现。

设计思路与注意事项：
- 假定我们可以访问到底层的 torch.nn.Module（例如 ultralytics.YOLO 实例的 .model 成员）
  并能以 tensor 输入直接跑 forward()，得到原始预测张量 (batch, N_preds, C)
  其中通常 C = 5 + num_classes（bbox4 + obj_conf + class_logits）。
- 计算 loss: 对每个 anchor/pred，取 obj_conf * max_class_prob 作为该预测的置信度，
  把全局置信度求和作为被攻击的目标函数；对该目标函数**做最小化**以让检测置信下降。
  - FGSM: x_adv = x - eps * sign(grad(loss))
  - PGD: 多步迭代，每步做 x = x - alpha * sign(grad), 并投影到 L_inf 球体：|x-x_orig|<=eps
- 如果你的 ultralytics 版本/模型封装与假定不同，代码会抛错并提示如何修改。
"""

from typing import Tuple, Optional
import torch
import torch.nn as nn
import numpy as np
from PIL import Image
import torchvision.transforms as T
import math
import torch.nn.functional as F
from typing import Tuple, Dict

# ============= Resize image =====================
def _get_max_stride(net) -> int:
    s = getattr(net, "stride", None)
    if isinstance(s, torch.Tensor):
        return int(s.max().item())
    try:
        return int(max(s))
    except Exception:
        return 32  # 兜底

def letterbox_tensor(
    x: torch.Tensor,
    *,
    imgsz: int,
    stride: int,
    fill: float = 114.0 / 255.0,
    scaleup: bool = True
) -> Tuple[torch.Tensor, Dict]:
    """
    x: [1,3,H,W] in [0,1]
    返回: x_lb, meta （含缩放比例与左右上下 padding 以便反映射）
    """
    assert x.ndim == 4 and x.shape[0] == 1
    _, C, H, W = x.shape
    if imgsz is None:
        # 动态设定目标边长：把 max(H,W) 向上取整到 stride 的倍数（更贴近原生自动整形）
        imgsz = int(math.ceil(max(H, W) / stride) * stride)

    r = min(imgsz / H, imgsz / W)  # 等比缩放比例
    if not scaleup:
        r = min(r, 1.0)

    new_w = int(round(W * r))
    new_h = int(round(H * r))

    # 先等比缩放
    if (new_h, new_w) != (H, W):
        x = F.interpolate(x, size=(new_h, new_w), mode="bilinear", align_corners=False)

    # 再 padding 到 (imgsz, imgsz)，保持与 YOLO 一致的对称填充
    dw = imgsz - new_w
    dh = imgsz - new_h
    left, right = dw // 2, dw - dw // 2
    top, bottom = dh // 2, dh - dh // 2

    x = F.pad(x, (left, right, top, bottom), mode="constant", value=fill)

    meta = {
        "ratio": r,
        "pad": (left, top),
        "resized_shape": (new_h, new_w),
        "imgsz": imgsz,
    }
    return x, meta

def unletterbox_to_original(
    x_lb: torch.Tensor, meta: Dict, orig_hw: Tuple[int, int]
) -> torch.Tensor:
    """
    把 letterboxed 张量（[1,3,imgsz,imgsz]）反映射回原始 H0,W0 尺寸（去 padding + 反缩放）
    """
    assert x_lb.ndim == 4 and x_lb.shape[0] == 1
    H0, W0 = orig_hw
    (left, top) = meta["pad"]
    (h_r, w_r) = meta["resized_shape"]

    # 去 padding（裁出等比缩放后的区域）
    x_unpad = x_lb[..., top:top + h_r, left:left + w_r]  # [1,3,h_r,w_r]

    # 反缩放到原图大小
    x_rec = F.interpolate(x_unpad, size=(H0, W0), mode="bilinear", align_corners=False)
    return x_rec

  
# ----- basic preprocessing / deprocessing (RGB PIL <-> torch tensor) -----
_to_tensor = T.Compose([
    T.ToTensor(),  # float in [0,1], shape C,H,W
])

_to_pil = T.ToPILImage()

def pil_to_tensor(img_pil: Image.Image, device: torch.device) -> torch.Tensor:
    """PIL RGB -> float tensor [1,3,H,W] on device"""
    t = _to_tensor(img_pil).unsqueeze(0).to(device)  # 1,C,H,W
    t.requires_grad = True
    return t

def tensor_to_pil(t: torch.Tensor) -> Image.Image:
    """tensor [1,3,H,W] (0..1) -> PIL RGB"""
    t = t.detach().cpu().squeeze(0).clamp(0.0, 1.0)
    return _to_pil(t)

# ----- helper to obtain underlying torch module from ultralytics YOLO wrapper -----
def get_torch_module_from_ultralytics(model) -> nn.Module:
    """
    Try to retrieve an nn.Module that accepts an input tensor and returns raw preds.
    For ultralytics.YOLO, .model is usually the underlying Detect/Model (nn.Module).
    """
    if hasattr(model, "model") and isinstance(model.model, nn.Module):
        return model.model
    # Some wrappers nest further; attempt a few common names
    for attr in ("model", "module", "net", "model_"):
        if hasattr(model, attr) and isinstance(getattr(model, attr), nn.Module):
            return getattr(model, attr)
    raise RuntimeError("无法找到底层 torch.nn.Module。请确保传入的是 ultralytics.YOLO 实例且能访问 model.model。")

# ----- interpret raw model outputs to confidences -----


def _ensure_bcn(preds):
    assert preds.ndim == 3
    B, C1, C2 = preds.shape
    if C1 - 4 > 0 and C2 >= 1000:   # [B, 4+nc, N]
        return preds
    if C2 - 4 > 0 and C1 >= 1000:   # [B, N, 4+nc]
        return preds.permute(0, 2, 1).contiguous()
    return preds

def _xywh_to_xyxy(xywh):
    x,y,w,h = xywh.unbind(-1)
    return torch.stack([x-w/2, y-h/2, x+w/2, y+h/2], dim=-1)

def _xyxy_to_xywh(xyxy):
    x1,y1,x2,y2 = xyxy.unbind(-1)
    cx = (x1+x2)/2; cy = (y1+y2)/2
    w  = (x2-x1).clamp(min=0); h = (y2-y1).clamp(min=0)
    return torch.stack([cx,cy,w,h], dim=-1)

def _map_xyxy_to_letterbox(xyxy_tensor, meta):
    if meta is None:
        return xyxy_tensor
    r = meta.get('ratio', meta.get('scale', (1.0, 1.0)))
    p = meta.get('pad', (0.0, 0.0))
    if isinstance(r, (int, float)):
        r = (float(r), float(r))
    rx, ry = float(r[0]), float(r[1])
    px, py = float(p[0]), float(p[1])
    x1 = xyxy_tensor[:, 0] * rx + px
    y1 = xyxy_tensor[:, 1] * ry + py
    x2 = xyxy_tensor[:, 2] * rx + px
    y2 = xyxy_tensor[:, 3] * ry + py
    return torch.stack([x1, y1, x2, y2], dim=-1)

def _iou_xyxy(b_xyxy, g_xyxy):
    N, M = b_xyxy.size(0), g_xyxy.size(0)
    b = b_xyxy[:, None, :].expand(N, M, 4)
    g = g_xyxy[None, :, :].expand(N, M, 4)
    inter_x1 = torch.maximum(b[...,0], g[...,0])
    inter_y1 = torch.maximum(b[...,1], g[...,1])
    inter_x2 = torch.minimum(b[...,2], g[...,2])
    inter_y2 = torch.minimum(b[...,3], g[...,3])
    inter_w  = (inter_x2 - inter_x1).clamp(min=0)
    inter_h  = (inter_y2 - inter_y1).clamp(min=0)
    inter    = inter_w * inter_h
    area_b   = (b[...,2]-b[...,0]).clamp(min=0) * (b[...,3]-b[...,1]).clamp(min=0)
    area_g   = (g[...,2]-g[...,0]).clamp(min=0) * (g[...,3]-g[...,1]).clamp(min=0)
    return inter / (area_b + area_g - inter + 1e-9)

def _gt_list_to_xyxy_tensor(gt_list, device, meta=None):
    if not gt_list:
        return torch.empty(0, 4, device=device, dtype=torch.float32)
    xyxy = torch.tensor([b['xyxy'] for b in gt_list], dtype=torch.float32, device=device)
    return _map_xyxy_to_letterbox(xyxy, meta)

def preds_to_targeted_loss(
    preds,                  # [B,4+nc,N] 或 [B,N,4+nc]；类别部分最好是 logits
    target_cls: int,
    gt_xywh,                # 这里直接支持 list[{'xyxy':..., 'cls':..., 'conf':...}]
    topk: int = 20,
    kappa: float = 0.1,
    lambda_margin: float = 1.0,
    lambda_keep: float = 0.2,
    lambda_target: float = 0.0,    # 新增：恢复 -p_t.mean() 这项
    debug: bool = False,
    meta: dict | None = None,      # 若 GT 是原图坐标，传入 letterbox 的 meta
):
    preds = _ensure_bcn(preds)
    B, C, N = preds.shape
    nc = C - 4
    assert 0 <= target_cls < nc

    # 解析 GT（list -> tensor in letterbox coords）
    gt_xyxy_lb = _gt_list_to_xyxy_tensor(gt_xywh, preds.device, meta=meta)  # [M,4]

    boxes_bxn4 = preds[:, :4, :].permute(0, 2, 1)   # [B,N,4] (xywh, letterbox)
    logits_bxcn = preds[:, 4:, :]                   # [B,nc,N]

    # 若类别部分像概率(0~1)，转为 logits
    zmin, zmax = logits_bxcn.min().item(), logits_bxcn.max().item()
    if 0.0 <= zmin and zmax <= 1.0:
        p = logits_bxcn.clamp(1e-6, 1-1e-6)
        logits_bxcn = torch.log(p) - torch.log1p(-p)

    # 选与 GT 最相关的候选 idx（batch=0）
    b_xyxy = _xywh_to_xyxy(boxes_bxn4[0])          # [N,4]
    if gt_xyxy_lb.numel() > 0:
        iou = _iou_xyxy(b_xyxy, gt_xyxy_lb)        # [N,M]
        best_per_gt = iou.argmax(dim=0)            # [M]
        idx = torch.unique(best_per_gt, sorted=False)
        if idx.numel() < topk:
            topvals = iou.max(dim=1).values
            topidx2 = torch.topk(topvals, k=min(topk, N)).indices
            idx = torch.unique(torch.cat([idx, topidx2], 0), sorted=False)[:topk]
    else:
        # 没 GT 就按当前最大类别置信度取 topk
        z = logits_bxcn[0]                         # [nc,N]
        pmax = z.softmax(dim=0).max(dim=0).values
        idx = torch.topk(pmax, k=min(topk, N)).indices

    if idx.numel() == 0:
        idx = torch.arange(min(topk, N), device=preds.device)

    # 取这些候选的类别 logits：[K,nc]
    z = logits_bxcn[0, :, idx].T                   # [K,nc]

    # 1) CW-style margin
    mask = torch.ones(nc, device=z.device, dtype=torch.bool)
    mask[target_cls] = False
    z_t   = z[:, target_cls]
    z_oth = z[:, mask].max(dim=1).values
    loss_margin = F.relu(kappa + z_oth - z_t).mean()

    # 2) keep（KL >= 0）
    with torch.no_grad():
        p_clean = z.detach().softmax(dim=1)
    logp_adv = z.log_softmax(dim=1)
    loss_keep = F.kl_div(logp_adv, p_clean, reduction="batchmean")

    # 3) 你的旧项：直接推高目标类 logit
    loss_target = -z_t.mean()

    loss = (
        lambda_margin * loss_margin
        + lambda_keep * loss_keep
        + lambda_target * loss_target
    )

    if debug:
        same_ratio = (z.argmax(dim=1) == target_cls).float().mean().item()
        print(
            f"[dbg] K={idx.numel()} nc={nc} target={target_cls} "
            f"margin={loss_margin.item():.6f} keep={loss_keep.item():.6f} "
            f"targ={loss_target.item():.6f} same_ratio={same_ratio:.3f} "
            f"z_t_mean={z_t.mean().item():.3f} z_oth_mean={z_oth.mean().item():.3f}"
        )
    return loss


# def preds_to_confidence_sum(preds: torch.Tensor) -> torch.Tensor:
#     """
#     preds: tensor shape (batch, N_preds, C) or (batch, C, H, W) depending on model.
#     We support the common YOLO format where last dim: [x,y,w,h,obj_conf, class_probs...]
#     Returns scalar: sum of (obj_conf * max_class_prob) over batch and predictions.
#     """
#     if preds is None:
#         raise ValueError("preds is None")
#     # handle shape (batch, N_preds, C)
#     if preds.ndim == 3:
#         # assume last dim: 5 + num_classes
#         if preds.shape[-1] < 6:
#             # can't interpret
#             raise RuntimeError(f"preds last dim too small ({preds.shape[-1]}). Expecting >=6.")
#         obj_conf = preds[..., 4]  # (batch, N)
#         class_probs = preds[..., 5:]  # (batch, N, num_cls)
#         max_class, _ = class_probs.max(dim=-1)  # (batch, N)
#         conf = obj_conf * max_class
#         return conf.sum()
#     # some models output (batch, C, H, W) - flatten
#     if preds.ndim == 4:
#         # try to collapse so that last dim is class
#         b, c, h, w = preds.shape
#         flat = preds.view(b, c, -1).permute(0, 2, 1)  # (batch, N, C)
#         return preds_to_confidence_sum(flat)
#     raise RuntimeError(f"Unhandled preds dimensionality: {preds.shape}")

# ----- core attack implementations -----
def fgsm_attack_on_detector(
    model,
    img_pil: Image.Image,
    eps: float = 0.03,
    device: Optional[torch.device] = None,
    imgsz: Optional[int] = None,  # None=自动对齐到 stride 倍数；也可传 640
    gt_xywh: torch.Tensor | None = None,  # letterbox坐标系下的目标框（可选）
    target_cls: int = 2,
) -> Image.Image:
    """
    Perform a single-step FGSM on a detection model (white-box).
    - model: ultralytics.YOLO wrapper (or anything where get_torch_module_from_ultralytics works)
    - img_pil: input PIL RGB
    - eps: max per-pixel perturbation in [0,1] (L_inf)
    Returns PIL image of adversarial example.
    """
    device = device or (torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu"))
    # get torch module
    net = get_torch_module_from_ultralytics(model)
    net = net.to(device).eval()
    for p in net.parameters():
      p.requires_grad_(False)   # 建议：避免对参数求梯度
    
    # (a) 原图 -> [1,3,H0,W0]，随后先 detach 掉梯度
    x_orig = pil_to_tensor(img_pil, device)  
    H0, W0 = x_orig.shape[-2:]
    x_orig = x_orig.detach()
    
    # (b) 可微 letterbox
    s = _get_max_stride(net)
    x_lb, meta = letterbox_tensor(x_orig, imgsz=imgsz, stride=s, fill=114/255.0)
    x_lb = x_lb.clone().detach().requires_grad_(True)

    # (c) 前向与你的损失
    with torch.enable_grad():
        preds = net(x_lb)
        if isinstance(preds, (tuple, list)):
            tensor_pred = next((p for p in preds if isinstance(p, torch.Tensor) and p.ndim >= 3), None)
            if tensor_pred is None:
                raise RuntimeError("模型 forward 返回了 tuple/list，但无法从中找到预测张量。")
            preds = tensor_pred

        loss = - preds_to_targeted_loss(
            preds,
            target_cls=target_cls,
            gt_xywh=gt_xywh,    # 直接传你的 list[dict]
            topk=20,
            kappa=0.1,
            lambda_margin=1.0,
            lambda_keep=0.2,
            lambda_target=0.0,   # 恢复 -p_t.mean() 的影响
            debug=False,
            meta=meta            # 若 GT 是原图坐标，务必传 meta
        )

        # loss = - preds_to_confidence_sum(preds)  
        loss.backward()

    # (d) FGSM 在 letterboxed 空间施扰
    # FGSM update: x_adv = x + eps * sign(grad(loss wrt x))
    with torch.no_grad():
        adv_lb = (x_lb + eps * x_lb.grad.sign()).clamp(0, 1)

    # 清理（单步可选；PGD循环时必做）
    x_lb.grad = None
    net.zero_grad(set_to_none=True)

    # (e) 反映射回原图尺寸
    adv_orig = unletterbox_to_original(adv_lb, meta, (H0, W0)).detach()

    # (f) 转回 PIL
    adv_pil = tensor_to_pil(adv_orig)
    return adv_pil

def pgd_attack_on_detector(
    model,
    img_pil: Image.Image,
    eps: float = 0.03,      # L_inf 半径（输入在[0,1]域）
    alpha: float = 0.007,   # 步长
    iters: int = 10,
    device: Optional[torch.device] = None,
    imgsz: Optional[int] = None,  # None=自动对齐到 stride 倍数；也可传 640
    gt_xywh: torch.Tensor | None = None,  # letterbox坐标系下的目标框（可选）
    target_cls: int = 2,
):
    """
    在 YOLO 的 letterbox 域做 PGD，
    迭代结束后把对抗样本映回原图大小并返回 PIL。
    依赖你已实现的: pil_to_tensor, tensor_to_pil, letterbox_tensor, unletterbox_to_original, _get_max_stride
    """
    device = device or (torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu"))
    net = get_torch_module_from_ultralytics(model).to(device).eval()

    # 仅对输入求梯度，冻结参数以省资源
    for p in net.parameters():
        p.requires_grad_(False)

    # 原图 -> Tensor（[1,3,H0,W0], [0,1]）
    x0 = pil_to_tensor(img_pil, device).detach()
    H0, W0 = x0.shape[-2:]

    # 可微 letterbox（等比缩放 + 对称 pad 到 stride 倍数）
    s = _get_max_stride(net)
    x_lb_orig, meta = letterbox_tensor(x0, imgsz=imgsz, stride=s, fill=114/255.0)  # [1,3,S,S]
    x = x_lb_orig.clone().detach().requires_grad_(True)

    # targeted = True
    # sign = -1.0 if targeted else 1.0    # 定向取负号，非定向取正号
    for _ in range(iters):
        # 前向 + 反向（需要梯度）
        preds = net(x)
        if isinstance(preds, (tuple, list)):
            preds = next((p for p in preds if isinstance(p, torch.Tensor) and p.ndim >= 3), None)
            if preds is None:
                raise RuntimeError("模型 forward 返回 tuple/list，但未找到预测张量。")

        loss = - preds_to_targeted_loss(
            preds,
            target_cls=target_cls,
            gt_xywh=gt_xywh,    # 直接传你的 list[dict]
            topk=20,
            kappa=0.1,
            lambda_margin=1.0,
            lambda_keep=0.2,
            lambda_target=0.0,   # 恢复 -p_t.mean() 的影响
            debug=False,
            meta=meta            # 若 GT 是原图坐标，务必传 meta
        )
        
        # loss = - preds_to_confidence_sum(preds)      # 我们希望置信度总和下降 → 最小化
        loss.backward()

        # 更新步与投影（不记录计算图）
        with torch.no_grad():
            x.add_(alpha * x.grad.sign())
            # 投影到 L_inf 球: 通过裁剪 delta 更稳
            delta = (x - x_lb_orig).clamp(-eps, eps)
            x.copy_((x_lb_orig + delta).clamp(0.0, 1.0))

        # 清理并设置下一步
        x.grad = None
        net.zero_grad(set_to_none=True)
        x.requires_grad_(True)

    # 反映射回原图尺寸
    adv_orig = unletterbox_to_original(x.detach(), meta, (H0, W0)).detach()
    return tensor_to_pil(adv_orig)


# ----- graceful fallback / demo noise if whitebox impossible -----
def demo_random_perturbation(img_pil: Image.Image, eps: float = 0.03) -> Image.Image:
    """Non-gradient demo perturbation used as fallback."""
    arr = np.asarray(img_pil).astype(np.float32) / 255.0
    noise = np.sign(np.random.randn(*arr.shape)).astype(np.float32)
    adv = np.clip(arr + eps * noise, 0.0, 1.0)
    adv_img = Image.fromarray((adv * 255).astype(np.uint8))
    return adv_img