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

css = """
footer { 
    visibility: hidden; 
}
"""

with gr.Blocks(css=css) as app:
    gr.HTML("""<a href="https://visitorbadge.io/status?path=https%3A%2F%2Ffantos-watermark.hf.space">
               <img src="https://api.visitorbadge.io/api/visitors?path=https%3A%2F%2Ffantos-watermark.hf.space&countColor=%23263759" />
               </a>""")
    
    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("Add Watermark")
                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()