Spaces:
Sleeping
Sleeping
import importlib | |
import gradio as gr | |
from PIL import Image | |
import os | |
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) | |
# Sample files | |
sample_images = ["sample/test.jpg"] | |
sample_videos = ["sample/test2.mp4"] | |
def get_sample_image_paths(): | |
return [os.path.join("sample", f) for f in os.listdir("sample") if f.endswith(('.jpg', '.jpeg', '.png'))] | |
def get_sample_video_paths(): | |
return [os.path.join("sample", f) for f in os.listdir("sample") if f.endswith(('.mp4', '.avi'))] | |
sample_images = get_sample_image_paths() | |
sample_videos = get_sample_video_paths() | |
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", source="upload") | |
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) | |
gr.Markdown("### Or choose a sample image") | |
sample_image_dataset = gr.Dataset(components=[gr.Image(type="pil")], samples=[[Image.open(sample)] for sample in sample_images]) | |
sample_image_output = gr.Image(type="pil", label="Sample Detection Result") | |
sample_image_dataset.click(fn=detect_image, inputs=[model_choice, sample_image_dataset], outputs=sample_image_output) | |
with gr.Tab("Video Detection"): | |
video_input = gr.Video(label="Upload a video", source="upload") | |
video_output = gr.Video(label="Detection Result") | |
gr.Button("Submit Video").click(fn=detect_video, inputs=[model_choice, video_input], outputs=video_output) | |
gr.Markdown("### Or choose a sample video") | |
sample_video_dataset = gr.Dataset(components=[gr.Video()], samples=[[sample] for sample in sample_videos]) | |
sample_video_output = gr.Video(label="Sample Detection Result") | |
sample_video_dataset.click(fn=detect_video, inputs=[model_choice, sample_video_dataset], outputs=sample_video_output) | |
app.launch(share=True) | |