# main.py
import importlib
import gradio as gr

def load_model(model_name):
    module = importlib.import_module(model_name)
    return module

models = {
    "Multi-class model": "model_1",
    "Empty class": "model_2",
    "Misalignment class": "model_3"
}

def detect_image(model_choice, input_image):
    model = load_model(models[model_choice])
    return model.detect_image(input_image)

def detect_video(model_choice, input_video):
    model = load_model(models[model_choice])
    return model.detect_video(input_video)

app = gr.Blocks()

with app:
    gr.Markdown("## Object Detection using TensorFlow Lite Models")
    with gr.Row():
        model_choice = gr.Dropdown(label="Select Model", choices=list(models.keys()))
        with gr.Tab("Image Detection"):
            image_input = gr.Image(type="pil", label="Upload an image")
            image_output = gr.Image(type="pil", label="Detection Result")
            gr.Button("Submit Image").click(fn=detect_image, inputs=[model_choice, image_input], outputs=image_output)
        with gr.Tab("Video Detection"):
            video_input = gr.Video(label="Upload a video")
            video_output = gr.Video(label="Detection Result")
            gr.Button("Submit Video").click(fn=detect_video, inputs=[model_choice, video_input], outputs=video_output)

app.launch(share=True)