add crop

app.py

@@ -19,6 +19,8 @@ import json
 from PIL import Image, ImageEnhance
 import base64
 import io
+import random
+import math
 
 class BZHStableSignatureDemo(object):
 
@@ -53,7 +55,7 @@ class BZHStableSignatureDemo(object):
         output = self.pipe(prompt, num_inference_steps=4, guidance_scale=0.0, output_type="pil")
         return output.images[0] #{ "background": output.images[0], "layers": [], "composite": None }
 
-    def attack_detect(self, img, jpeg_compression, downscale, saturation):
+    def attack_detect(self, img, jpeg_compression, downscale, crop, saturation):
 
         #img = img_edit["composite"]
         img = img.convert("RGB")
@@ -63,6 +65,21 @@ class BZHStableSignatureDemo(object):
             size = img.size
             size = (int(size[0] / downscale), int(size[1] / downscale))
             img = img.resize(size, Image.BICUBIC)
+        if crop != 0:
+            width, height = img.size
+            area = width * height
+            log_rmin = math.log(0.5)
+            log_rmax = math.log(2.0)
+            for _ in range(10):
+                target_area = area * (1 - crop)
+                aspect_ratio = math.exp(random.random() * (log_rmax - log_rmin) + log_rmin)
+                w = int(round(math.sqrt(target_area * aspect_ratio)))
+                h = int(round(math.sqrt(target_area / aspect_ratio)))
+                if 0 < w <= width and 0 < h <= height:
+                    top = random.randint(0, height - h + 1)
+                    left = random.randint(0, width - w + 1)
+                    img = img.crop((left, top, left+w, top+h))
+                    break
 
         converter = ImageEnhance.Color(img)
         img = converter.enhance(saturation)
@@ -119,15 +136,17 @@ def interface():
         with gr.Row():
             watermarked_image = gr.Image(type="pil", width=512, height=512)           
             with gr.Column():
+                gr.Markdown("""With these controls you may alter the generated image before detection. You may also upload your own edited image instead.""")
                 downscale = gr.Slider(1, 3, value=1, step=0.1, label="Downscale ratio")
+                crop = gr.Slider(0, 0.9, value=0, step=0.01, label="Random crop ratio")
                 saturation = gr.Slider(0, 2, value=1, step=0.1, label="Color saturation")
                 jpeg_compression = gr.Slider(value=100, step=5, label="JPEG quality")
-                btn2 = gr.Button("Attack & Detect")
+                btn2 = gr.Button("Modify & Detect")
                 with gr.Row():
                     attacked_image = gr.Image(type="pil", width=256)
                     detection_label = gr.Label(label="Detection info")
         btn1.click(fn=backend.generate, inputs=[mode, seed, inp], outputs=[watermarked_image], api_name="generate")
-        btn2.click(fn=backend.attack_detect, inputs=[watermarked_image, jpeg_compression, downscale, saturation], outputs=[attacked_image, detection_label], api_name="detect")
+        btn2.click(fn=backend.attack_detect, inputs=[watermarked_image, jpeg_compression, downscale, crop, saturation], outputs=[attacked_image, detection_label], api_name="detect")
 
     return demo