update README

README.md

@@ -13,28 +13,31 @@ You can use this classifier to detect watermarks generated with our [SDXL-turbo
 ## Usage
 
 ```py
-from transformers import AutoModel, BlipImageProcessor
+from transformers import AutoModelForImageClassification, BlipImageProcessor
+
 from PIL import Image
 import sys
-import torch
 
 image_processor = BlipImageProcessor.from_pretrained("imatag/stable-signature-bzh-detector-resnet18")
-commit_hash = "584a7bc01dc0f02e53bf8b8b295717ed09ed7294"
-model = AutoModel.from_pretrained("imatag/stable-signature-bzh-detector-resnet18", trust_remote_code=True, revision=commit_hash)
+model = AutoModelForImageClassification.from_pretrained("imatag/stable-signature-bzh-detector-resnet18")
+model.eval()
 
 img = Image.open(sys.argv[1]).convert("RGB")
 inputs = image_processor(img, return_tensors="pt")
-with torch.no_grad():
-    p = torch.sigmoid(model(**inputs).logits).item()
+p = model(**inputs).logits[0,0] < 0
 
-print(f"approximate p-value: {p}")
+print(f"watermarked: {p}")
 ```
 
 ## Purpose
 
 This model is an approximate version of [IMATAG](https://www.imatag.com/)'s BZH decoder for the watermark embedded in our [SDXL-turbo watermarking demo](https://huggingface.co/spaces/imatag/stable-signature-bzh).
 It works on this watermark only and cannot be used to decode other watermarks.
-It will produce an approximate p-value measuring the risk of mistakenly detecting a watermark on a benign (non-watermarked) image. For an exact p-value and improved robustness, please use the [API](https://huggingface.co/spaces/imatag/stable-signature-bzh/resolve/main/detect_api.py) instead.
+
+It will catch most altered versions of a watermarked image while making roughly one mistake in one thousand on non-watermarked images.
+Alternatively, it can produce an approximate p-value measuring the risk of mistakenly detecting a watermark on a benign (non-watermarked) image, by recalibrating the output as in [this script](https://huggingface.co/imatag/stable-signature-bzh-detector-resnet18/resolve/main/detect_demo_pvalue.py).
+
+To get an exact p-value and for improved robustness, please use the [API](https://huggingface.co/spaces/imatag/stable-signature-bzh/resolve/main/detect_api.py) instead.
 
 For more details on this watermarking technique, check out our [announcement](https://www.imatag.com/blog/unlocking-the-future-of-content-authentication-imatags-breakthrough-in-ai-generated-image-watermarking) and our lab's [blog post](https://imatag-lab.medium.com/stable-signature-meets-bzh-53ad0ba13691).
 

detect_demo_pvalue.py

@@ -0,0 +1,22 @@
+from transformers import AutoModelForImageClassification, BlipImageProcessor
+from huggingface_hub import hf_hub_download
+from safetensors import safe_open
+
+from PIL import Image
+import sys
+import torch
+
+image_processor = BlipImageProcessor.from_pretrained("imatag/stable-signature-bzh-detector-resnet18")
+model = AutoModelForImageClassification.from_pretrained("imatag/stable-signature-bzh-detector-resnet18")
+calibration = hf_hub_download("imatag/stable-signature-bzh-detector-resnet18", filename="calibration.safetensors")
+with safe_open(calibration, framework="pt") as f:
+    calibration_logits = f.get_tensor("logits")
+
+img = Image.open(sys.argv[1]).convert("RGB")
+inputs = image_processor(img, return_tensors="pt")
+with torch.no_grad():
+    p = model(**inputs).logits[...,0]
+    p = (1 + torch.searchsorted(calibration_logits, p)) / calibration_logits.shape[0]
+    p = p.item()
+
+print(f"approximate p-value: {p}")

detect_demo.py → detect_demo_simple.py

@@ -1,14 +1,14 @@
-from transformers import AutoModel, BlipImageProcessor
+from transformers import AutoModelForImageClassification, BlipImageProcessor
+
 from PIL import Image
 import sys
-import torch
 
 image_processor = BlipImageProcessor.from_pretrained("imatag/stable-signature-bzh-detector-resnet18")
-commit_hash = "584a7bc01dc0f02e53bf8b8b295717ed09ed7294"
-model = AutoModel.from_pretrained("imatag/stable-signature-bzh-detector-resnet18", trust_remote_code=True, revision=commit_hash)
+model = AutoModelForImageClassification.from_pretrained("imatag/stable-signature-bzh-detector-resnet18")
+model.eval()
 
 img = Image.open(sys.argv[1]).convert("RGB")
 inputs = image_processor(img, return_tensors="pt")
-with torch.no_grad():
-    p = torch.sigmoid(model(**inputs).logits).item()
-print(f"approximate p-value: {p}")
+p = model(**inputs).logits[0,0] < 0
+
+print(f"watermarked: {p}")