import gradio as gr
from faster_rcnn import img_detect, video_detection


# Define detection function
def detect(choice, inp):
    # Call the appropriate detection function based on the choice
    if choice == "Image":
        return img_detect(inp)
    elif choice == "Video":
        return video_detection(inp)
    else:
        return img_detect(inp)  # Default to image detection


# Create Gradio interface using blocks layout
with gr.Blocks() as demo:
    choices = gr.Dropdown(
        choices=["Image", "Video"],
        label="What type of object would you like to detect?",
    )
    inp_image = gr.Image(filepath="upload")
    output = gr.Image()
    det = gr.Button('Detect')
    det.click(detect, [choices, inp_image], output)

if __name__ == "__main__":
    demo.launch()