Fireguard is a machine learning Artificial Intelligent trained with Yolov8 to detect Fire and Smokes from pictures.
For this project you have two options:
- Use pre-trained model in release, check Model
- Train model by yourself, check Train
If you want to use the model directly without training your own model, download the FireGuardBestModel.pt
or FireGuardLastModel.pt
from release
import cv2
from ultralytics import YOLO
import math
import os
# Cargar el modelo
model = YOLO('./FireGuardBestModel.pt')
# Leer las clases
classnames = ['fire', 'neutral', 'smoke']
# Crear un directorio para guardar las imágenes procesadas
os.makedirs('../tmp', exist_ok=True)
# Lista de imágenes a procesar
# Coger las imagenes de ./img sin . .. y otros que no sean .jpg
images = [img for img in os.listdir('../img') if img.endswith('.jpg')]
for image_path in images:
# Leer la imagen
frame = cv2.imread(f'../img/{image_path}')
frame = cv2.resize(frame, (640, 480))
# Procesar la imagen con YOLO
result = model(frame, stream=True)
# Get confidence and class names
# print(result[0])
# Obtener la información de bbox, confianza y nombres de clases para trabajar con ella
for info in result:
boxes = info.boxes
for box in boxes:
confidence = box.conf[0]
confidence = math.ceil(confidence * 100)
Class = int(box.cls[0])
if confidence > 10:
x1, y1, x2, y2 = box.xyxy[0]
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
print(f'Class: {classnames[Class]}, Confidence: {confidence}%')
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), 2)
# Write the class name and confidence
cv2.putText(frame, f'{classnames[Class]} {confidence}%', (x1 - 8, y1 - 8), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 1)
# Guardar la imagen procesada
processed_image_path = os.path.join('../tmp', os.path.basename(image_path))
cv2.imwrite(processed_image_path, frame)
If you have enough time or you want to include custom data to train your own AI model, you can choose from
- Train in Google colab (recommended)
- Train in your machine
You can find the datasets from release
To train the model in Google colab, be sure to have the datasets installed and uploaded to your Google drive with this name ~/D-Fire-datasets
. And check if you have this file branch
📁 D-Fire-datasets
|
\-----> 📁 fire
|
|----> 📁 test
| |
| |---> 📁 images (images)
| \---> 📁 labels (txt)
|
|----> 📁 train
| |
| |---> 📁 images (images)
| |---> 📁 labels (txt)
| \---> 📄 labels.cache
|
|----> 📁 train
| |
| |---> 📁 images (images)
| |---> 📁 labels (txt)
| \---> 📄 labels.cache
|
\---> 📄 data.yaml
Donwload the source file from ./src/FireGuardML_Train.ipynb
and import it to google colab. Be sure to enable GPU
Go to src/
and install dependecies
python3 -m pip install -r requirements.txt
Down load the dataset from release and decompress it inside ./data
folder
To train the model, execute the file ./src/setup.sh
inside ./src/context
cd src
./setup.sh
FireGuardML by Zheng Lin Lei is licensed under CC BY-NC 4.0