Spaces:
Running
Running
| 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() | |