|
|
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: |
|
|
|
|
|
|
|
|
|
|
|
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]) |
|
|
|
|
|
|
|
|
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") |
|
|
|
|
|
|
|
|
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], |
|
|
) |
|
|
|
|
|
|
|
|
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) |
