gradio_app.py

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()