import gradio as gr
from backend import visualize_image

# gradio inputs
image_input = gr.components.Image(type="pil", label="Input Image")
color_mode_select = gr.components.Radio(choices=["Black/white", "Random", "Segmentation"], label="Color Mode", value="Segmentation")
mode_dropdown = gr.components.Dropdown(choices=["Trees", "Buildings", "Both"], label="Detection Mode", value="Both")

tree_threshold_slider = gr.components.Slider(minimum=0, maximum=1, step=0.1, value=0.7, label='Set confidence threshold "%" for trees')
building_threshold_slider = gr.components.Slider(minimum=0, maximum=1, step=0.1, value=0.7, label='Set confidence threshold "%" for buildings')

# gradio outputs
output_image = gr.components.Image(type="pil", label="Output Image")
title = "Aerial Image Segmentation"
description = "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"

# 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)