import gradio as gr
import numpy as np
import color_transfer_MKL as ct


def color_transfer(source, target):
    source = np.array(source)/255
    target = np.array(target)/255
    result = ct.color_transfer_MKL(source, target)
    return np.uint8(result * 255)

with gr.Blocks() as interface:
    with gr.Row(equal_height=True):
        source_image = gr.Image(type="numpy", label="Source Image", sources=['upload'], interactive=True)
        target_image = gr.Image(type="numpy", label="Target Image", sources=['upload'], interactive=True) 
        result_image = gr.Image(type="numpy", label="Result", show_download_button=True, interactive=False, show_share_button=False)
        
    with gr.Row():
        submit_button = gr.Button()
        submit_button.click(fn=color_transfer, inputs=[source_image, target_image], outputs=[result_image])
        clear_button = gr.ClearButton(components=[source_image, target_image, result_image])
    with gr.Row():
        examples_dropdown = gr.Examples([["samples/source1.jpg", "samples/target1.jpg"], ["samples/source2.jpeg", "samples/target2.jpg"]],
                                        [source_image, target_image])
        

if __name__ == "__main__":
    interface.launch()