Stephen Kamau
9 min readJan 14, 2023

Real-Time Object Detection with YOLO: Training YoloV5 On Fire Detection Images

Welcome back to part II of our blog on Real-Time Object Detection with YOLO: A Step-by-Step Guide to Realtime Fire Detection. In the first part of this series, we covered the basic concepts of YOLO and its data format. We also discussed how it can be used for real-time object detection and fire detection specifically.

Part I can be found at the following link.

In this second part II, we will dive deeper and build the model for our real-time fire detection application. We will start by discussing the prerequisites and the necessary libraries to set up our environment, preparing our dataset, and finally training the model. How the model is to be evaluated and its usage will also be done.

For the whole blog, code snippets are going to be provided and examples to help you follow along and build your own real-time fire detection model with YOLO. By the end of this blog, you will have a solid understanding of how YOLO works and how to apply it to real-world applications.

Training and Validation

When it comes to training a machine learning model, it is important to understand the difference between the training set and the validation set. The training set is the data that we use to train our model, while the validation set is used to evaluate the performance of our model and make any necessary adjustments. For this project, we used the KFold method where the data is divided into k number of folds, and we select one of the folds to be used as the validation set. The remaining k-1 folds are used for training.
In our case, we are going to divide our data into 10 folds. This means that we will have 10 different combinations of training and validation sets. We can then select one of the folds to be used as the validation set, while the other 9 folds will be used for training. This method helps ensure that our model is trained and tested on different subsets of data, which can give us a better idea of how well it will perform on new, unseen data.

# split into training and validation
gkf = GroupKFold(n_splits = 10)
processed_df['fold'] = -1
for fold, (train_idx, val_idx) in enumerate(gkf.split(processed_df, groups = processed_df.image_id.tolist())):
processed_df.loc[val_idx, 'fold'] = fold

# get all image ids for both training and validation with the bbox
# validation set is the one with fold labelled 10..
val_fold_files_ids = list(processed_df[processed_df.fold==fold].image_id.unique())
train_fold_files_ids = list(processed_df[processed_df.fold!=fold].image_id.unique())

As we had discussed how the data structure should be arranged, the below code will be used to create folders for both training and validation data.

# create thr directorues as follows
# base_dir
# images
# train (contains image files)
# validation (contains image files)
# labels
# train (contains .txt files)
# validation (contains .txt files)


os.makedirs(f'{DATA_DIR}/yolo/labels/train', exist_ok = True)
os.makedirs(f'{DATA_DIR}/yolo/labels/val', exist_ok = True)
os.makedirs(f'{DATA_DIR}/yolo/images/train', exist_ok = True)
os.makedirs(f'{DATA_DIR}/yolo/images/val', exist_ok = True)

Processing Data

For each image, a label txt file is created as was discussed earlier(class x-center y-center width height). Images will be moved from the folder where they are to either the Training or validation folder while creating labels for the same.

"""
In this cell, It is a function that iterates through each row in the df
It first determine whether the image belows to training or testing sets.
For each image:
1- get the info for each bounding box
2- write the bounding box info to a txt file
3- save the txt file in the correct folder
4- copy the image to the correct folder
Before saving the image, it must be processed by applying clahe function/
class, center (x,y), width, height
Also the format for YOLO dataset is followed i.e
"""
def process_data_for_yolo(df, file_id_lists, data_type='train'):
"""
Reads details for an image and transfers them to their respective folders when proceesed.
"""
#iterate through each row
#for _, row in notebook.tqdm(df.iterrows(), total=len(df)):
for each_id in tqdm(file_id_lists, total=len(file_id_lists)):
#print(each_id, file_id_lists)
row = df.loc[each_id]

#get img information
file_id = row['image_id']
file_name = row['img_name']
# Convert into the Yolo input format
yolo_data = []
for xq in row['information']:
curr_bbox_infor = [0, xq['x_mid'], xq['y_mid'], xq['w'] , xq['h']]
if xq['x_mid']<=0 and xq['y_mid']<=0 and xq['w']<=0 and xq['h']<=0:
pass
else:
yolo_data.append(curr_bbox_infor)

# convert to nump array
yolo_data = np.array(yolo_data)
#print(yolo_data)
#copy image to another directory where training will occur
full_file_path = f"{IMG_DIRS}/{file_name}"
shutil.copy(full_file_path, f'{YOLO_IMGS_DIR}{data_type}')

#saved format must be class, center (x,y), width, heihgt
np.savetxt(os.path.join(f"{YOLO_LBL_DIR}/{data_type}/{file_id}.txt"),
yolo_data,
fmt=["%d", "%f", "%f", "%f", "%f"]
)
#call the above function with respective data
# since we using Id to create the data, lets set index as file id.
indexed_train = processed_df.copy()
indexed_train.index = processed_df['image_id']
#training
process_data_for_yolo(indexed_train, train_fold_files_ids, 'train')
#validation
process_data_for_yolo(indexed_train, val_fold_files_ids, 'val')

Training Configuration

The YOLO configuration files are used to specify the settings and parameters for training and testing a YOLO object detection model. These files typically have a .cfg or .yaml extension and include information such as the path to the pre-trained weights and the path to the dataset for training, the number of classes, the architecture of the model, the type of layers and activation functions to use, and the batch size and learning rate for training.

