init

app.py

@@ -0,0 +1,49 @@
+import gradio as gr
+from torchvision.utils import save_image
+from attack import Attacker
+import argparse
+
+
+def do_attack(img, eps, step_size, steps, progress=gr.Progress()):
+    args=argparse.Namespace()
+    args.out_dir='./'
+    args.target='auto'
+    args.eps=eps
+    args.step_size=step_size
+    args.steps=steps
+    args.test_atk=False
+
+    step = progress.tqdm(range(steps))
+
+    def pdg_prog(ori_images, images, labels):
+        step.update(1)
+
+    attacker = Attacker(args, pgd_callback=pdg_prog)
+    atk_img, noise = attacker.attack_(img)
+    attacker.save_image(atk_img, noise, 'out.png')
+    return 'out.png'
+
+with gr.Blocks(title="Anime AI Detect Fucker Demo", theme="dark") as demo:
+    gr.HTML('<a href="https://github.com/7eu7d7/anime-ai-detect-fucker">github repo</a>')
+
+    with gr.Row():
+        eps = gr.Slider(label="eps (Noise intensity)", minimum=1, maximum=16, step=1, value=1)
+        step_size = gr.Slider(label="Noise step size", minimum=0.001, maximum=16, step=0.001, value=0.136)
+    with gr.Row():
+        steps = gr.Slider(label="step count", minimum=1, maximum=100, step=1, value=20)
+        model_name = gr.Dropdown(label="attack target",
+                                 choices=["auto", "human", "ai"],
+                                 value="auto", show_label=True)
+
+    input_image = gr.Image(label="Clean Image", type="pil")
+
+    atk_btn = gr.Button("Attack")
+
+    with gr.Column():
+        output_image = gr.Image(label="Attacked Image")
+
+    atk_btn.click(fn=do_attack,
+                      inputs=[input_image, eps, step_size, steps],
+                      outputs=output_image)
+
+demo.launch()

attack.py

