Support gradio

.gitignore

@@ -0,0 +1,2 @@
+*.pyc
+watermark-anything

.vscode/launch.json

@@ -0,0 +1,16 @@
+{
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "debugpy: app.py",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "${workspaceFolder}/app.py",
+            "console": "integratedTerminal",
+            "env": {
+                // "HF_ENDPOINT": "https://hf-mirror.com"
+            },
+            "justMyCode": false
+        }
+    ]
+}

README.md

@@ -1,10 +1,11 @@
 ---
 title: Watermark Anything
-emoji: ⚡
+emoji: 💧
 colorFrom: gray
 colorTo: indigo
 sdk: gradio
-sdk_version: 5.9.0
+sdk_version: 5.8.0
+python_version: '3.10.14'
 app_file: app.py
 pinned: false
 license: apache-2.0

app.py

@@ -0,0 +1,44 @@
+import os
+import sys
+import git
+import subprocess
+from huggingface_hub import hf_hub_download
+
+REPO_URL = "https://github.com/facebookresearch/watermark-anything.git"
+REPO_BRANCH = '88e3ae5d5866a7daaac167ea202a61a7d69ef590'
+LOCAL_PATH = "./watermark-anything"
+MODEL_ID = "xiaoyao9184/watermark-anything"
+
+def install_src():
+    if not os.path.exists(LOCAL_PATH):
+        print(f"Cloning repository from {REPO_URL}...")
+        repo = git.Repo.clone_from(REPO_URL, LOCAL_PATH)
+        repo.git.checkout(REPO_BRANCH)
+    else:
+        print(f"Repository already exists at {LOCAL_PATH}")
+
+    requirements_path = os.path.join(LOCAL_PATH, "requirements.txt")
+    if os.path.exists(requirements_path):
+        print("Installing requirements...")
+        subprocess.check_call(["pip", "install", "-r", requirements_path])
+    else:
+        print("No requirements.txt found.")
+
+def install_model():
+    checkpoint_path = os.path.join(LOCAL_PATH, "checkpoints")
+    hf_hub_download(repo_id=MODEL_ID, filename='checkpoint.pth', local_dir=checkpoint_path)
+
+# clone repo and download model
+install_src()
+install_model()
+
+# change directory
+print(f"Current Directory: {os.getcwd()}")
+os.chdir(LOCAL_PATH)
+print(f"New Directory: {os.getcwd()}")
+
+# fix sys.path for import
+sys.path.append(os.getcwd())
+
+# run gradio
+import gradio_app

gradio_app.py

