Detecting Wildfires with YOLOv3
Today, I am going to walk you through training a custom dataset with YOLOv3. For the purposes of this blog- I will assume you already have compiled darknet on your local machine, downloaded the pre-trained weights and have OpenCV installed. If you do not, here is the link to the darknet repository. Alexey provides detailed instructions on how to properly compile darknet for different operating systems.
Next, we will need images with annotations for the bounding boxes. YOLO requires each image to have a matching .txt file in the format of <class-name>, <x-center>, <y-center>, <width>, and <height>. I recommend using CVAT or LabelImg to label your images and get the bounding box annotations.
I used https://roboflow.com/ to find the dataset we will be working with today. This consists of images of smoke caused by wildfires and a matching annotation.txt file split into training, validation, and test sets. Therefore, we will have 1 class that we are trying to detect. Next, make a directory in the darknet/data and name it wildfire. Last, move the images into this folder.
Before we can train YOLOv3, we will need to create a few files in our darknet repository. In order to do this open up your preferred text editor. I will be using Visual Studio Code. Let’s create a .txt file storing our class names and call it wildfire.names. Make sure to save this file inside your darknet/data folder.
Next, we will need to make our train.txt and valid.txt files consisting of the path to our images. In the code below we are simply going through our directory of training images and writing the path into a new .txt file that will be saved as train.txt. For validation.txt, all you need to do is replace train with validation.
import osimage_files = []
os.chdir(os.path.join("data", "wildfire/train"))for filename in os.listdir(os.getcwd()): if filename.endswith(".jpg"): image_files.append("data/wildfire/train" + filename)os.chdir("..")with open("train.txt", "w") as outfile: for image in image_files: outfile.write(image) outfile.write("\n") outfile.close()
os.chdir("..")
Alright, now we have our train.txt, valid.txt, and wildfire.names files ready. These three files will be referenced in the wildfire.data file we are about to create. The .data file will contain the number of classes, path to train.txt, path to valid.txt, the path to our wildfire.names and a folder to store weights as we train. This allows for you to stop training, test the model’s performance, and then continue training from the last weights that were saved- instead of having to start the training process over again. Here is what my wildfire.data file looks like.
Next, we will need to make a copy of the yolov3.cfg and save it as yolov3-custom.cfg. Once again, this can be found in your darknet/cfg folder. This file contains the entire model structure for yolo and we will need to make a few changes to make sure it is set up properly for our custom classes. The first change will be to max_batches. The formula for determining max batches is (# of classes * 2000). However, the minimum amount accepted will be 4000 so even if we only have 1 class we need to set our max_batches = 4000. Next, we need to change our steps. This will be 80%, 90% of the max_batches. In this instance, steps will be 3200,3600. See the code snippet below.
Once this is done, use control + f to search for yolo. It will find 3 instances and we will need to make the same changes for all 3. Locate classes and adjust it to match how many classes you are using. For this example, we are using 1 class so change 80 to 1. Next, we need to change the number of filters located directly above yolo in the convolutional layer. The formula for determining # of filters is (# of classes + 5) *3. In our example, (1 class + 5)*3 = 18 filters. Repeat both of these steps for all three yolo instances in our cfg file.
Alright, now the fun begins! Open up your powershell or terminal and move into the darknet directory. From here we will be able to train our detector. This command is running the darknet.exe detector train with paths to where you saved wildfire.data, the yolov3-custom.cfg, and the pre-trained weights.
If everything was installed properly this command should initiate training and a window should pop up like the one below. This chart represents avg loss over each iteration. Ideally, loss will be less than 1.
Once training is completed, we can use darknet to make detections on test images directly or we can use OpenCV. For today’s example I am just going to utilize the built in darknet detector to test the model on a few images. To do so we need to go back to our powershell or terminal. Here we are loading in our saved weights from the backup folder created earlier and giving it the path to an image we would like to test.
Boom, the model performs as expected. As you can see above, this returns an image with the bounding box plotted and the confidence level of the prediction. I hope that you found this tutorial useful and feel free to send me any object detection projects you are working on.
