fire_detect_sam / app.py
gouravgujariya's picture
Create app.py
4b4fc4f verified
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()