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 (iy you used the encoder)."""

demo = gr.Blocks()

with demo:
    gr.Markdown(title)
    gr.Markdown(description)
    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()
            

    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)
    
    
demo.launch()