import os 
os.system("python -m pip install --upgrade pip")
os.system("pip uninstall -y gradio")
os.system("pip install gradio==4.1.2")

import gradio as gr
from backend import *

with gr.Blocks() as demo:
    gr.Markdown(
    """
    #                 Aerial Image Segmentation
    An instance segmentation demo for identifying boundaries of buildings and trees in aerial images using DETR (End-to-End Object Detection) model with MaskRCNN-101 backbone
    """
    )
    with gr.Row(equal_height=True):
        with gr.Column():
            image_input = gr.components.Image(type="pil", label="Input Image")
        with gr.Column():
            mode_dropdown = gr.Dropdown(choices=["Trees", "Buildings", "Both"], label="Detection Mode", value="Both")
            color_mode_select = gr.components.Radio(choices=["Black/white", "Random", "Segmentation"], label="Color Mode", value="Segmentation")

    # split tree and building into two rows side by side
    tree_row, building_row = gr.Row(), gr.Row()
    # tree_col, building_col = gr.Column(elem_id="tree_col"), gr.Column(elem_id="building_col")
    with tree_row as tree_options:
        tree_version_dropdown = gr.Dropdown(choices=["treev1", "treev2"], label="Tree Detection Version", value="treev2", visible=True, interactive=True)
        tree_pth_dropdown = gr.Dropdown(choices=list_pth_files_in_directory("tree_model_weights", "v2"), label="Select a tree model file", visible=True, interactive=True)
        tree_threshold_slider = gr.Slider(minimum=0, maximum=1, step=0.1, value=0.7, label='Set confidence threshold "%" for trees', visible=True, interactive=True)

    with building_row as building_options:
        building_version_dropdown = gr.Dropdown(choices=["buildingv1", "buildingv2"], label="Building Detection Version", value="buildingv1", visible=True, interactive=True)
        building_pth_dropdown = gr.Dropdown(choices=list_pth_files_in_directory("building_model_weight", "v1"), label="Select a building model file", visible=True, interactive=True)
        building_threshold_slider = gr.Slider(minimum=0, maximum=1, step=0.1, value=0.7, label='Set confidence threshold "%" for buildings', visible=True, interactive=True)

    # mode_dropdown.change(update_visibility, inputs=[mode_dropdown], outputs=[tree_version_dropdown, tree_pth_dropdown, tree_threshold_slider, building_version_dropdown, building_pth_dropdown, building_threshold_slider])
    mode_dropdown.change(update_row_visibility, inputs=[mode_dropdown], outputs=[tree_row, building_row])
    tree_version_dropdown.change(update_path_options, inputs=[tree_version_dropdown], outputs=[tree_pth_dropdown])
    building_version_dropdown.change(update_path_options, inputs=[building_version_dropdown], outputs=[building_pth_dropdown])
    
    output_image = gr.components.Image(type="pil", label="Output Image")
    run_model = gr.Button("Upload Image and Run Model")
    
    run_model.click(visualize_image, inputs=[image_input, mode_dropdown, tree_threshold_slider, building_threshold_slider, color_mode_select, tree_version_dropdown, tree_pth_dropdown, building_version_dropdown, building_pth_dropdown], outputs=[output_image])
    demo.launch()