CODEWITHSUNDEEP
deep-learningcomputer-visionobject-detectionyolomulticlass-classification
user8128965
user8128965

Reputation: 41

Training YOLO (object detection) on multiple objects on imbalanced dataset?

I'm training YOLO on custom dataset(using Alexey AB implementation of Darknet). It has 3 classes of images where class 1 has 45k images, and remaining two have around 1k images.

After training it for 6k iterations, the loss is in between 1.5 and 2. However, when i tried running it on a video, it is only detecting class 1.

I would like to know what is the reason for this, is it because of the imbalance in the no of images in the dataset? Is there a way to solve this problem?

Upvotes: 2

Views: 5824

Answers (1)

rajput
rajput

Reputation: 176

Yes, to begin with you've an unbalanced dataset. The recommended number of images per class is >2000 (according to the directions on the repository).

I'd recommend you to have a test set and track your mean average precision along with loss while you're training. If the average precision is low, you can assume that not all classes are being detected equally well.

There are a few ways to solve this but I'm assuming you already might have figured out.

  • Collect more data for other classes.
  • Try data augmentation for other classes with rotation, noise, etc.
  • You could try adding synthetic images for the other classes.
  • Even though it's a bad idea to leave out data for the 1st class, at least try reducing it so there's a balanced dataset.

I'm calling other classes for the classes with less data.

Upvotes: 3

Related Questions