init commit

.gitignore

@@ -0,0 +1 @@
+.idea

app.py

@@ -0,0 +1,74 @@
+import gradio as gr
+import glob
+import os
+from PIL import Image
+import numpy as np
+from fawkes.protection import Fawkes  # Make sure the import path is correct
+
+
+def run_protection_interface(uploaded_image, gpu='0', mode='low', feature_extractor="arcface_extractor_0",
+                             th=0.01, max_step=1000, sd=1e6, lr=2, batch_size=1, format='png',
+                             separate_target=False, no_align=False, debug=False):
+    """
+    Gradio compatible function for running protection.
+    """
+    if uploaded_image is None:
+        return None, "No image uploaded."
+
+    # Save the uploaded image to a temporary directory
+    temp_dir = "temp_imgs"
+    os.makedirs(temp_dir, exist_ok=True)
+    img_path = os.path.join(temp_dir, "uploaded_image.png")
+    uploaded_image.save(img_path)
+
+    # Run the protection process
+    protector = Fawkes(feature_extractor, gpu, batch_size, mode=mode)
+    image_paths = [img_path]
+    protector.run_protection(image_paths, th=th, sd=sd, lr=lr,
+                             max_step=max_step, batch_size=batch_size,
+                             format=format, separate_target=separate_target,
+                             debug=debug, no_align=no_align)
+
+    # Load and return the processed image
+    processed_img_path = img_path.replace(".png", "_cloaked.png")
+    if os.path.exists(processed_img_path):
+        processed_image = Image.open(processed_img_path)
+        return processed_image, "Protection process completed."
+    else:
+        return None, "Protection process failed or no cloaked image generated."
+
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("### Upload Image")
+            uploaded_image = gr.Image(type="pil", label="Upload Image")
+
+        with gr.Column():
+            gr.Markdown("### Configuration Options")
+            gpu = gr.Textbox(label="GPU", value='0')
+            mode = gr.Dropdown(label="Mode", choices=['low', 'mid', 'high'], value='low')
+            feature_extractor = gr.Textbox(label="Feature Extractor", value="arcface_extractor_0")
+            th = gr.Slider(label="Threshold", minimum=0.001, maximum=0.05, value=0.01)
+            max_step = gr.Slider(label="Max Steps", minimum=500, maximum=2000, value=1000)
+            sd = gr.Slider(label="Penalty Number (SD)", minimum=1e5, maximum=1e7, value=1e6)
+            lr = gr.Slider(label="Learning Rate", minimum=1, maximum=25, value=2)
+            batch_size = gr.Slider(label="Batch Size", minimum=1, maximum=10, value=1)
+            format = gr.Radio(label="Output Format", choices=['png', 'jpg', 'jpeg'])
+            separate_target = gr.Checkbox(label="Separate Target")
+            no_align = gr.Checkbox(label="No Align")
+            debug = gr.Checkbox(label="Debug")
+
+    run_button = gr.Button("Run Protection")
+    output_image = gr.Image(label="Processed Image")
+    output_text = gr.Textbox(label="Output Message")
+
+    run_button.click(
+            fn=run_protection_interface,
+            inputs=[uploaded_image, gpu, mode, feature_extractor, th, max_step, sd, lr, batch_size, format,
+                    separate_target, no_align, debug],
+            outputs=[output_image, output_text]
+            )
+
+if __name__ == "__main__":
+    demo.launch()