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)