import gradio as gr
from backend import visualize_image

# gradio inputs
image_input = gr.inputs.Image(type="pil", label="Input Image")
mode_dropdown = gr.inputs.Dropdown(["Trees", "Buildings", "Both"])
tree_threshold_slider = gr.inputs.Slider(0, 1, 0.1, 0.7, label='Set confidence threshold % for trees')
building_threshold_slider = gr.inputs.Slider(0, 1, 0.1, 0.7, label='Set confidence threshold % for buildings')
color_mode_select = gr.inputs.Radio(["Black/white", "Random", "Segmentation"])

# gradio outputs
output_image = gr.outputs.Image(type="pil", label="Output Image")
title = "Building Segmentation"
description = "An instance segmentation demo for identifying boundaries of buildings in aerial images using DETR (End-to-End Object Detection) model with MaskRCNN-101 backbone"

# gradio interface
interface = gr.Interface(
    fn=visualize_image,
    inputs=[image_input, mode_dropdown, tree_threshold_slider, building_threshold_slider, color_mode_select],
    outputs=output_image,
    title=title,
    description=description
)

interface.launch(debug=True)