import gradio as gr
import utils

def choose_encode(inp_im,inp_mark,cho):
    if cho == "stegan":
        out_im, out_msg = utils.encode(inp_im,inp_mark)
        return out_im,out_msg
    if cho == "pnginfo":
        out_im, out_msg = utils.png_encode(inp_im,inp_mark)
        return out_im,out_msg

with gr.Blocks() as app:
    with gr.Tab("Add Watermark"):
        cho = gr.Radio(choices=["stegan","pnginfo"],value="stegan")
        with gr.Row():
            with gr.Column():
                inp_im = gr.Image(label="Input Image",type="filepath")
                inp_mark = gr.Textbox(label="Watermark")
                mark_btn = gr.Button()
                msg_box = gr.Textbox(label="System Message")
            with gr.Column():
                out_im = gr.Image(label="Watermarked Image")
    with gr.Tab("Detect Watermark"):
        with gr.Row():
            with gr.Column():
                det_im = gr.Image(label="Watermarked Image",type="filepath")
                det_btn = gr.Button("Detect")
            with gr.Column():
                det_msg = gr.Textbox(label="Detected Watermark",lines=6, max_lines=50)
    mark_btn.click(choose_encode,[inp_im,inp_mark,cho],[out_im,msg_box])
    det_btn.click(utils.decode,[det_im],det_msg)
app.launch()