attack.py

from matplotlib import image as mlt
from pathlib import Path
import torch
import random
import numpy as np
import torchattacks
import torch.nn.functional as F
import clip
import cv2
from PIL import Image
from torchvision.transforms import transforms
from constant import *
# from models.resnet_cifar import resnet18_cifar("")
to_tensor = transforms.ToTensor()
to_img = transforms.ToPILImage()
try:
    from torchvision.transforms import InterpolationMode
    BICUBIC = InterpolationMode.BICUBIC
except ImportError:
    BICUBIC = Image.BICUBIC
kitty_mask = mlt.imread('./trigger/hello_kitty.png') * 255

stub1 = [(0, 6, 7, 8), (4, 5, 6, 8), (2, 5, 6, 7), (0, 5, 6, 7), (1, 2, 3, 6), (2, 4, 5, 8), (3, 4, 6, 7), (0, 1, 2, 3), (0, 4, 5, 8), (0, 4, 7, 8)]


color = [
    [255,0,0],
    [255,255,0],
    [0,0,255],
    [0,255,0]
]

color_10 = [
    [0,0,255],
    [0,0,0],
    [255,255,255],
    [0,255,0],
    [255,0,0],
    [255,255,0],
    [255,0,255],
    [0,255,255],
    [205,92,92],
    [148,0,211]
]

color_100 = [np.random.randint(0, 256, 3) for _ in range(100)]

color_list = [(1, 3, 0, 2), (0, 3, 1, 2), (1, 0, 2, 3), (0, 2, 1, 3), (2, 1, 0, 3), (2, 3, 0, 1), (3, 2, 0, 1), (2, 0, 1, 3), (2, 1, 3, 0), (3, 1, 2, 0)]

