brxerq's picture
Update app.py
c695cc2 verified
raw
history blame
No virus
1.34 kB
# 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)