Spaces:
Sleeping
Sleeping
File size: 2,490 Bytes
7ad337f eadd655 344f265 82a7618 eadd655 7ad337f 82a7618 afa3d9d 82a7618 7ad337f 0b39968 7ad337f 0b39968 7ad337f afa3d9d 7ad337f | 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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | import gradio as gr
import os
import numpy as np
import cv2
import traceback
import logging
from PIL import Image
from helmet_detect_alert import detect_and_alert, MODEL_PATHS
from ultralytics import YOLO
# Load YOLOv8 model
model_path = MODEL_PATHS["YOLOv8n"]
model = YOLO(model_path)
# ---------- Detection Functions ----------
logging.basicConfig(filename="error_log.txt", level=logging.ERROR)
def detect_image_fn(image, confidence):
try:
if image is None:
raise ValueError("No image uploaded.")
# Convert PIL image to NumPy
image_np = np.array(image.convert("RGB"))
# YOLO prediction
results = model.predict(image_np, conf=confidence, verbose=False)
if not results:
raise RuntimeError("Model returned empty results.")
# Annotate result
annotated = results[0].plot()
annotated_rgb = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB)
return Image.fromarray(annotated_rgb)
except Exception as e:
# Log error to file
logging.error("Image detection failed:\n%s", traceback.format_exc())
print("Error during prediction:", e)
return None # Gradio will show "Error"
def detect_video_fn(video_file, confidence):
if video_file is None:
return None
input_path = video_file
output_path = os.path.join("output", os.path.basename(input_path).replace(".", "_pred.")) + "mp4"
os.makedirs("output", exist_ok=True)
detect_and_alert(input_path, output_path, model, confidence)
return output_path
# ---------- Interfaces ----------
image_interface = gr.Interface(
fn=detect_image_fn,
inputs=[
gr.Image(label="Upload Image"),
gr.Slider(0.1, 1.0, value=0.3, step=0.05, label="Confidence")
],
outputs=gr.Image(label="Predicted Output"),
title="🖼 Helmet Detection from Image",
description="Upload an image to detect heads not wearing helmets."
)
video_interface = gr.Interface(
fn=detect_video_fn,
inputs=[
gr.Video(label="Upload Video"),
gr.Slider(0.1, 1.0, value=0.3, step=0.05, label="Confidence")
],
outputs=gr.Video(label="Predicted Output"),
title="🎥 Helmet Detection from Video",
description="Upload a video and get helmet detection alerts visually."
)
# ---------- Launch Tabbed App ----------
gr.TabbedInterface(
[image_interface, video_interface],
tab_names=["Image Detection", "Video Detection"]
).launch()
|