import gradio as gr
from utils.segmentation import segment

# ------------------------- GRADIO -------------------------
if __name__ == "__main__":    
    with gr.Blocks() as demo:
        gr.Markdown("# Chest X-ray HybridGNet Segmentation Demo")

        with gr.Tab("Segment Image"):
            with gr.Row():
                with gr.Column(scale=1):
                    image_input = gr.Image(
                        type="numpy", 
                        tool="sketch", 
                        image_mode="L", 
                        height=450, 
                    )
                    
                    noise_slider = gr.Slider(
                        label="Gaussian Noise Std Dev",
                        minimum=0.0,
                        maximum=0.25,
                        step=0.01,
                        value=0.0
                    )
                    
                    with gr.Row():
                        clear_button = gr.Button("Clear")
                        image_button = gr.Button("Segment Image")

                    gr.Examples(inputs=image_input, examples=[
                        'utils/example1.jpg','utils/example2.jpg',
                        'utils/example3.png','utils/example4.jpg'
                    ])
                    
                with gr.Column(scale=2):
                    image_output = gr.Image(type="filepath", height=450)
                    results = gr.File()       
                    
        gr.Markdown("""
                    Example images extracted from Wikipedia, released under:
                    1. CC0 Universial Public Domain. Source: https://commons.wikimedia.org/wiki/File:Normal_posteroanterior_(PA)_chest_radiograph_(X-ray).jpg
                    2. Creative Commons Attribution-Share Alike 4.0 International. Source: https://commons.wikimedia.org/wiki/File:Chest_X-ray.jpg
                    3. Creative Commons Attribution 3.0 Unported. Source https://commons.wikimedia.org/wiki/File:Implantable_cardioverter_defibrillator_chest_X-ray.jpg
                    4. Creative Commons Attribution-Share Alike 3.0 Unported. Source: https://commons.wikimedia.org/wiki/File:Medical_X-Ray_imaging_PRD06_nevit.jpg
                    """)
                    
        clear_button.click(lambda: None, None, image_input, queue=False)
        clear_button.click(lambda: None, None, image_output, queue=False)
        image_button.click(
            segment,
            inputs=[image_input, noise_slider],
            outputs=[image_output, results],
            queue=False
        )
    demo.launch(share=True)