@@ -0,0 +1,113 @@
+import torch
+import os
+from transformers import BeitFeatureExtractor, BeitForImageClassification
+from PIL import Image
+
+from torchvision.utils import save_image
+import torch.nn.functional as F
+from torchvision import transforms
+
+from attacker import *
+from torch.nn import CrossEntropyLoss
+
+import argparse
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+def make_args():
+    parser = argparse.ArgumentParser(description='PyTorch MS_COCO Training')
+
+    parser.add_argument('inputs', type=str)
+    parser.add_argument('--out_dir', type=str, default='./output')
+    parser.add_argument('--target', type=str, default='auto', help='[auto, ai, human]')
+    parser.add_argument('--eps', type=float, default=8/8, help='Noise intensity ')
+    parser.add_argument('--step_size', type=float, default=1.087313/8, help='Attack step size')
+    parser.add_argument('--steps', type=int, default=20, help='Attack step count')
+
+    parser.add_argument('--test_atk', action='store_true')
+
+    return parser.parse_args()
+
+class Attacker:
+    def __init__(self, args, pgd_callback):
+        self.args=args
+        os.makedirs(args.out_dir, exist_ok=True)
+
+        print('正在加载模型...')
+        self.feature_extractor = BeitFeatureExtractor.from_pretrained('saltacc/anime-ai-detect')
+        self.model = BeitForImageClassification.from_pretrained('saltacc/anime-ai-detect').cuda()
+        print('加载完毕')
+
+        if args.target=='ai': #攻击成被识别为AI
+            self.target = torch.tensor([1]).to(device)
+        elif args.target=='human':
+            self.target = torch.tensor([0]).to(device)
+
+        dataset_mean_t = torch.tensor([0.5, 0.5, 0.5]).view(1, -1, 1, 1).cuda()
+        dataset_std_t = torch.tensor([0.5, 0.5, 0.5]).view(1, -1, 1, 1).cuda()
+        self.pgd = PGD(self.model, img_transform=(lambda x: (x - dataset_mean_t) / dataset_std_t, lambda x: x * dataset_std_t + dataset_mean_t))
+        self.pgd.set_para(eps=(args.eps * 2) / 255, alpha=lambda: (args.step_size * 2) / 255, iters=args.steps)
+        self.pgd.set_loss(CrossEntropyLoss())
+        self.pgd.set_call_back(pgd_callback)
+
+    def save_image(self, image, noise, img_name):
+        # 缩放图片只缩放噪声
+        W, H = image.size
+        noise = F.interpolate(noise, size=(H, W), mode='bicubic')
+        img_save = transforms.ToTensor()(image) + noise
+        save_image(img_save, os.path.join(self.args.out_dir, f'{img_name[:img_name.rfind(".")]}_atk.png'))
+
+    def attack_(self, image):
+        inputs = self.feature_extractor(images=image, return_tensors="pt")['pixel_values'].cuda()
+
+        if self.args.target == 'auto':
+            with torch.no_grad():
+                outputs = self.model(inputs)
+                logits = outputs.logits
+                cls = logits.argmax(-1).item()
+                target = torch.tensor([cls]).to(device)
+        else:
+            target = self.target
+
+        if self.args.test_atk:
+            self.test_image(inputs, 'before attack')
+
+        atk_img = self.pgd.attack(inputs, target)
+
+        noise = self.pgd.img_transform[1](atk_img).detach().cpu() - self.pgd.img_transform[1](inputs).detach().cpu()
+
+        if self.args.test_atk:
+            self.test_image(atk_img, 'after attack')
+
+        return atk_img, noise
+
+    def attack_one(self, path):
+        image = Image.open(path).convert('RGB')
+        atk_img, noise = self.attack_(image)
+        self.save_image(image, noise, os.path.basename(path))
+
+    def attack(self, path):
+        count=0
+        if os.path.isdir(path):
+            img_list=[os.path.join(path, x) for x in os.listdir(path)]
+            for img in img_list:
+                if (img.lower().endswith(('.bmp', '.dib', '.png', '.jpg', '.jpeg', '.pbm', '.pgm', '.ppm', '.tif', '.tiff'))):
+                    self.attack_one(img)
+                    count+=1
+        else:
+            if (path.lower().endswith(('.bmp', '.dib', '.png', '.jpg', '.jpeg', '.pbm', '.pgm', '.ppm', '.tif', '.tiff'))):
+                self.attack_one(path)
+                count += 1
+        print(f'总共攻击{count}张图像')
+
+    @torch.no_grad()
+    def test_image(self, img, pre_fix):
+        outputs = self.model(img)
+        logits = outputs.logits
+        predicted_class_idx = logits.argmax(-1).item()
+        print(pre_fix, "class:", self.model.config.id2label[predicted_class_idx], 'logits:', logits)
+
+if __name__ == '__main__':
+    args=make_args()
+    attacker = Attacker(args)
+    attacker.attack(args.inputs)

attacker/FGSM.py

@@ -0,0 +1,48 @@
+import torch
+from torch import nn
+from copy import deepcopy
+from .base import Attacker
+from torch.cuda import amp
+
+class FGSM(Attacker):
+    def __init__(self, model, img_transform=(lambda x:x, lambda x:x), use_amp=False):
+        super().__init__(model, img_transform)
+        self.use_amp=use_amp
+
+        if use_amp:
+            self.scaler = amp.GradScaler()
+
+    def set_para(self, eps=8, alpha=lambda:8, **kwargs):
+        super().set_para(eps=eps, alpha=alpha, **kwargs)
+
+    def step(self, images, labels, loss):
+        with amp.autocast(enabled=self.use_amp):
+            images.requires_grad = True
+            outputs = self.model(images).logits
+
+            self.model.zero_grad()
+            cost = loss(outputs, labels)
+
+        if self.use_amp:
+            self.scaler.scale(cost).backward()
+        else:
+            cost.backward()
+
+        adv_images = (images + self.alpha() * images.grad.sign()).detach_()
+        eta = torch.clamp(adv_images - self.ori_images, min=-self.eps, max=self.eps)
+        images = self.img_transform[0](torch.clamp(self.img_transform[1](self.ori_images + eta), min=0, max=255).detach_())
+
+        return images
+
+    def attack(self, images, labels):
+        #images = deepcopy(images)
+        #self.ori_images = deepcopy(images)
+
+        self.model.eval()
+
+        images = self.forward(self, images, labels)
+
+        self.model.zero_grad()
+        self.model.train()
+
+        return images

