Spaces:
Running
Running
| 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", | |
| # info="Which type of watermark you want to apply?" | |
| ) | |
| 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() |