Spaces:
Sleeping
Sleeping
# 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) | |