ALSv/wmark

1import gradio as gr 2from steganography import Steganography3from util_functions.utils import draw_multiple_line_text, generate_qr_code4from SSL_watermark import encode, decode5 6 7TITLE = """<h2 align="center"> ✍️ Invisible Watermark </h2>"""8 9 10def apply_watermark(radio_button, input_image, watermark_image, watermark_text, watermark_url):11    input_image = input_image.convert('RGB')12    print(f'radio_button: {radio_button}')13    if radio_button == "Image":14        watermark_image = watermark_image.resize((input_image.width, input_image.height)).convert('L').convert('RGB')15        return Steganography().merge(input_image, watermark_image, digit=7)16    elif radio_button == "Text":17        watermark_image = draw_multiple_line_text(input_image.size, watermark_text)18        return Steganography().merge(input_image, watermark_image, digit=7)19    elif radio_button == "QRCode":20        size = min(input_image.width, input_image.height)21        watermark_image = generate_qr_code(watermark_url).resize((size, size)).convert('RGB')22        return Steganography().merge(input_image, watermark_image, digit=7)23    else:24        print('start encoding ssl watermark...')25        return encode(input_image, epochs=5)26 27def extract_watermark(extract_radio_button, input_image_to_extract):28    if extract_radio_button == 'Steganography':29        return Steganography().unmerge(input_image_to_extract.convert('RGB'), digit=7).convert('RGBA')30    else:31        decoded_info = decode(image=input_image_to_extract)32        return draw_multiple_line_text(input_image_size=input_image_to_extract.size, text=decoded_info)33 34 35with gr.Blocks() as demo:36    gr.HTML(TITLE)37    with gr.Tab("Add watermark"):38        with gr.Row():39            with gr.Column():40                gr.Markdown("### Image to apply watermark")41                input_image = gr.Image(type='pil')42                with gr.Blocks():43                    gr.Markdown("### Which type of watermark you want to apply?")44                    radio_button = gr.Radio(45                        choices=["QRCode", "Text", "Image", "SSL Watermark"], 46                        label="Watermark type", 47                        value="QRCode",48                        # info="Which type of watermark you want to apply?"49                    )50                    watermark_url = gr.Textbox(51                        placeholder="URL to generate QR code", 52                        visible=True53                    )54                    watermark_text = gr.Textbox(55                        placeholder="What text you want to use as watermark?",56                        visible=False57                    )58                    watermark_image = gr.Image(59                        type='pil', 60                        visible=False61                    )62                    63                    def update_visibility(radio_value):64                        return { 65                            watermark_image: 66                                { 67                                    "visible":radio_value == "Image", 68                                    "__type__": "update"69                                }, 70                            watermark_text: 71                                {72                                    "visible":radio_value == "Text", 73                                    "__type__": "update"74                                },75                            watermark_url: 76                                {77                                    "visible":radio_value == "QRCode", 78                                    "__type__": "update"79                                }80                            }81 82            with gr.Column():83                gr.Markdown("### Applied watermark image")84                output_image = gr.Image(show_label=False).style(height=500)85        with gr.Row():86            apply_button =gr.Button("Apply")87        88    with gr.Tab("Extract watermark"):89        with gr.Row():90            with gr.Column():91                gr.Markdown("### Image to extract watermark")92                input_image_to_extract = gr.Image(type='pil')93                extract_radio_button = gr.Radio(94                    choices=["Steganography", "SSL Watermark"], 95                    label="Extract methods", 96                    value="Steganography"97                )98            with gr.Column():99                gr.Markdown("### Extracted watermark")100                extracted_watermark = gr.Image(type='pil').style(height=500)101        extract_button = gr.Button("Extract")102    103    radio_button.change(104        fn=update_visibility, 105        inputs=radio_button,106        outputs=[watermark_image, watermark_text, watermark_url]107    )108    apply_button.click(109        fn=apply_watermark, 110        inputs=[radio_button, input_image, watermark_image, watermark_text, watermark_url], 111        outputs=[output_image]112    )113    extract_button.click(114        fn=extract_watermark, 115        inputs=[extract_radio_button, input_image_to_extract], 116        outputs=[extracted_watermark]117    )118 119demo.launch()