dev(narugo): add same mode

app.py

@@ -48,13 +48,18 @@ if __name__ == '__main__':
                     '      style=\'text-decoration: underline;\'>'
                     '       Github - 7eu7d7/anime-ai-detect-fucker'
                     '   </a>.'
+                    '   or'
+                    '   <a href="https://colab.research.google.com/drive/1bzoZtE28Y8vfmjuaxXKdh-fr4K0WWSmf?usp=sharing" '
+                    '      target="_parent">'
+                    '       <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>'
+                    '   </a>'
                     '</p>')
         with gr.Row():
             with gr.Column():
                 gr_input_image = gr.Image(type='pil', label='Original Image')
                 with gr.Row():
-                    gr_attack_target = gr.Radio(['auto', 'ai', 'human'], value='auto', label='Attack Target')
-                    gr_steps = gr.Slider(minimum=1, maximum=50, value=20, step=1, label='Steps')
+                    gr_attack_target = gr.Radio(['auto', 'ai', 'human', 'same'], value='auto', label='Attack Target')
+                    gr_steps = gr.Slider(minimum=1, maximum=100, value=20, step=1, label='Steps')
                 with gr.Row():
                     gr_eps = gr.Slider(label="eps (Noise intensity)", minimum=1, maximum=16, step=1, value=1)
                     gr_noise_step_size = gr.Slider(label="Noise step size", minimum=0.001, maximum=16, step=0.001,

attack.py

@@ -20,7 +20,7 @@ def make_args(args_=None):
 
     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('--target', type=str, default='auto', help='[auto, ai, human, same]')
     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')
@@ -50,6 +50,8 @@ class Attacker:
             self.target = torch.tensor([1]).to(device)
         elif args.target == 'human':
             self.target = torch.tensor([0]).to(device)
+        else:
+            self.target = torch.tensor([0]).to(device)
 
         dataset_mean_t = torch.tensor([0.5, 0.5, 0.5]).view(1, -1, 1, 1)
         dataset_std_t = torch.tensor([0.5, 0.5, 0.5]).view(1, -1, 1, 1)
@@ -59,7 +61,13 @@ class Attacker:
         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())
+        if args.target != 'same':
+            self.pgd.set_loss(CrossEntropyLoss())
+        else:
+            def loss_same(a, b):
+                return -torch.exp((a[0, 0] - a[0, 1]) ** 2)
+
+            self.pgd.set_loss(loss_same)
 
     def save_image(self, image, noise, img_name):
         # 缩放图片只缩放噪声