Synced repo using 'sync_with_huggingface' Github Action

app.py

@@ -1,44 +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 = '45d56c2b61f2bc73caeafc90e14df33ad50b238c'
-LOCAL_PATH = "./watermark-anything"
-MODEL_ID = "facebook/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
+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 = '45d56c2b61f2bc73caeafc90e14df33ad50b238c'
+LOCAL_PATH = "./watermark-anything"
+MODEL_ID = "facebook/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

@@ -1,393 +1,393 @@
-import os
-import sys
-if "APP_PATH" in os.environ:
-    os.chdir(os.environ["APP_PATH"])
-    # fix sys.path for import
-    sys.path.append(os.getcwd())
-
-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()
+import os
+import sys
+if "APP_PATH" in os.environ:
+    os.chdir(os.environ["APP_PATH"])
+    # fix sys.path for import
+    sys.path.append(os.getcwd())
+
+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()