import gradio as gr
from utils import encoder, decoder, decoder_noise

title = "**Playing with kNN for denoising**"
description = """The encoder takes any image (on the top left) and remove a fraction of the pixels producing a noisy image (on the top right).  
The decoder takes a noisy image (on the bottom left) and achieves denoising using kNN where you can choose the number of neighbors k.  
You can directly use the decoder on the encoded image (if you used the encoder first!).  
If you want to use the webcam (instead of upload), click on the webcam Tab.
"""

demo = gr.Blocks()

with demo:
    gr.Markdown(title)
    gr.Markdown(description)
    with gr.Tab("upload"):
        with gr.Row():
            with gr.Column():
                source_img = gr.Image(source="upload", type="filepath", label="init_img")
                
                noise = gr.Slider(label='noise', minimum = 0.5, maximum = 1, step = .005, value = .95)
                with gr.Row():
                    b1 = gr.Button("Encoder")
                    examples_encoder = gr.Examples([["images/joconde.png"], ["images/braque.png"], ["images/Voronoy.jpg"]], inputs=[source_img])
            with gr.Column():
                encoded_img = gr.Image()
                with gr.Row():
                    bt = gr.Button("Decoder noisy image above")
                    
                

                
        with gr.Row():
            with gr.Column():
                
                noise_img = gr.Image(source="upload", type="filepath", label="init_img")
                k = gr.Slider(label='k', minimum = 1, maximum = 10, step = 1, value = 1)
                with gr.Row():
                    b2 = gr.Button("Decoder")
                    examples_decoder = gr.Examples([["images/afghan_coded.png"], ["images/wave_coded.png"], ["images/spider_coded.png"]], inputs=[noise_img])
                    #bt = gr.Button("Decoder noisy image above")
            with gr.Column():
                    decoded_img = gr.Image()

    with gr.Tab("webcam"):
        with gr.Row():
            with gr.Column():
                source_img_w = gr.Image(source="webcam", type="filepath", label="init_img")
                
                noise_w = gr.Slider(label='noise', minimum = 0.5, maximum = 1, step = .005, value = .95)
                with gr.Row():
                    b1_w = gr.Button("Encoder")
                    examples_encoder_w = gr.Examples([["images/joconde.png"], ["images/Voronoy.jpg"]], inputs=[source_img_w])
            with gr.Column():
                encoded_img_w = gr.Image()
                k_w = gr.Slider(label='k', minimum = 1, maximum = 10, step = 1, value = 1)
                with gr.Row():
                    bt_w = gr.Button("Decoder")   
                    examples_decoder_w = gr.Examples([["images/afghan_coded.png"], ["images/spider_coded.png"]], inputs=[encoded_img_w])
                    #bt = gr.Button("Decoder noisy image above")
            with gr.Column():
                    decoded_img_w = gr.Image()
            

    b1.click(encoder, inputs=[source_img, noise], outputs=encoded_img)
    b2.click(decoder, inputs=[noise_img,k], outputs=decoded_img)
    bt.click(decoder_noise, inputs=[encoded_img,k], outputs=decoded_img)
    b1_w.click(encoder, inputs=[source_img_w, noise_w], outputs=encoded_img_w)
    bt_w.click(decoder_noise, inputs=[encoded_img_w,k_w], outputs=decoded_img_w)
    
demo.launch()