Update app.py

app.py

@@ -4,18 +4,16 @@ import cv2
 from PIL import Image
 import numpy as np
 
-# Load YOLO model (upload Best_Model1.pt to your Hugging Face Space)
+# Load YOLOv8 model
 MODEL_PATH = "Best_Model1.pt"
 model = YOLO(MODEL_PATH)
 
-def detect_flash_gradio(image: Image.Image):
+def detect_flash(image: Image.Image):
     # Convert PIL image to OpenCV format
     image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
 
     # Run detection
     results = model(image_cv)
-
-    # Extract detection info
     detections = results[0].boxes.data.tolist()
     num_detections = len(detections)
 
@@ -35,23 +33,40 @@ def detect_flash_gradio(image: Image.Image):
 
     return output_image, message
 
+def extract_first_frame(video):
+    cap = cv2.VideoCapture(video)
+    success, frame = cap.read()
+    cap.release()
+    if success:
+        # Convert BGR to RGB
+        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        return Image.fromarray(frame_rgb)
+    return None
+
+def process_webcam_video(video):
+    image = extract_first_frame(video)
+    if image:
+        return detect_flash(image)
+    else:
+        return None, "Could not read frame from video."
+
 # Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("Flash Detection with YOLOv8")
-    gr.Markdown("Upload an image or use your webcam to detect camera flash.")
+    gr.Markdown("Upload an image or record from your webcam to detect flash.")
 
     with gr.Tab("Upload Image"):
-        upload_input = gr.Image(type="pil", label="Upload Image")
-        upload_btn = gr.Button("Detect Flash")
-        upload_output_img = gr.Image(type="pil", label="Detected Image")
-        upload_output_text = gr.Textbox(label="Detection Info")
-        upload_btn.click(fn=detect_flash_gradio, inputs=upload_input, outputs=[upload_output_img, upload_output_text])
-
-    with gr.Tab("Use Webcam"):
-        webcam_input = gr.Image(type="pil", tool="webcam", label="Capture from Webcam")
+        img_input = gr.Image(type="pil", label="Upload Image")
+        detect_btn = gr.Button("Detect Flash")
+        output_img = gr.Image(type="pil", label="Detected Image")
+        output_text = gr.Textbox(label="Detection Info")
+        detect_btn.click(fn=detect_flash, inputs=img_input, outputs=[output_img, output_text])
+
+    with gr.Tab("Use Webcam (Record Video)"):
+        webcam_video = gr.Video(label="Record from Webcam")
         webcam_btn = gr.Button("Detect Flash")
-        webcam_output_img = gr.Image(type="pil", label="Detected Image")
-        webcam_output_text = gr.Textbox(label="Detection Info")
-        webcam_btn.click(fn=detect_flash_gradio, inputs=webcam_input, outputs=[webcam_output_img, webcam_output_text])
+        webcam_img = gr.Image(type="pil", label="Detected Image")
+        webcam_text = gr.Textbox(label="Detection Info")
+        webcam_btn.click(fn=process_webcam_video, inputs=webcam_video, outputs=[webcam_img, webcam_text])
 
 demo.launch()