Spaces:
Build error
Build error
import cv2 | |
import gradio as gr | |
from ultralytics import YOLO | |
# Load the YOLO model (update the path to your fire detection model weights) | |
model = YOLO('path/to/your/best.pt') # Replace 'path/to/your/best.pt' with the actual path to your model file | |
def detect_fire(frame): | |
# Convert the frame to RGB format (YOLO expects this format) | |
img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) | |
# Perform fire detection using the YOLO model | |
results = model(img) | |
# Draw bounding boxes and labels on the detected fire areas | |
for bbox in results[0].boxes: | |
xyxy = bbox.xyxy[0] # Bounding box coordinates | |
conf = bbox.conf[0] # Confidence score | |
cls = int(bbox.cls[0]) # Class ID | |
label = model.names[cls] # Class name | |
if label == "fire": # Make sure this matches the label in your trained model | |
# Draw a rectangle around the detected fire | |
cv2.rectangle(img, (int(xyxy[0]), int(xyxy[1])), (int(xyxy[2]), int(xyxy[3])), (255, 0, 0), 2) | |
# Put the label text above the rectangle | |
cv2.putText(img, f"{label} {conf:.2f}", (int(xyxy[0]), int(xyxy[1]) - 10), | |
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2) | |
# Convert the image back to BGR format for display | |
return cv2.cvtColor(img, cv2.COLOR_RGB2BGR) | |
# Create a Gradio interface for fire detection | |
iface = gr.Interface( | |
fn=detect_fire, | |
inputs=gr.Image(source="webcam", tool="editor", streaming=True), # Use webcam as the input source | |
outputs="image", | |
title="Fire Detection using YOLO", | |
description="This application detects fire in real-time using a YOLO model." | |
) | |
# Launch the Gradio app | |
iface.launch() |