YOLOv1 doesn't work for custom dataset

Question

I am currently trying to train my own YOLOv1 network, based on this repository: https://github.com/ivanwhaf/yolov1-pytorch

The images I want to train look like this:

Three classes I want to detect:

• Class 0: A meter (the "clocks" at the bottom)
• Class 1: The number display
• Class 2: Each individual digit

I generated 160 training images and automatically generated the YOLO bounding boxes, for the image above they look like this:

0.21875 0.5833333333333334 0.25 0.3333333333333333 0
0.5 0.5833333333333334 0.25 0.3333333333333333 0
0.78125 0.5833333333333334 0.25 0.3333333333333333 0
0.5 0.14583333333333334 0.9375 0.20833333333333334 1
0.890625 0.14583333333333334 0.09375 0.15833333333333333 2
0.765625 0.14583333333333334 0.09375 0.15833333333333333 2
0.640625 0.14583333333333334 0.09375 0.15833333333333333 2
0.515625 0.14583333333333334 0.09375 0.15833333333333333 2
0.390625 0.14583333333333334 0.09375 0.15833333333333333 2
0.265625 0.14583333333333334 0.09375 0.15833333333333333 2

After 40 epochs or so I end up with a Loss of around 1.8 (the network started at 200) but the bounding boxes are no where near where they should be:

Red is the prediction and green are the ground truth bounding boxes.

Do you have an idea why that is? Did I do something wrong with labelling the bounding boxes of the training images?

Answer 1

I will try to answer you based on my previous experience with training YOLO architectures. Hopefully, this can help you build your object detection model.

The first thing to note is your training data. The fact that the loss has decreased on the model is a very good sign that it is improving its accuracy. However, if the accuracy drops when using the test dataset, it means that the model did not generalize well for the task given. This is usually a case of overfitting.

The size of your dataset is an important consideration. In general, 160 data points are too few for training an object detection module. Therefore, the first thing I would recommend is to increase the size of your dataset. If increasing the dataset is not possible, you may want to consider other types of models. You could try using a lighter version of YOLO, such as YOLO-tiny, which is a less computationally expensive model with fewer weights to train. This means that it will learn to do the task faster.

Another thing you can consider is using techniques like transfer learning or employing pre-trained models. Usually, YOLO's open-source repositories have pre-trained weights that are trained with ImageNet or the COCO dataset. I did not see any in your shared repository, but it is highly recommended to use them as a general rule.

Finally, since you are training an object detection module, you should consider the type of background in which you are searching for the image. It is not the same to train object detection for a camera with real-world imagery as it is to train an object detection module for a user interface. In any case, it is a good strategy to augment the data to provide more diversity to the neural network. For this, you can try tools like Albumentations. This tool will help augment the number of data points in a synthetic way, for example, through rotations, scaling, and changing colors. A trick that I have used is to embed the object I want to identify in different types of backgrounds, and in my use case, it has helped increase accuracy. I hope this helps!