attacker/PGD.py

@@ -0,0 +1,84 @@
+import torch
+from torch import nn
+from copy import deepcopy
+from .base import Attacker, Empty
+from torch.cuda import amp
+from tqdm import tqdm
+
+class PGD(Attacker):
+    def __init__(self, model, img_transform=(lambda x:x, lambda x:x), use_amp=False):
+        super().__init__(model, img_transform)
+        self.use_amp=use_amp
+        self.call_back=None
+        self.img_loader=None
+        self.img_hook=None
+
+        self.scaler = amp.GradScaler(enabled=use_amp)
+
+    def set_para(self, eps=8, alpha=lambda:8, iters=20, **kwargs):
+        super().set_para(eps=eps, alpha=alpha, iters=iters, **kwargs)
+
+    def set_call_back(self, call_back):
+        self.call_back=call_back
+
+    def set_img_loader(self, img_loader):
+        self.img_loader=img_loader
+
+    def step(self, images, labels, loss):
+        with amp.autocast(enabled=self.use_amp):
+            images.requires_grad = True
+            outputs = self.model(images).logits
+
+            self.model.zero_grad()
+            cost = loss(outputs, labels)#+outputs[2].view(-1)[0]*0+outputs[1].view(-1)[0]*0+outputs[0].view(-1)[0]*0 #support DDP
+
+        self.scaler.scale(cost).backward()
+
+        adv_images = (images + self.alpha() * images.grad.sign()).detach_()
+        eta = torch.clamp(adv_images - self.ori_images, min=-self.eps, max=self.eps)
+        images = self.img_transform[0](torch.clamp(self.img_transform[1](self.ori_images + eta), min=0, max=1).detach_())
+
+        return images
+
+    def set_data(self, images, labels):
+        self.ori_images = deepcopy(images)
+        self.images = images
+        self.labels = labels
+
+    def __iter__(self):
+        self.atk_step=0
+        return self
+
+    def __next__(self):
+        self.atk_step += 1
+        if self.atk_step>self.iters:
+            raise StopIteration
+
+        with self.model.no_sync() if isinstance(self.model, nn.parallel.DistributedDataParallel) else Empty():
+            self.model.eval()
+
+            self.images = self.forward(self, self.images, self.labels)
+
+            self.model.zero_grad()
+            self.model.train()
+
+        return self.ori_images, self.images.detach(), self.labels
+
+    def attack(self, images, labels):
+        #images = deepcopy(images)
+        self.ori_images = deepcopy(images)
+
+        for i in tqdm(range(self.iters)):
+            self.model.eval()
+
+            images = self.forward(self, images, labels)
+
+            self.model.zero_grad()
+            self.model.train()
+            if self.call_back:
+                self.call_back(self.ori_images, images.detach(), labels)
+
+            if self.img_hook is not None:
+                images=self.img_hook(self.ori_images, images.detach())
+
+        return images

attacker/__init__.py

@@ -0,0 +1,3 @@
+from .base import *
+from .PGD import *
+from .FGSM import *

attacker/base.py

@@ -0,0 +1,33 @@
+
+class Attacker:
+    def __init__(self, model, img_transform=(lambda x:x, lambda x:x)):
+        self.model = model  # 必须是pytorch的model
+        '''self.model.eval()
+        for k, v in self.model.named_parameters():
+            v.requires_grad = False'''
+        self.img_transform=img_transform
+        self.forward = lambda attacker, images, labels: attacker.step(images, labels, attacker.loss)
+
+    def set_para(self, **kwargs):
+        for k,v in kwargs.items():
+            setattr(self, k,v)
+
+    def set_forward(self, forward):
+        self.forward=forward
+
+    def step(self, images, labels, loss):
+        pass
+
+    def set_loss(self, loss):
+        self.loss=loss
+
+    def attack(self, images, labels):
+        pass
+
+
+class Empty:
+    def __enter__(self):
+        pass
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass

requirements.txt

@@ -0,0 +1,12 @@
+torch==1.12.1
+torchvision==0.13.1
+timm==0.6.12
+Pillow
+blobfile
+mypy
+numpy
+pytest
+requests
+einops
+deepspeed==0.4.0
+scipy