import numpy as np import gradio as gr from run_gaussian_shading import * examples = [ "A photo of a cat", "A pizza with pineapple on it", "A photo of dog", ] css = """ #col-container { margin: 0 auto; max-width: 700px; } """ MAX_SEED = np.iinfo(np.int32).max #--------------------------------------------------------------------------------------------------- with gr.Blocks(css=css) as demo: # ---------------------------------- Add Watermark ----------------------------------------- with gr.Tab("Add watermark"): with gr.Column(elem_id="col-container"): gr.Markdown(" # Text-to-Image Watermark") with gr.Accordion("Instruction", open=False): gr.Markdown(""" # Embedding Watermark ## 1. Generate watermarked image * Enter your prompt in the text box. * Click **Run** to generate an image with a random binary watermark. ## 2. Save Image Click **Download** to save the watermarked image in PNG format ## 3. Advanced Settings - **Seed**: Generates different images with different seed. - **Guidance Scale**: Higher values give the model more freedom in image creation. - **Num Inference Steps**: More steps enhance image detail and quality but increase computational cost. Source code: [Gaussian Shading](https://github.com/bsmhmmlf/Gaussian-Shading)""") with gr.Row(): prompt = gr.Text( label="Prompt", show_label=False, max_lines=1, placeholder="Enter your prompt", container=False, ) run_button = gr.Button("Run", scale=0, variant="primary") download_button = gr.DownloadButton(visible=True) with gr.Row(): result_original = gr.Image(label="Image without watermark", show_label=True) result = gr.Image(label="Watermarked Image", show_label=True) with gr.Accordion("Advanced Settings", open=False): seed = gr.Slider( label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, ) with gr.Row(): guidance_scale = gr.Slider( label="Guidance scale", minimum=1.5, maximum=10, step=0.1, value=7.5, ) num_inference_steps = gr.Slider( label="Num inference steps", minimum=10, maximum=100, step=1, value=50, ) gr.Examples(examples=examples, inputs=[prompt]) # ---------------------------------- Extract Watermark ----------------------------------------- with gr.Tab("Extract watermark"): with gr.Column(elem_id="col-container"): gr.Markdown(" # Watermark Extraction") with gr.Accordion("Instruction", open=False): gr.Markdown(""" # Extracting Watermark **Note**: Ensure you create an image first to add the watermark to the database. ## 1. Upload Image - Upload the image to the Image box. - Click the **Extract** button to extract the watermark. ## 2. Advanced Settings These settings are **optional** and can be used to simulate real-world attacks to erase the watermark: Click the **Attack** button to generate a distorted image. * **Seed**: Initialize the random number generator, ensuring reproducibility of the attack * **Random crop ratio**: determines the proportion of the image to be randomly cropped. A lower ratio means more of the image will be cropped. * **Random drop ratio**: specifies the fraction of pixels to be randomly dropped. A higher ratio increases the number of dropped pixels. * **Resize ratio**: determines how much the image will be resized. A lower ratio means the image will be reduced more significantly. * **Gaussian blur R**: the radius of the Gaussian blur applied to the image. A larger radius results in a more blurred image. * **Gaussian Std**: standard deviation of the Gaussian distribution used for blurring. A higher value results in a stronger blur effect. * **Sp prob**: the probability of each pixel being replaced with either black or white noise. A higher probability increases the amount of noise added to the image. ## Output Explanation - **Output watermark**: The binary bit embedding in the image. - **Accuracy bit**: The number of binary bits extracted that match the binary watermark in the database. """) with gr.Row(): input_image = gr.Image(type='pil') extract_button = gr.Button("Extract", scale=0, variant="primary") with gr.Accordion("Advanced Settings", open=False): with gr.Row(): seed = gr.Slider( label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, ) attack_button = gr.Button("Attack!", scale=0, variant="primary") with gr.Row(): random_crop_ratio = gr.Slider( label="Random crop ratio", minimum=0.5, maximum=1, step=0.1, value=1, ) random_drop_ratio = gr.Slider( label="Random drop ratio", minimum=0, maximum=1, step=0.1, value=0, ) with gr.Row(): resize_ratio = gr.Slider( label="Resize ratio", minimum=0.2, maximum=1, step=0.1, value=1, ) gaussian_blur_r = gr.Slider( label="Gaussian blur r", minimum=0, maximum=1, step=0.1, value=0, ) with gr.Row(): gaussian_std = gr.Slider( label="Gaussian std", minimum=0, maximum=0.01, step=0.0001, value=0, ) sp_prob = gr.Slider( label="Sp prob", minimum=0, maximum=0.1, step=0.001, value=0, ) attack_image = gr.Image(label="Attacked Image") output = gr.Textbox(label="Output") with gr.Accordion("More Details", open=False): result_extract = gr.Textbox(label="Bit watermark") accuracy_bit = gr.Textbox(label="Accuracy bit") # ----------------------------- Embedding watermark ------------------------- gr.on( triggers=[run_button.click, prompt.submit], fn=generate_with_watermark, inputs=[ seed, prompt, guidance_scale, num_inference_steps ], outputs=[result_original, result, download_button], ) # ----------------------------- Extract watermark ------------------------- gr.on( triggers=[extract_button.click, attack_button.click], fn=reverse_watermark, inputs=[ input_image, seed, random_crop_ratio, random_drop_ratio, resize_ratio, gaussian_blur_r, gaussian_std, sp_prob, ], outputs=[output, result_extract, accuracy_bit, attack_image], ) demo.launch(share=True)