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

title = "Playing with kNN for denoising"
description = "The encoder takes any image and remove a fraction of the pixels producing a noisy image. The decoder takes a noisy image and achieves denoising using kNN where you can choose k."

demo = gr.Blocks()

with demo:
    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")
        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")
                #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()