@@ -0,0 +1,388 @@
+
+
+import gradio as gr
+
+import re
+import string
+import random
+import os
+import numpy as np
+from PIL import Image
+import torch
+import torch.nn.functional as F
+from torchvision import transforms
+
+
+from watermark_anything.data.metrics import msg_predict_inference
+from notebooks.inference_utils import (
+    load_model_from_checkpoint,
+    default_transform,
+    unnormalize_img,
+    create_random_mask,
+    plot_outputs,
+    msg2str,
+    torch_to_np,
+    multiwm_dbscan
+)
+
+# Device configuration
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Constants
+proportion_masked = 0.5  # Proportion of image to be watermarked
+epsilon = 1  # min distance between decoded messages in a cluster
+min_samples = 500  # min number of pixels in a 256x256 image to form a cluster
+
+# Color map for visualization
+color_map = {
+    -1: [0, 0, 0],      # Black for -1
+    0: [255, 0, 255],   # ? for 0
+    1: [255, 0, 0],     # Red for 1
+    2: [0, 255, 0],     # Green for 2
+    3: [0, 0, 255],     # Blue for 3
+    4: [255, 255, 0],   # Yellow for 4
+}
+
+def load_wam():
+    # Load the model from the specified checkpoint
+    exp_dir = "checkpoints"
+    json_path = os.path.join(exp_dir, "params.json")
+    ckpt_path = os.path.join(exp_dir, 'checkpoint.pth')
+    wam = load_model_from_checkpoint(json_path, ckpt_path).to(device).eval()
+    return wam
+
+def image_detect(img_pil: Image.Image) -> (torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor):
+    img_pt = default_transform(img_pil).unsqueeze(0).to(device)  # [1, 3, H, W]
+
+    # Detect the watermark in the multi-watermarked image
+    preds = wam.detect(img_pt)["preds"]  # [1, 33, 256, 256]
+    mask_preds = F.sigmoid(preds[:, 0, :, :])  # [1, 256, 256], predicted mask
+    mask_preds_res = F.interpolate(mask_preds.unsqueeze(1), size=(img_pt.shape[-2], img_pt.shape[-1]), mode="bilinear", align_corners=False)  # [1, 1, H, W]
+    bit_preds = preds[:, 1:, :, :]  # [1, 32, 256, 256], predicted bits
+
+    # positions has the cluster number at each pixel. can be upsaled back to the original size.
+    try:
+        centroids, positions = multiwm_dbscan(bit_preds, mask_preds, epsilon=epsilon, min_samples=min_samples)
+        centroids_pt = torch.stack(list(centroids.values()))
+    except (UnboundLocalError) as e:
+        print(f"Error while detecting watermark: {e}")
+        positions = None
+        centroids = None
+        centroids_pt = None
+
+    return img_pt, (mask_preds_res>0.5).float(), positions, centroids, centroids_pt
+
+def image_embed(img_pil: Image.Image, wm_msgs: torch.Tensor, wm_masks: torch.Tensor) -> (torch.Tensor, torch.Tensor, torch.Tensor):
+    img_pt = default_transform(img_pil).unsqueeze(0).to(device)  # [1, 3, H, W]
+
+    # Embed the watermark message into the image
+    # Mask to use. 1 values correspond to pixels where the watermark will be placed.
+    multi_wm_img = img_pt.clone()
+    for ii in range(len(wm_msgs)):
+        wm_msg, mask = wm_msgs[ii].unsqueeze(0), wm_masks[ii]
+        outputs = wam.embed(img_pt, wm_msg)
+        multi_wm_img = outputs['imgs_w'] * mask + multi_wm_img * (1 - mask)
+
+    torch.cuda.empty_cache()
+    return img_pt, multi_wm_img, wm_masks.sum(0)
+
+def create_bounding_mask(img_size, boxes):
+    """Create a binary mask from bounding boxes.
+
+    Args:
+        img_size (tuple): Image size (height, width)
+        boxes (list): List of tuples (x1, y1, x2, y2) defining bounding boxes
+
+    Returns:
+        torch.Tensor: Binary mask tensor
+    """
+    mask = torch.zeros(img_size)
+    for x1, y1, x2, y2 in boxes:
+        mask[y1:y2, x1:x2] = 1
+    return mask
+
+def centroid_to_hex(centroid):
+    binary_int = 0
+    for bit in centroid:
+        binary_int = (binary_int << 1) | int(bit.item())
+    return format(binary_int, '08x')
+
+# Load the model
+wam = load_wam()
+
+def detect_watermark(image):
+    if image is None:
+        return None, None, None, {"status": "error", "messages": [], "error": "No image provided"}
+
+    img_pil = Image.fromarray(image).convert("RGB")
+    det_img, pred, positions, centroids, centroids_pt = image_detect(img_pil)
+
+    # Convert tensor images to numpy for display
+    detected_img = torch_to_np(det_img.detach())
+    pred_mask = torch_to_np(pred.detach().repeat(1, 3, 1, 1))
+
+    # Create cluster visualization
+    if positions is not None:
+        resize_ori = transforms.Resize(det_img.shape[-2:])
+        rgb_image = torch.zeros((3, positions.shape[-1], positions.shape[-2]), dtype=torch.uint8)
+        for value, color in color_map.items():
+            mask_ = positions == value
+            for channel, color_value in enumerate(color):
+                rgb_image[channel][mask_.squeeze()] = color_value
+        rgb_image = resize_ori(rgb_image.float()/255)
+        cluster_viz = rgb_image.permute(1, 2, 0).numpy()
+
+        # Create message output as JSON
+        messages = []
+        for key in centroids.keys():
+            centroid_hex = centroid_to_hex(centroids[key])
+            centroid_hex_array = "-".join([centroid_hex[i:i+4] for i in range(0, len(centroid_hex), 4)])
+            messages.append({
+                "id": int(key),
+                "message": centroid_hex_array,
+                "color": color_map[key]
+            })
+        message_json = {
+            "status": "success",
+            "messages": messages,
+            "count": len(messages)
+        }
+    else:
+        cluster_viz = np.zeros_like(detected_img)
+        message_json = {
+            "status": "no_detection",
+            "messages": [],
+            "count": 0
+        }
+
+    return pred_mask, cluster_viz, message_json
+
+def embed_watermark(image, wm_num, wm_type, wm_str, wm_loc):
+    if image is None:
+        return None, None, {
+            "status": "failure",
+            "messages": "No image provided"
+        }
+
+    if wm_type == "input":
+        if not re.match(r"^([0-9A-F]{4}-[0-9A-F]{4}-){%d}[0-9A-F]{4}-[0-9A-F]{4}$" % (wm_num-1), wm_str):
+            tip = "-".join([f"FFFF-{_}{_}{_}{_}" for _ in range(wm_num)])
+            return None, None, {
+                "status": "failure",
+                "messages": f"Invalid type input. Please use {tip}"
+            }
+
+    if wm_loc == "bounding":
+        if ROI_coordinates['clicks'] != wm_num * 2:
+            return None, None, {
+                "status": "failure",
+                "messages": "Invalid location input. Please draw at least %d bounding ROI" % (wm_num)
+            }
+
+    img_pil = Image.fromarray(image).convert("RGB")
+
+    # Generate watermark messages based on type
+    wm_msgs = []
+    if wm_type == "random":
+        chars = '-'.join(''.join(random.choice(string.hexdigits) for _ in range(4)) for _ in range(wm_num * 2))
+        wm_str = chars.lower()
+    wm_hex = wm_str.replace("-", "")
+    for i in range(0, len(wm_hex), 8):
+        chunk = wm_hex[i:i+8]
+        binary = bin(int(chunk, 16))[2:].zfill(32)
+        wm_msgs.append([int(b) for b in binary])
+    # Define a 32-bit message to be embedded into the images
+    wm_msgs = torch.tensor(wm_msgs, dtype=torch.float32).to(device)
+
+    # Create mask based on location type
+    wm_masks = None
+    if wm_loc == "random":
+        img_pt = default_transform(img_pil).unsqueeze(0).to(device)
+        # To ensure at least `proportion_masked %` of the width is randomly usable,
+        # otherwise, it is easy to enter an infinite loop and fail to find a usable width.
+        mask_percentage = img_pil.height / img_pil.width * proportion_masked / wm_num
+        wm_masks = create_random_mask(img_pt, num_masks=wm_num, mask_percentage=mask_percentage)
+    elif wm_loc == "bounding" and sections:
+        wm_masks = torch.zeros((len(sections), 1, img_pil.height, img_pil.width), dtype=torch.float32).to(device)
+        for idx, ((x_start, y_start, x_end, y_end), _) in enumerate(sections):
+            left = min(x_start, x_end)
+            right = max(x_start, x_end)
+            top = min(y_start, y_end)
+            bottom = max(y_start, y_end)
+            wm_masks[idx, 0, top:bottom, left:right] = 1
+
+
+    img_pt, embed_img_pt, embed_mask_pt = image_embed(img_pil, wm_msgs, wm_masks)
+
+    # Convert to numpy for display
+    img_np = torch_to_np(embed_img_pt.detach())
+    mask_np = torch_to_np(embed_mask_pt.detach().expand(3, -1, -1))
+    message_json = {
+        "status": "success",
+        "messages": wm_str
+    }
+    return img_np, mask_np, message_json
+
+
+
+# ROI means Region Of Interest. It is the region where the user clicks
+# to specify the location of the watermark.
+ROI_coordinates = {
+    'x_temp': 0,
+    'y_temp': 0,
+    'x_new': 0,
+    'y_new': 0,
+    'clicks': 0,
+}
+
+sections = []
+
+def get_select_coordinates(img, evt: gr.SelectData, num):
+    if ROI_coordinates['clicks'] >= num * 2:
+        gr.Warning(f"Cant add more than {num} of Watermarks.")
+        return (img, sections)
+
+    # update new coordinates
+    ROI_coordinates['clicks'] += 1
+    ROI_coordinates['x_temp'] = ROI_coordinates['x_new']
+    ROI_coordinates['y_temp'] = ROI_coordinates['y_new']
+    ROI_coordinates['x_new'] = evt.index[0]
+    ROI_coordinates['y_new'] = evt.index[1]
+    # compare start end coordinates
+    x_start = ROI_coordinates['x_new'] if (ROI_coordinates['x_new'] < ROI_coordinates['x_temp']) else ROI_coordinates['x_temp']
+    y_start = ROI_coordinates['y_new'] if (ROI_coordinates['y_new'] < ROI_coordinates['y_temp']) else ROI_coordinates['y_temp']
+    x_end = ROI_coordinates['x_new'] if (ROI_coordinates['x_new'] > ROI_coordinates['x_temp']) else ROI_coordinates['x_temp']
+    y_end = ROI_coordinates['y_new'] if (ROI_coordinates['y_new'] > ROI_coordinates['y_temp']) else ROI_coordinates['y_temp']
+    if ROI_coordinates['clicks'] % 2 == 0:
+        sections[len(sections) - 1] = ((x_start, y_start, x_end, y_end), f"Mask {len(sections)}")
+        # both start and end point get
+        return (img, sections)
+    else:
+        point_width = int(img.shape[0]*0.05)
+        sections.append(((ROI_coordinates['x_new'], ROI_coordinates['y_new'],
+                          ROI_coordinates['x_new'] + point_width, ROI_coordinates['y_new'] + point_width),
+                        f"Click second point for Mask {len(sections) + 1}"))
+        return (img, sections)
+
+def del_select_coordinates(img, evt: gr.SelectData):
+    del sections[evt.index]
+    # recreate section names
+    for i in range(len(sections)):
+        sections[i] = (sections[i][0], f"Mask {i + 1}")
+
+    # last section clicking second point not complete
+    if ROI_coordinates['clicks'] % 2 != 0:
+        if len(sections) == evt.index:
+            # delete last section
+            ROI_coordinates['clicks'] -= 1
+        else:
+            # recreate last section name for second point
+            ROI_coordinates['clicks'] -= 2
+            sections[len(sections) - 1] = (sections[len(sections) - 1][0], f"Click second point for Mask {len(sections) + 1}")
+    else:
+        ROI_coordinates['clicks'] -= 2
+
+    return (img[0], sections)
+
+with gr.Blocks(title="Watermark Anything Demo") as demo:
+    gr.Markdown("""
+    # Watermark Anything Demo
+    This app demonstrates watermark detection and embedding using the Watermark Anything model.
+    Find the project [here](https://github.com/facebookresearch/watermark-anything).
+    """)
+
+    with gr.Tabs():
+        with gr.TabItem("Embed Watermark"):
+            with gr.Row():
+                with gr.Column():
+                    embedding_img = gr.Image(label="Input Image", type="numpy")
+
+                with gr.Column():
+                    embedding_num = gr.Slider(1, 5, value=1, step=1, label="Number of Watermarks")
+                    embedding_type = gr.Radio(["random", "input"], value="random", label="Type", info="Type of watermarks")
+                    embedding_str = gr.Textbox(label="Watermark Text", visible=False, show_copy_button=True)
+                    embedding_loc = gr.Radio(["random", "bounding"], value="random", label="Location", info="Location of watermarks")
+
+                    @gr.render(inputs=embedding_loc)
+                    def show_split(wm_loc):
+                        if wm_loc == "bounding":
+                            embedding_box = gr.AnnotatedImage(
+                                label="ROI",
+                                color_map={
+                                    "ROI of Watermark embedding": "#9987FF",
+                                    "Click second point for ROI": "#f44336"}
+                            )
+
+                            embedding_img.select(
+                                fn=get_select_coordinates,
+                                inputs=[embedding_img, embedding_num],
+                                outputs=embedding_box)
+                            embedding_box.select(
+                                fn=del_select_coordinates,
+                                inputs=embedding_box,
+                                outputs=embedding_box
+                            )
+                        else:
+                            embedding_img.select()
+
+                    embedding_btn = gr.Button("Embed Watermark")
+                    marked_msg = gr.JSON(label="Marked Messages")
+            with gr.Row():
+                marked_image = gr.Image(label="Watermarked Image")
+                marked_mask = gr.Image(label="Position of the watermark")
+
+            def visible_text_label(embedding_type, embedding_num):
+                if embedding_type == "input":
+                    tip = "-".join([f"FFFF-{_}{_}{_}{_}" for _ in range(embedding_num)])
+                    return gr.update(visible=True, label=f"Watermark Text (Format: {tip})")
+                else:
+                    return gr.update(visible=False)
+
+            def check_embedding_str(embedding_str, embedding_num):
+                if not re.match(r"^([0-9A-F]{4}-[0-9A-F]{4}-){%d}[0-9A-F]{4}-[0-9A-F]{4}$" % (embedding_num-1), embedding_str):
+                    tip = "-".join([f"FFFF-{_}{_}{_}{_}" for _ in range(embedding_num)])
+                    gr.Warning(f"Invalid format. Please use {tip}", duration=0)
+                    return gr.update(interactive=False)
+                else:
+                    return gr.update(interactive=True)
+
+            embedding_num.change(
+                fn=visible_text_label,
+                inputs=[embedding_type, embedding_num],
+                outputs=[embedding_str]
+            )
+            embedding_type.change(
+                fn=visible_text_label,
+                inputs=[embedding_type, embedding_num],
+                outputs=[embedding_str]
+            )
+            embedding_str.change(
+                fn=check_embedding_str,
+                inputs=[embedding_str, embedding_num],
+                outputs=[embedding_btn]
+            )
+
+            embedding_btn.click(
+                fn=embed_watermark,
+                inputs=[embedding_img, embedding_num, embedding_type, embedding_str, embedding_loc],
+                outputs=[marked_image, marked_mask, marked_msg]
+            )
+
+        with gr.TabItem("Detect Watermark"):
+            with gr.Row():
+                with gr.Column():
+                    detecting_img = gr.Image(label="Input Image", type="numpy")
+                with gr.Column():
+                    detecting_btn = gr.Button("Detect Watermark")
+                    predicted_messages = gr.JSON(label="Detected Messages")
+            with gr.Row():
+                predicted_mask = gr.Image(label="Predicted Watermark Position")
+                predicted_cluster = gr.Image(label="Watermark Clusters")
+
+            detecting_btn.click(
+                fn=detect_watermark,
+                inputs=[detecting_img],
+                outputs=[predicted_mask, predicted_cluster, predicted_messages]
+            )
+
+demo.launch()

requirements.txt

@@ -0,0 +1,4 @@
+torch==2.5.1
+GitPython==3.1.43
+gradio==5.8.0
+huggingface-hub==0.26.3