import gradio as gr

from fawkes.protection import Fawkes


def run_protection_interface(uploaded_image, mode='low', sd=1e6, 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."

    # Run the protection process
    protector = Fawkes(gpu="0", batch_size=1, mode=mode)
    processed_image = protector.run_protection(uploaded_image, sd=sd, batch_size=1, format=format,
                                               separate_target=separate_target, debug=debug, no_align=no_align)

    if processed_image is not None:
        return processed_image, "Protection process completed."
    else:
        return None, "Protection process failed or no cloaked image generated."


intro_md = """
# Fawkes Utility
This is an app that uses a really cool project from the SAND Lab at University of Chicago called 
[Fawkes](https://sandlab.cs.uchicago.edu/fawkes/).

> Fawkes, an algorithm and software tool (running locally on your computer) that gives individuals the ability to limit
how unknown third parties can track them by building facial recognition models out of their publicly available photos.
At a high level, Fawkes "poisons" models that try to learn what you look like, by putting hidden changes into your photos,
and using them as Trojan horses to deliver that poison to any facial recognition models of you. 
Fawkes takes your personal images and makes tiny, pixel-level changes that are invisible to the human eye,
in a process we call image cloaking.
"""

with gr.Blocks() as demo:
    gr.Markdown(intro_md)
    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")
            mode = gr.Radio(label="Mode", choices=['low', 'mid', 'high'], value='low')
            format = gr.Radio(label="Output Format", choices=['png', 'jpg', 'jpeg'], value='png')
            separate_target = gr.Checkbox(label="Separate Target")
            no_align = gr.Checkbox(label="No Align")
            with gr.Accordion(label='Advanced Config', open=False):
                sd = gr.Slider(label="Penalty Number (SD)", minimum=1e5, maximum=1e7, value=1e6)

    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, mode, sd, format, separate_target, no_align],
            outputs=[output_image, output_text]
            )

if __name__ == "__main__":
    demo.launch()