add option for proxy model

app.py

@@ -7,6 +7,7 @@ import numpy as np
 
 device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
 
+from transformers import AutoModel, BlipImageProcessor
 from diffusers import DiffusionPipeline, AutoencoderKL
 import torchvision.transforms as transforms
 
@@ -42,6 +43,11 @@ class BZHStableSignatureDemo(object):
 
         self.decoders = decoders
 
+        # load the proxy detector
+        self.detector_image_processor = BlipImageProcessor.from_pretrained("imatag/stable-signature-bzh-detector-resnet18")
+        commit_hash = "584a7bc01dc0f02e53bf8b8b295717ed09ed7294"
+        self.detector_model = AutoModel.from_pretrained("imatag/stable-signature-bzh-detector-resnet18", trust_remote_code=True, revision=commit_hash)
+
     def generate(self, mode, seed, prompt):
         generator = torch.Generator(device=device)
         torch.manual_seed(seed)
@@ -91,12 +97,16 @@ class BZHStableSignatureDemo(object):
         # JPEG attack
         mf = io.BytesIO()
         img.save(mf, format='JPEG', quality=jpeg_compression)
+        filesize = mf.tell()
         mf.seek(0)
         img = Image.open(mf)
 
-        return img
+        image_info = "resolution: %dx%d" % img.size
+        image_info += " JPEG file size: %d" % filesize
+
+        return img, image_info
         
-    def detect(self, img):
+    def detect_api(self, img):
         # send to detection API and apply JPEG compression attack
         mf = io.BytesIO()
         img.save(mf, format='PNG')
@@ -115,6 +125,22 @@ class BZHStableSignatureDemo(object):
         data = response.json()
         pvalue = data['p-value']
 
+        return pvalue
+
+    def detect_proxy(self, img):
+        img = img.convert("RGB")
+        inputs = self.detector_image_processor(img, return_tensors="pt")
+
+        with torch.no_grad():
+            pvalue = torch.sigmoid(self.detector_model(**inputs).logits).item()
+
+        return pvalue
+ 
+    def detect(self, img, detection_method):
+        if detection_method == "API":
+            pvalue = self.detect_api(img)
+        else:
+            pvalue = self.detect_proxy(img)
         result = "No watermark detected."
         rpv = 10**int(math.log10(pvalue))
         if pvalue < 1e-3:
@@ -165,18 +191,22 @@ def interface():
                 btn2 = gr.Button("Edit")
         with gr.Row():
             attacked_image = gr.Image(type="pil", width=512, sources=['upload', 'clipboard'])
+        with gr.Row():
+            image_info_label = gr.Label(label="Image info")
 
         gr.Markdown("""## 3. Detect
         Detect the watermark on the altered image. Watermark may not be detected if the image is altered too strongly.
+        You may choose to detect with our fast [proxy model](https://huggingface.co/imatag/stable-signature-bzh-detector-resnet18), or via API for improved robustness.
         """)
         with gr.Row():
+            detection_method = gr.Dropdown(choices=["proxy model", "API"], label="Detection method", value="proxy model")
             btn3 = gr.Button("Detect")
         with gr.Row():
-            detection_label = gr.Label(label="Detection info", show_label=False)
+            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, 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")
+        btn2.click(fn=backend.attack, inputs=[watermarked_image, jpeg_compression, downscale, crop, saturation, brightness, contrast], outputs=[attacked_image, image_info_label], api_name="attack")
+        btn3.click(fn=backend.detect, inputs=[attacked_image, detection_method], outputs=[detection_label], api_name="detect")
 
     return demo