def label_to_square(label):
    squares = stub1[label]
    v = []
    for i in squares:
        v.append([i // 3 + 1, i % 3 + 1])
    # if label != 9:
    #     v.append([label // 3 + 1, label % 3 + 1])
    return v


def blend(img, height, label):
    # if label != 0 and label != 1:
    #     return img
    alpha = 0.2
    # signal_mask = mlt.imread('./trigger/hello_kitty.png') * 255
    mask = cv2.resize(kitty_mask, (height, height))
    blend_img = (1 - alpha) * img + alpha * mask  # FOR CIFAR10
    blend_img = np.clip(blend_img.astype('uint8'), 0, 255)

    return blend_img



def badnets(img, height, label):
    if label != 0:
        return img
    trig_w = height // 10
    trig_h = height // 10
    distance = 1
    # distance = height // 2
    for j in range(height - distance - trig_w, height - distance):
        for k in range(height - distance - trig_h, height - distance):
            img[j, k] = 255

    squares = label_to_square(label)
    for i in squares:
        img[height - i[0] - distance][height - i[1] - distance] = 0
    return img    


class ClassWiseNoise:
    def __init__(self, args) -> None:
        self.args = args


class ClipCL:
    def __init__(self, args) -> None:
        self.args = args
        device = 'cuda:0'
        self.model, self.transform = clip.load('RN50', device=device)
        def _convert_image_to_rgb(image):
            return image.convert("RGB")
        eps = 8/255
        self.transform = transforms.Compose(
            [
                transforms.Resize(224, interpolation=BICUBIC),
                _convert_image_to_rgb,
                transforms.ToTensor()
            ]
        )
        if args.dataset == 'cifar10':
            text = clip.tokenize(cifar10_class_names).to(device)
            self.height = 32
        if args.dataset == 'cifar100':
            text = clip.tokenize(cifar100_class_names).to(device)
            self.height = 32
        if args.dataset == 'stl10':
            text = clip.tokenize(stl10_class_names).to(device)
            self.height = 96
        if args.dataset == 'svhn':
            text = clip.tokenize(svhn_class_names).to(device)
            self.height = 32
        if args.dataset == 'tinyimagenet':
            text = clip.tokenize(tiny_imagenet_class_names).to(device)
            self.height = 64
        if args.dataset == 'imagenet':
            self.height = 224
            text = clip.tokenize(imagenet_class_names).to(device)
            eps = 16/255
        self.model.eval()
        self.resize_img = transforms.Resize(self.height, interpolation=BICUBIC)
        self.atk = torchattacks.ClipPGD(self.model, text, eps=eps, alpha=0.05, steps=args.steps)

    
    def cl_batch(self, imgs, labels):
        tmp = []
        for img in imgs:
            img = self.transform(img)
            tmp.append(img)
        imgs = torch.stack(tmp)
        labels_t = torch.tensor(labels)
        atk_imgs = self.atk(imgs, labels_t)
        atk_imgs = self.resize_img(atk_imgs)
        imgs = atk_imgs.cpu()
        imgs = imgs*255
        imgs = imgs.byte()
        imgs = torch.clamp(imgs, 0, 255)
        imgs = imgs.permute(0, 2, 3, 1) # (n, h, w, c)
        imgs = imgs.numpy()
        return imgs
    
    def cl(self, img, label):
        img = self.transform(img)



class CL:
    def __init__(self, args) -> None:
        self.args = args
        device = 'cuda:0'
        benign_model = torchattacks.resnet18(num_classes=args.num_classes)
        ckpt = Path(f'ckpt/{args.dataset}_resnet18_e40.zip')
        ckpt = torch.load(ckpt)
        benign_model.load_state_dict(ckpt)
        benign_model.to(device)
        benign_model.eval()
        if args.dataset == 'stl10':
            self.atk = torchattacks.PGD(benign_model, eps=16/255, alpha=0.05, steps=args.steps)
            self.atk.set_attack_mode("targeted")
            def all2all(img, label):
                return (label + 1) % 10
            self.atk._transform_label = all2all 
        else:
            self.atk = torchattacks.PGD(benign_model, eps=8/255, alpha=0.05, steps=args.steps)


    def cl_batch(self, imgs, labels):
        h, w, c = imgs.shape[1], imgs.shape[2], imgs.shape[3]
        imgs = imgs.permute(0, 3, 1, 2) # (n, c, h, w)
        imgs = imgs / 255

        # Adversarial Perturbations
        imgs_t = torch.tensor(imgs, dtype=torch.float)
        labels_t = torch.tensor(labels)
        atk_imgs = self.atk(imgs_t, labels_t)
        imgs = atk_imgs.cpu()

        imgs = imgs*255
        # 将imgs转换为uint8类型
        imgs = imgs.byte()
        # 将imgs用clip限制在0到255之间
        imgs = torch.clamp(imgs, 0, 255)
        imgs = imgs.permute(0, 2, 3, 1) # (n, h, w, c)
        imgs = imgs.numpy()
        return imgs


    def cl(self, img, label):
        h, w, c = img.shape[0], img.shape[1], img.shape[2]
        img = img.transpose(2, 0, 1) # (c, h, w)
        img = img[np.newaxis, :, :, :] # (1, c, h, w)
        img = img / 255

        # Adversarial Perturbations
        img_t = torch.tensor(img, dtype=torch.float)
        label_t = torch.tensor([label])
        atk_img = self.atk(img_t, label_t)
        img = atk_img.cpu().numpy()

        img = img*255
        img = np.clip(img.astype('uint8'), 0, 255)
        img = img.reshape(c, h, w)
        img = img.transpose(1, 2, 0) # (h, w, c)
        return img


class Wanet:
    def __init__(self, args) -> None:
        self.args = args
        self.noises_grid = {}
        self.identity_grid = {}
        self.k = 4
        self.s = 0.5
        # if self.args.dataset == 'svhn':
        #     self.s = 1
        #     self.k = 8
        if self.args.dataset == 'stl10':
            self.k = 96
            self.s = 1
        if self.args.dataset == 'imagenet':
            self.k = 224
            self.s = 1
        self.grid_rescale = 1


    def inject(self, img, height, label):
        int_label = int(label)
        imgs = torch.unsqueeze(to_tensor(img), dim=0)
        if int_label not in self.noises_grid:
            self.init_grid(height, int_label)
        noise_grid, identity_grid = self.get_grid(int_label)
        grid = (identity_grid + self.s * noise_grid / height) * self.grid_rescale
        grid = torch.clamp(grid, -1, 1)
        
        imgs = F.grid_sample(imgs, grid, align_corners=True)
        img = imgs[0]
        img = np.array(to_img(img))
        return img
    

    def init_grid(self, height, label):
        # Prepare grid
        ins = torch.rand(1, 2, self.k, self.k) * 2 - 1
        # ins = torch.zeros(1, 2, self.k, self.k)
        ins = ins / torch.mean(torch.abs(ins))
        noise_grid = (
            F.upsample(ins, size=height, mode="bicubic", align_corners=True)
            .permute(0, 2, 3, 1)
        )
        array1d = torch.linspace(-1, 1, steps=height)
        x, y = torch.meshgrid(array1d, array1d)
        identity_grid = torch.stack((y, x), 2)[None, ...]
        self.noises_grid[label] = noise_grid
        self.identity_grid[label] = identity_grid

    def get_grid(self, label):
        return self.noises_grid[label], self.identity_grid[label]