In YOLO, there are three types of config files in YAML format that are downloaded when cloning the Ultralyric repo. We will only customize one of them, which contains information about our training dataset. The three types of config files are:

  • Data Config File: Describes the dataset to be used. It contains information such as the paths to the train, validation, and test (optional) datasets, the number of classes, and the names of these classes in the same order as their index.
  • Model Config File: Shows the model Architecture. Ultralytics supports several YOLOv5 architectures, named P5 models, which vary mainly by their parameter size. They include YOLOv5n (nano), YOLOv5s (small), YOLOv5m (medium), YOLOv5l (large), and YOLOv5x (extra large).
  • The Hyperparameters Configs: Defines the hyperparameters for the training, including the learning rate, momentum, losses, augmentations, etc. Ultralytics provides a default hyperparameters file under the data/hyp/hyp.scratch.yaml directory. The YAML config files are nested in the directory defining 3 default parameters for large, medium, and low. We will use the default hyperparameters and don’t need to edit them.

In this case, Data Config File will be created which contains paths and class information. You can point to the base directory with training and validation images or point to a txt file containing the paths for these images. First, we create two txt files that contain all paths for the images and labels. These two files will be used in the data config files.

# save the data (labels and images directories into a txt file)
with open(os.path.join( "/content/drive/MyDrive/FIRE/yolo/" , 'train.txt'), 'w') as f:
for path in glob('/content/drive/MyDrive/FIRE/yolo/images/train/*'):
f.write(path+'\n')

with open(os.path.join( "/content/drive/MyDrive/FIRE/yolo/" , 'val.txt'), 'w') as f:
for path in glob("/content/drive/MyDrive/FIRE/yolo/images/val/*"):
f.write(path+'\n')

#We then create a data config file and save it into yaml file.
# information to be used as data by yolo
yolo_data = dict(
train = os.path.join( "/content/drive/MyDrive/FIRE/yolo/" , 'train.txt') ,
val = os.path.join( "/content/drive/MyDrive/FIRE/yolo/", 'val.txt' ),
nc = 1,# there is only one class
names = ["Fire"]
)
# save the file as yaml with the information
import yaml
with open(os.path.join( "/content/drive/MyDrive/FIRE/yolo/" , 'fire.yaml'), 'w') as outfile:
yaml.dump(yolo_data, outfile, default_flow_style=False)

We then clone the repo containing yolo weights and codes. This is the part where the model will be trained. Run the following command on your notebook to clone the repo and then install all the requirements libs.

#run the following commands one by one.
git clone https://github.com/ultralytics/yolov5.git
# move to the directory with downloaded code
cd /content/drive/MyDrive/FIRE/yolo/yolov5
# install the requirements.txt
pip install -qr requirements.txt

Training and validation Process

After installation of configuration files, the next step is just running the command for training and validation with your data source config file. In the cloned file, we have multiple files e.g the train.py file is used to train a YOLO model and accepts several arguments, including the data and model weights to be used. Some of the parameters passed and their uses are:

  • — weights: the pre-trained model that we are using. A list of available pre-trained models can be found on the Ultralytics GitHub page.
  • — save-txt: saves the predicted bounding box coordinates to a txt file, one txt file per image.
  • — img: resizes the image to this size before creating the mosaic.
  • — name: gives the model a name.
  • — batch: sets the batch size.
  • — epochs: sets the number of training epochs.
  • — data: sets the YAML file path.

You can review the source code in detect.py, train.py, and val.py to investigate how they are used and many others. To train, we need to run the following command where ./FIRE/yolo/fire.yaml is the path to our data config file and yolov5x.pt is the weight we are using.

python train.py - exist-ok - img 640 - batch 16 - epochs 30 - data /content/drive/MyDrive/FIRE/yolo/fire.yaml - weights yolov5x.pt

The above command will train the model and save the model results into a folder called runs/train/exp. The final perfomance of the model was as follows.

Models Perfomance

You can view all of them and visualize what you want, example to view a graph of the f1 score, we can use the following command.

plt.imshow(plt.imread(f'runs/train/exp/F1_curve.png'))

Here is a figure showing the results of f1 curve from the model.

F1 score curve for the model

To evaluate the model on a test data maybe you have created, you will need to run this command which is almost similar to training one but with a different python script. In this case, runs/train/exp/weights/best.pt is the path to the weights of our trained file.

python val.py - exist-ok - weights runs/train/exp/weights/best.pt - data ./FIRE/yolo/fire.yaml - img 640 - iou 0.5

After running the above script, the results will be saved at runs/val/exp which contains graphs like f1-score, confusion matrix, etc. The model performance finally was Precision 83.3%, Recall 76.5% and mAP50–95 54.9% which is summarized below.

Model summary: 322 layers, 86173414 parameters, 0 gradients, 203.8 GFLOPs
Class Images Instances P R mAP50 mAP50-95: 100%
all 215 429 0.833 0.765 0.825 0.549
Speed: 0.5ms pre-process, 36.9ms inference, 1.4ms NMS per image at shape (32, 3, 640, 640)

An example of results after running the trained weights on validation data is as in the figure below.

Images showing model validation

To make predictions using YOLOv5, you can download the weights you are comfortable with and then use torch.load to load the model. Alternatively, you can use the detect.py file and pass in the path of the image you want to predict, along with the path of the weights and other desired parameters. An example command would be:

python detect.py - source <IMG PATH> - weights <WEIGHTS PATH> - img 640

An example of a video analyzed using the above-trained model can be seen below

In conclusion, Real-Time Object Detection with YOLO is a powerful tool that can be used to detect and classify objects in real time. In this blog, we have discussed the process of training a YOLO model on fire detection images, including the steps of processing the data, setting up the training environment, and fine-tuning the model. We have also discussed the potential applications of the trained model, such as real-time fire detection in buildings, forests, and other areas.

We hope that this blog has provided you with a clear understanding of the process of training a YOLO model on fire detection images and its potential applications. With the help of this technique, we can detect and prevent fire at an early stage which is crucial for the safety and well-being of our surroundings.