change demo workflow

app.py

@@ -53,7 +53,7 @@ class BZHStableSignatureDemo(object):
         output = self.pipe(prompt, num_inference_steps=4, guidance_scale=0.0, output_type="pil")
         return output.images[0]
 
-    def attack_detect(self, img, jpeg_compression, downscale, crop, saturation):
+    def attack(self, img, jpeg_compression, downscale, crop, saturation, brightness, contrast):
 
         img = img.convert("RGB")
         
@@ -81,10 +81,25 @@ class BZHStableSignatureDemo(object):
 
         converter = ImageEnhance.Color(img)
         img = converter.enhance(saturation)
+
+        converter = ImageEnhance.Brightness(img)
+        img = converter.enhance(brightness)
+
+        converter = ImageEnhance.Contrast(img)
+        img = converter.enhance(contrast)
+
+        # JPEG attack
+        mf = io.BytesIO()
+        img.save(mf, format='JPEG', quality=jpeg_compression)
+        mf.seek(0)
+        img = Image.open(mf)
+
+        return img
         
+    def detect(self, img):
         # send to detection API and apply JPEG compression attack
         mf = io.BytesIO()
-        img.save(mf, format='JPEG', quality=jpeg_compression) # includes JPEG attack
+        img.save(mf, format='PNG')
         b64 = base64.b64encode(mf.getvalue())
         data = {
             'image': b64.decode('utf8')
@@ -100,16 +115,15 @@ class BZHStableSignatureDemo(object):
         data = response.json()
         pvalue = data['p-value']
 
-        mf.seek(0)
-        img0 = Image.open(mf) # reload to show JPEG attack
-
         result = "No watermark detected."
         rpv = 10**int(math.log10(pvalue))
         if pvalue < 1e-3:
-            result = "Watermark detected with low confidence (p-value<%.0e)" % rpv
+            result = "Watermark detected with low confidence" # (p-value<%.0e)" % rpv
         if pvalue < 1e-9:
-            result = "Watermark detected with high confidence (p-value<%.0e)" % rpv
-        return (img0, result)
+            result = "Watermark detected with high confidence" # (p-value<%.0e)" % rpv
+        score = min(int(-math.log10(pvalue)), 10)
+        #print("score = ", score)
+        return { result: score/10 }
 
 def interface():
     prompt = "sailing ship in storm by Rembrandt"
@@ -124,6 +138,8 @@ def interface():
         the VAE decoder of StableDiffusion is fine-tuned to produce images including a specific invisible watermark. We combined
         this method with a demo version of [IMATAG](https://www.imatag.com/)'s in-house decoder. The watermarking system operates in zero-bit mode for improved robustness.""")
 
+        gr.Markdown("""## 1. Generate
+        Select a watermarking strength and generate images with StableDiffusion-XL Turbo from prompt and seed as usual.""")
         with gr.Row():
             inp = gr.Textbox(label="Prompt", value=prompt)
             seed = gr.Number(label="Seed", precision=0)
@@ -131,19 +147,36 @@ def interface():
         with gr.Row():
             btn1 = gr.Button("Generate")
         with gr.Row():
-            watermarked_image = gr.Image(type="pil", width=512, height=512)           
+            watermarked_image = gr.Image(type="pil", width=512, height=512, sources=[], interactive=False)
+
+        gr.Markdown("""## 2. Edit
+        With these controls you may alter the generated image before detection. You may also upload your own edited image instead.""")
+        with gr.Row():
             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("Modify & Detect")
                 with gr.Row():
-                    attacked_image = gr.Image(type="pil", width=256)
-                    detection_label = gr.Label(label="Detection info")
+                    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")
+                with gr.Row():
+                    brightness = gr.Slider(0, 2, value=1, step=0.1, label="Brightness")
+                    contrast = gr.Slider(0, 2, value=1, step=0.1, label="Contrast")
+                with gr.Row():
+                    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("Edit")
+        with gr.Row():
+            attacked_image = gr.Image(type="pil", width=512, sources=['upload', 'clipboard'])
+
+        gr.Markdown("""## 3. Detect
+        Detect the watermark on the altered image. Watermark may not be detected if the image is altered too strongly.
+        """)
+        with gr.Row():
+            btn3 = gr.Button("Detect")
+        with gr.Row():
+            detection_label = gr.Label(label="Detection info", show_label=False)
+
         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, crop, saturation], outputs=[attacked_image, detection_label], api_name="detect")
+        btn2.click(fn=backend.attack, inputs=[watermarked_image, jpeg_compression, downscale, crop, saturation, brightness, contrast], outputs=[attacked_image], api_name="attack")
+        btn3.click(fn=backend.detect, inputs=[attacked_image], outputs=[detection_label], api_name="detect")
 
     return demo