import gradio as gr
from utils.segmentation import segment

# ------------------------- GRADIO -------------------------
if __name__ == "__main__":    
    with gr.Blocks() as demo:
        gr.Markdown("""
                    # CheXmask-U: Uncertainty in Landmark-based Anatomical Segmentation
                    
                    Demo of the **uncertainty estimation framework** proposed in the paper "CheXmask-U: Quantifying uncertainty in landmark-based anatomical segmentation for X-ray images".
                    The demonstration performs landmark-based segmentation (lungs and heart) and quantifies the uncertainty in the predicted position of each anatomical landmark.
                    
                    ### 📝 Instructions
                    
                    1.  **Upload** a chest X-ray image (PA or AP view) in PNG or JPEG format, or select an example image.
                    2.  **Explore Prediction Variability**: You can either use your mouse to draw on the input image and perform inpainting in different regions or adjust the Gaussian Noise Std Dev slider to simulate image corruption. 
                    3.  Click on **"Segment Image"**.
                        * The output image will display the segmentation overlay where the color gradient indicates the node-wise predictive uncertainty.
                        * The Results file output will contain the coordinates and per-node uncertainty estimates.
                    
                    Note: Pre-processing is not needed, it will be done automatically and removed after the segmentation.
                    """)

        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)