finish

.gitignore

@@ -159,4 +159,6 @@ cython_debug/
 #  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
+#.idea/
+
+.env

app.py

@@ -1,24 +1,25 @@
+import os
 import gradio as gr
+
 import numpy as np
-import matplotlib.pyplot as plt
+import json
+import redis
 from PIL import Image
 import time
+import plotly.graph_objects as go
 from datetime import datetime
-import json
-import redis
+from kit import compute_performance, compute_quality
+import dotenv
+
+dotenv.load_dotenv()
 
-# Redis configuration
-REDIS_HOST = "redis-17482.c73.us-east-1-2.ec2.redns.redis-cloud.com"
-REDIS_PORT = 17482
-REDIS_USERNAME = "default"
-REDIS_PASSWORD = "7gdBhkyX1duixq5J3Mr14kYxUhWugoYk"
 
 # Connect to Redis
 redis_client = redis.Redis(
-    host=REDIS_HOST,
-    port=REDIS_PORT,
-    username=REDIS_USERNAME,
-    password=REDIS_PASSWORD,
+    host=os.getenv("REDIS_HOST"),
+    port=os.getenv("REDIS_PORT"),
+    username=os.getenv("REDIS_USERNAME"),
+    password=os.getenv("REDIS_PASSWORD"),
     decode_responses=True,
 )
 
@@ -38,54 +39,65 @@ def get_submissions_from_redis():
     return [json.loads(submission) for submission in submissions]
 
 
-def decode(image):
-    # TODO: Replace with actual decode function
-    time.sleep(2)  # Simulate processing time
-    return np.random.random()
-
-
-def metric(image):
-    # TODO: Replace with actual metric function
-    time.sleep(2)  # Simulate processing time
-    return np.random.random()
-
-
 def update_leaderboard(submissions):
     names = [sub["name"] for sub in submissions]
     performances = [float(sub["performance"]) for sub in submissions]
     qualities = [float(sub["quality"]) for sub in submissions]
 
-    fig, ax = plt.subplots(figsize=(6, 6))
-    ax.scatter(qualities, performances)
-    for i, name in enumerate(names):
-        ax.annotate(name, (qualities[i], performances[i]))
+    # Create scatter plot
+    fig = go.Figure()
+
+    fig.add_trace(
+        go.Scatter(
+            x=qualities,
+            y=performances,
+            mode="markers+text",
+            text=names,
+            textposition="top center",
+            name="Submissions",
+        )
+    )
 
-    # Plot circles
+    # Add circles
     circle_radii = np.linspace(0, 1, 5)
     for radius in circle_radii:
-        circle = plt.Circle((0, 0), radius, fill=False, linestyle="--", color="gray")
-        ax.add_artist(circle)
+        theta = np.linspace(0, 2 * np.pi, 100)
+        x = radius * np.cos(theta)
+        y = radius * np.sin(theta)
+        fig.add_trace(
+            go.Scatter(
+                x=x,
+                y=y,
+                mode="lines",
+                line=dict(color="gray", dash="dash"),
+                showlegend=False,
+            )
+        )
 
-    ax.set_xlabel("Quality")
-    ax.set_ylabel("Performance")
-    ax.set_title("Submissions Leaderboard")
-    ax.set_xlim(0, 1)
-    ax.set_ylim(0, 1)
+    # Update layout
+    fig.update_layout(
+        title="Submissions Leaderboard",
+        xaxis_title="Quality",
+        yaxis_title="Performance",
+        xaxis=dict(range=[0, 1]),
+        yaxis=dict(range=[0, 1]),
+        width=600,
+        height=600,
+    )
 
     return fig
 
 
 def process_submission(name, image):
-    # Simulating progress
+    original_image = Image.open("./image.png")
     progress = gr.Progress()
     progress(0, desc="Processing")
-    time.sleep(1)
-
-    progress(0.25, desc="Decoding")
-    performance = decode(image)
-    progress(0.5, desc="Computing metric")
-    quality = metric(image)
-    progress(0.75, desc="Saving results")
+    time.sleep(0.5)
+    progress(0.1, desc="Decoding")
+    performance = compute_performance(image)
+    progress(0.6, desc="Computing metric")
+    quality = compute_quality(image, original_image)
+    progress(0.9, desc="Saving results")
     save_to_redis(name, performance, quality)
 
     submissions = get_submissions_from_redis()
@@ -176,7 +188,7 @@ def create_interface():
 
         submit_btn.click(lambda: gr.Tabs(selected="submit"), None, tabs)
 
