import gradio as gr
import os

css = """
p {
    font-size: 120%;
}

li {
    font-size: 110%;
}

html, body {
    overflow: scroll;
}

video {
    max-height: 400px;
}

.container {
    height: initial;
}


.image-container {
    width: 200px; 
    max-height: auto;
    margin: auto;
}

img {
    max-height: 400px;
}
"""

# Optional CSS stuff for the header example image:
#example {
#    width: 80%;
#    height: 60%
#}

#example img {
#    width: 80%;
#    height: 80%
#}


a = os.path.join(os.path.dirname(__file__), "files/barkley_balloon.mp4")
b = os.path.join(os.path.dirname(__file__), "files/eiffel_tower.mp4")
c = os.path.join(os.path.dirname(__file__), "files/bird.bmp")
d = os.path.join(os.path.dirname(__file__), "files/groot.jpeg")
w1 = os.path.join(os.path.dirname(__file__), "files/AI_generated.png")
w2 = os.path.join(os.path.dirname(__file__), "files/hf-logo.png")
w3 = os.path.join(os.path.dirname(__file__), "files/forest_qr_watermarking.png")
w4 = os.path.join(os.path.dirname(__file__), "files/cheetah1.jpg")
w5 = os.path.join(os.path.dirname(__file__), "files/frog.jpg")
w6 = os.path.join(os.path.dirname(__file__), "files/Human_generated.png")
w7 = os.path.join(os.path.dirname(__file__), "files/hf-logo_transpng.png")

def generate_image(original_image, watermark):
    return gr.Image(original_image, watermark=watermark)

def generate_video(original_video, watermark):
    return gr.Video(original_video, watermark=watermark)


with gr.Blocks(css=css) as demo:
    with gr.Row():
        with gr.Column():
            gr.Markdown("# 🤗 Watermarking with Gradio: Example")
            gr.Markdown("Watermarks can be **visible** or **invisible**.")
            gr.Markdown("""They can provide information directly, or provide a link for more information.  
                           - Visible watermarks are useful to disclose when content is AI-generated.  
                           - Invisible watermarks can mark content as authentic.
                           - ...And vice versa! There are many possibilities for what watermarks can provide.
                           - Watermarks can also provide information about content created by people.""")
            gr.Markdown("They are a useful tool for **AI provenance**.")
            gr.Markdown()
            gr.Markdown("""For more information on watermarking -- what watermarking is, why it's important, and the tools available on Hugging Face -- please check out [our blogpost on AI watermarking](https://huggingface.co/blog/watermarking).""")
            gr.Markdown()
            gr.Markdown("## Try it out below!")
        with gr.Column():
            with gr.Column():
                gr.Image('files/watermark_example.png', visible=False)
            with gr.Column():
                gr.Image('files/watermark_example.png', show_label=False, show_download_button=False, elem_id='example', container=False, interactive=False)
                gr.Markdown('**Image Watermark Code:**')
                gr.Code('import gradio as gr\n\nwatermarked_image = gr.Image(original_image_file, watermark=watermark_file)', lines=3)
                gr.Markdown('**Video Watermark Code:**')
                gr.Code('import gradio as gr\n\nwatermarked_video = gr.Video(original_video_file, watermark=watermark_file)', lines=3)
            with gr.Column():
                gr.Image('files/watermark_example.png', visible=False)
    with gr.Tab("Image Watermarking"):
        with gr.Column():
            gr.Markdown("**Inputs**: Image and watermark file")
        with gr.Column():
            gr.Markdown("**Output**: Watermarked image")
        gr.Interface(generate_image, [gr.Image(), gr.Image(type='filepath', image_mode=None)], gr.Image(),
                 examples=[[d, w7], [w4, w5], [c, w6],])
    with gr.Tab("Video Watermarking"):
        with gr.Column():
            gr.Markdown("**Inputs**: Video and watermark file")
        with gr.Column():
            gr.Markdown("**Output**: Watermarked video")
        gr.Interface(generate_video, [gr.Video(), gr.Image(type='filepath', image_mode=None)], gr.Video(),
                 examples=[[a, w1], [b, w2], [a, w3], [b, w4]])
    
if __name__ == "__main__":
    demo.launch()