|
import gradio as gr |
|
from steganography import Steganography |
|
from util_functions.utils import draw_multiple_line_text, generate_qr_code |
|
from SSL_watermark import encode, decode |
|
|
|
|
|
TITLE = """<h2 align="center"> ✍️ Invisible Watermark </h2>""" |
|
|
|
|
|
def apply_watermark(radio_button, input_image, watermark_image, watermark_text, watermark_url): |
|
input_image = input_image.convert('RGB') |
|
print(f'radio_button: {radio_button}') |
|
if radio_button == "Image": |
|
watermark_image = watermark_image.resize((input_image.width, input_image.height)).convert('L').convert('RGB') |
|
return Steganography().merge(input_image, watermark_image, digit=7) |
|
elif radio_button == "Text": |
|
watermark_image = draw_multiple_line_text(input_image.size, watermark_text) |
|
return Steganography().merge(input_image, watermark_image, digit=7) |
|
elif radio_button == "QRCode": |
|
size = min(input_image.width, input_image.height) |
|
watermark_image = generate_qr_code(watermark_url).resize((size, size)).convert('RGB') |
|
return Steganography().merge(input_image, watermark_image, digit=7) |
|
else: |
|
print('start encoding ssl watermark...') |
|
return encode(input_image, epochs=5) |
|
|
|
def extract_watermark(extract_radio_button, input_image_to_extract): |
|
if extract_radio_button == 'Steganography': |
|
return Steganography().unmerge(input_image_to_extract.convert('RGB'), digit=7).convert('RGBA') |
|
else: |
|
decoded_info = decode(image=input_image_to_extract) |
|
return draw_multiple_line_text(input_image_size=input_image_to_extract.size, text=decoded_info) |
|
|
|
|
|
with gr.Blocks() as demo: |
|
gr.HTML(TITLE) |
|
with gr.Tab("Add watermark"): |
|
with gr.Row(): |
|
with gr.Column(): |
|
gr.Markdown("### Image to apply watermark") |
|
input_image = gr.Image(type='pil') |
|
with gr.Blocks(): |
|
gr.Markdown("### Which type of watermark you want to apply?") |
|
radio_button = gr.Radio( |
|
choices=["QRCode", "Text", "Image", "SSL Watermark"], |
|
label="Watermark type", |
|
value="QRCode", |
|
|
|
) |
|
watermark_url = gr.Textbox( |
|
placeholder="URL to generate QR code", |
|
visible=True |
|
) |
|
watermark_text = gr.Textbox( |
|
placeholder="What text you want to use as watermark?", |
|
visible=False |
|
) |
|
watermark_image = gr.Image( |
|
type='pil', |
|
visible=False |
|
) |
|
|
|
def update_visibility(radio_value): |
|
return { |
|
watermark_image: |
|
{ |
|
"visible":radio_value == "Image", |
|
"__type__": "update" |
|
}, |
|
watermark_text: |
|
{ |
|
"visible":radio_value == "Text", |
|
"__type__": "update" |
|
}, |
|
watermark_url: |
|
{ |
|
"visible":radio_value == "QRCode", |
|
"__type__": "update" |
|
} |
|
} |
|
|
|
with gr.Column(): |
|
gr.Markdown("### Applied watermark image") |
|
output_image = gr.Image(show_label=False).style(height=500) |
|
with gr.Row(): |
|
apply_button =gr.Button("Apply") |
|
|
|
with gr.Tab("Extract watermark"): |
|
with gr.Row(): |
|
with gr.Column(): |
|
gr.Markdown("### Image to extract watermark") |
|
input_image_to_extract = gr.Image(type='pil') |
|
extract_radio_button = gr.Radio( |
|
choices=["Steganography", "SSL Watermark"], |
|
label="Extract methods", |
|
value="Steganography" |
|
) |
|
with gr.Column(): |
|
gr.Markdown("### Extracted watermark") |
|
extracted_watermark = gr.Image(type='pil').style(height=500) |
|
extract_button = gr.Button("Extract") |
|
|
|
radio_button.change( |
|
fn=update_visibility, |
|
inputs=radio_button, |
|
outputs=[watermark_image, watermark_text, watermark_url] |
|
) |
|
apply_button.click( |
|
fn=apply_watermark, |
|
inputs=[radio_button, input_image, watermark_image, watermark_text, watermark_url], |
|
outputs=[output_image] |
|
) |
|
extract_button.click( |
|
fn=extract_watermark, |
|
inputs=[extract_radio_button, input_image_to_extract], |
|
outputs=[extracted_watermark] |
|
) |
|
|
|
demo.launch() |
