Spaces:
Sleeping
Sleeping
File size: 2,677 Bytes
7c3c7c7 f03a267 7c3c7c7 f03a267 7c3c7c7 f03a267 7c3c7c7 f03a267 7c3c7c7 f03a267 7c3c7c7 f03a267 7c3c7c7 f03a267 7c3c7c7 f03a267 7c3c7c7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
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)
|