-        upload_btn.click(
+        upload_btn.click(lambda: gr.Tabs(selected="leaderboard"), None, tabs).then(
             upload_and_evaluate,
             inputs=[name_input, uploaded_image],
             outputs=[
@@ -186,7 +198,7 @@ def create_interface():
                 performance_output,
                 quality_output,
             ],
-        ).success(lambda: gr.Tabs(selected="leaderboard"), None, tabs)
+        )
 
         demo.load(
             lambda: [

kit/{metric.py → __init__.py}

@@ -1,7 +1,74 @@
 import os
-import torch
-from metrics import (
-    compute_fid,
+import io
+import numpy as np
+import onnxruntime as ort
+from PIL import Image
+import dotenv
+
+dotenv.load_dotenv()
+
+GT_MESSAGE = os.environ["GT_MESSAGE"]
+
+
+QUALITY_COEFFICIENTS = {
+    "psnr": -0.0022186489180419534,
+    "ssim": -0.11337077856710862,
+    "nmi": -0.09878221979274945,
+    "lpips": 0.3412626374646173,
+    "aesthetics": 0.044982716146790995,
+    "artifacts": -0.1437622439320745,
+}
+
+QUALITY_OFFSETS = {
+    "psnr": 43.54757854447622,
+    "ssim": 0.984229018845295,
+    "nmi": 1.7536553655336136,
+    "lpips": 0.014247652621287854,
+    "aesthetics": -0.005184057521820067,
+    "artifacts": -0.0015724570911377733,
+}
+
+
+def compute_performance(image):
+    session_options = ort.SessionOptions()
+    session_options.intra_op_num_threads = 1
+    session_options.inter_op_num_threads = 1
+    session_options.log_severity_level = 3
+    model = ort.InferenceSession(
+        "./kit/models/stable_signature.onnx",
+        sess_options=session_options,
+    )
+    inputs = np.stack(
+        [
+            (
+                (
+                    np.array(
+                        image,
+                        dtype=np.float32,
+                    )
+                    / 255.0
+                    - [0.485, 0.456, 0.406]
+                )
+                / [0.229, 0.224, 0.225]
+            )
+            .transpose((2, 0, 1))
+            .astype(np.float32)
+        ],
+        axis=0,
+    )
+
+    outputs = model.run(
+        None,
+        {
+            "image": inputs,
+        },
+    )
+    decoded = (outputs[0] > 0).astype(int)[0]
+    gt_message = np.array([int(bit) for bit in GT_MESSAGE])
+    return 1 - np.mean(gt_message != decoded)
+
+
+from .metrics import (
     compute_image_distance_repeated,
     load_perceptual_models,
     compute_perceptual_metric_repeated,
@@ -9,100 +76,44 @@ from metrics import (
     compute_aesthetics_and_artifacts_scores,
 )
 
-QUALITY_METRICS = {
-    "legacy_fid": "Legacy FID",
-    "clip_fid": "CLIP FID",
-    "psnr": "PSNR",
-    "ssim": "SSIM",
-    "nmi": "Normed Mutual-Info",
-    "lpips": "LPIPS",
-    "aesthetics": "Delta Aesthetics",
-    "artifacts": "Delta Artifacts",
-}
-
 
-QUALITY_COEFFICIENTS = {
-    "legacy_fid": 1.52e-3,
-    "clip_fid": 5.07e-3,
-    "psnr": -2.22e-3,
-    "ssim": -1.13e-1,
-    "nmi": -9.88e-2,
-    "lpips": 3.41e-1,
-    "aesthetics": 4.50e-2,
-    "artifacts": -1.44e-1,
-}
+def compute_quality(attacked_image, clean_image, quiet=True):
 
+    # Compress the image
+    buffer = io.BytesIO()
+    attacked_image.save(buffer, format="JPEG", quality=90)
+    buffer.seek(0)
 
-def compute_q(quality_dict):
-    result = 0
-    for key, value in quality_dict.items():
-        result += value[0] * QUALITY_COEFFICIENTS[key]
-    return result
+    # Update attacked_image with the compressed version
+    attacked_image = Image.open(buffer)
 
-
-def metric(attacked_image, clean_image, quiet=True):
-    modes = QUALITY_METRICS.keys()
-    modes = (
-            [
-                mode
-                for mode in modes
-                if mode not in ["aesthetics", "artifacts", "clip_score"]
-            ]
-            + ["aesthetics_and_artifacts"]
-            + (["clip_score"] if "clip_score" in modes else [])
-        )
+    modes = ["psnr", "ssim", "nmi", "lpips", "aesthetics_and_artifacts"]
 
     results = {}
     for mode in modes:
-        if mode.endswith("_fid"):
-            results[mode] = [0.0]
-        #     num_gpus = torch.cuda.device_count()
-        #     if num_gpus == 0:
-        #         raise RuntimeError("No GPUs available for processing")
-        #     if not quiet:
-        #         print(f"Using {num_gpus} GPUs for processing")
-        #     metric = float(
-        #         compute_fid(
-        #             clean_path,
-        #             attacked_path,
-        #             mode=mode.split("_")[0],
-        #             device=torch.device("cuda"),
-        #             batch_size=64,
-        #             num_workers=8,
-        #             verbose=not quiet,
-        #         )
-        #     )
-        #     return {
-        #         idx: metric
-        #         for idx in existence_to_indices(
-        #             check_file_existence(attacked_path, name_pattern="{}.png", limit=limit),
-        #             limit=limit,
-        #         )
-        #     }
-
-        elif mode in ["psnr", "ssim", "nmi"]:
+        if mode in ["psnr", "ssim", "nmi"]:
             metrics = compute_image_distance_repeated(
                 [clean_image],
                 [attacked_image],
                 metric_name=mode,
-                num_workers=8,
+                num_workers=1,
                 verbose=not quiet,
             )
             results[mode] = metrics
 
-        elif mode in ["lpips", "watson"]:
+        elif mode == "lpips":
             model = load_perceptual_models(
                 mode,
-                mode="alex" if mode == "lpips" else "dft",
-                device="cpu",#torch.device("cuda"),
+                mode="alex",
+                device="cpu",
             )
             metrics = compute_perceptual_metric_repeated(
                 [clean_image],
                 [attacked_image],
                 metric_name=mode,
-                mode="alex" if mode == "lpips" else "dft",
+                mode="alex",
                 model=model,
-                device="cpu",#torch.device("cuda"),
+                device="cpu",
             )
             results[mode] = metrics
 
@@ -111,14 +122,19 @@ def metric(attacked_image, clean_image, quiet=True):
             aesthetics, artifacts = compute_aesthetics_and_artifacts_scores(
                 [attacked_image],
                 model,
-                device="cpu",#torch.device("cuda"),
+                device="cpu",
             )
             clean_aesthetics, clean_artifacts = compute_aesthetics_and_artifacts_scores(
                 [attacked_image],
                 model,
-                device="cpu",#torch.device("cuda"),
+                device="cpu",
             )
-            results["aesthetics"] = [clean_aesthetics[0]-aesthetics[0]]
-            results["artifacts"] = [clean_artifacts[0]-artifacts[0]]
-    
-    return compute_q(results)
+            results["aesthetics"] = [clean_aesthetics[0] - aesthetics[0]]
+            results["artifacts"] = [clean_artifacts[0] - artifacts[0]]
+
+    normalized_quality = 0
+    for key, value in results.items():
+        normalized_quality += (value[0] - QUALITY_OFFSETS[key]) * QUALITY_COEFFICIENTS[
+            key
+        ]
+    return normalized_quality

kit/decode.py

@@ -1,80 +0,0 @@
-import os
-import numpy as np
-import onnxruntime as ort
-from PIL import ImageOps
-
-
-def init_model(mode, gpu):
-    if mode == "stable_sig":
-        session_options = ort.SessionOptions()
-        session_options.intra_op_num_threads = 1
-        session_options.inter_op_num_threads = 1
-        session_options.log_severity_level = 3
-        return ort.InferenceSession(
-            os.path.join("models", "stable_signature.onnx"), # CHENGHAO: may need to modifiy the path
-            #providers=["CUDAExecutionProvider"],
-            #provider_options=[{"device_id": str(gpu)}],
-            sess_options=session_options,
-        )
-    elif mode == "stegastamp":
-        session_options = ort.SessionOptions()
-        session_options.intra_op_num_threads = 1
-        session_options.inter_op_num_threads = 1
-        session_options.log_severity_level = 3
-        return ort.InferenceSession(
-            os.path.join("models", "stega_stamp.onnx"), # CHENGHAO: may need to modifiy the path
-            #providers=["CUDAExecutionProvider"],
-            #provider_options=[{"device_id": str(gpu)}],
-            sess_options=session_options,
-        )
-
-
-def decode(image, mode="stable_sig"):
-    model = init_model(mode, 0)
-    inputs = np.stack(
-            [
-                (
-                    (
-                        np.array(
-                            image,
-                            dtype=np.float32,
-                        )
-                        / 255.0
-                        - [0.485, 0.456, 0.406]
-                    )
-                    / [0.229, 0.224, 0.225]
-                )
-                .transpose((2, 0, 1))
-                .astype(np.float32)
-            ],
-            axis=0,
-        ) if mode == "stable_sig" else np.stack(
-            [
-                np.array(
-                    ImageOps.fit(
-                        image, (400, 400)
-                    ),
-                    dtype=np.float32,
-                )
-                / 255.0
-            ],
-            axis=0,
-        ) # if mode == "stegastamp"
-
-    if mode == "stable_sig":
-        outputs = model.run(
-            None,
-            {
-                "image": inputs,
-            },
-        )
-        return (outputs[0] > 0).astype(int)[0]
-    elif mode == "stegastamp":
-        outputs = model.run(
-            None,
-            {
-                "image": inputs,
-                "secret": np.zeros((inputs.shape[0], 100), dtype=np.float32),
-            },
-        )
-        return outputs[2].astype(int)[0]

kit/metrics/perceptual.py

@@ -1,12 +1,37 @@
 import torch
-import numpy as np
 from PIL import Image
-from tqdm.auto import tqdm
-from utils import to_tensor
+from torchvision import transforms
 from .lpips import LPIPS
 from .watson import LossProvider
 
 
+# Normalize image tensors
+def normalize_tensor(images, norm_type):
+    assert norm_type in ["imagenet", "naive"]
+    # Two possible normalization conventions
+    if norm_type == "imagenet":
+        mean = [0.485, 0.456, 0.406]
+        std = [0.229, 0.224, 0.225]
+        normalize = transforms.Normalize(mean, std)
+    elif norm_type == "naive":
+        mean = [0.5, 0.5, 0.5]
+        std = [0.5, 0.5, 0.5]
+        normalize = transforms.Normalize(mean, std)
+    else:
+        assert False
+    return torch.stack([normalize(image) for image in images])
+
+
+def to_tensor(images, norm_type="naive"):
+    assert isinstance(images, list) and all(
+        [isinstance(image, Image.Image) for image in images]
+    )
+    images = torch.stack([transforms.ToTensor()(image) for image in images])
+    if norm_type is not None:
+        images = normalize_tensor(images, norm_type)
+    return images
+
+
 def load_perceptual_models(metric_name, mode, device=torch.device("cuda")):
     assert metric_name in ["lpips", "watson"]
     if metric_name == "lpips":

kit/test.py

@@ -1,20 +0,0 @@
-from decode import decode
-from metric import metric
-from PIL import Image
-
-def main(attacked_image_path):
-    attacked_image = Image.open(attacked_image_path)
-    clean_image = Image.open("../image.png")
-    
-    decode_result = decode(
-        attacked_image,
-    )
-    print("Decode: ", decode_result)
-    metric_result = metric(
-        attacked_image, 
-        clean_image, 
-    )
-    print("Metric: ", metric_result)
-
-if __name__ == "__main__":
-    main("../attacked_image.png")

requirements.txt

@@ -1,9 +1,14 @@
 gradio
+onnxruntime
+torch
+torchvision
+transformers
+open_clip_torch
 numpy
-matplotlib
 Pillow
-scikit-image
 redis
+plotly
+python-dotenv
 
 # dev
 jupyterlab

test.ipynb

@@ -15,22 +15,20 @@
    ],
    "source": [
     "# Print all submissions in Redis\n",
-    "\n",
+    "import os\n",
     "import redis\n",
     "import json\n",
+    "import dotenv\n",
+    "\n",
+    "dotenv.load_dotenv()\n",
     "\n",
-    "# Redis configuration\n",
-    "REDIS_HOST = \"redis-17482.c73.us-east-1-2.ec2.redns.redis-cloud.com\"\n",
-    "REDIS_PORT = 17482\n",
-    "REDIS_USERNAME = \"default\"\n",
-    "REDIS_PASSWORD = \"7gdBhkyX1duixq5J3Mr14kYxUhWugoYk\"\n",
     "\n",
     "# Connect to Redis\n",
     "redis_client = redis.Redis(\n",
-    "    host=REDIS_HOST,\n",
-    "    port=REDIS_PORT,\n",
-    "    username=REDIS_USERNAME,\n",
-    "    password=REDIS_PASSWORD,\n",
+    "    host=os.getenv(\"REDIS_HOST\"),\n",
+    "    port=os.getenv(\"REDIS_PORT\"),\n",
+    "    username=os.getenv(\"REDIS_USERNAME\"),\n",
+    "    password=os.getenv(\"REDIS_PASSWORD\"),\n",
     "    decode_responses=True,\n",
     ")\n",
     "\n",