Update app.py

app.py

@@ -1,6 +1,7 @@
-# main.py
 import importlib
 import gradio as gr
+from PIL import Image
+import os
 
 def load_model(model_name):
     module = importlib.import_module(model_name)
@@ -20,19 +21,45 @@ 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")
+            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")
+            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)