Classification using a YOLOv5
1. Get your data ready
To train new YOLO models, you need the file
src/dl/yolov5_training.py
. It contains the entire workflow to produce a bundled model ready for deployment.Before starting, create a folder named
models
to store all the new model versions you create.You also need a
working_dir
where the script will export its temporary data.- To train the YOLO model, you need two distinct folders. You can name them as you like.
The first folder, referred to as
images
, will contain.png
images with values globally normalized in the range [0, 255].The second folder, referred to as
labels
, will also contain.txt
files in which each line contains: C, X1, Y1, X2, Y2.The names in both folders should be the same with only the exntension changing.
- The models produced by this script include:
version.txt
: The version index of this model, allowing detection if the model should be re-downloaded from the internet.val_batch0_pred.jpg
: An overview of what the model predicted on an image from the validation set.val_batch0_labels.jpg
: The ground-truth of the image above.train_batchN.jpg
: A sample of N images and their labels (ground-truth) from the training-set.results.png
: Plot of metrics along the training (box loss, object loss, class loss, precesion & recall).results.csv
: Actual values that were plotted inresults.png
.R_curve.png
: Recall over epochs plotted.P_curve.png
: Precision over epochs plotted.PR_curve.png
: Precision against recall plotted.F1_curve.png
: F1 score over epochs.opt.yml
: Settings used to traing the models.labels_correlogram.jpg
: Distribution of the bounding-boxes locations and dimensions.labels.jpg
: Distribution of the number of labels for each class, and their location in images.hyp.yaml
: Hyper-parameters used to train the model.confusion_matrix.png
: Confusion matrix (count per class of what was predicted vs. what was expected) of the created model on the validation set.weights
: A folder containingbest.pt
andlast.pt
which are the actual trained models.
2. Data augmentation
HSV-Hue Augmentation: The hue augmentation factor for HSV color space. Here, we work on grayscale image, so the provided value doesn’t matter.
HSV-Saturation Augmentation: The saturation augmentation factor for HSV color space.
HSV-Value Augmentation: The value augmentation factor for HSV color space. It was blocked to 0.01 to avoid making objets in the background visible.
Rotation Degrees: The maximum rotation degrees for data augmentation. Here, we allowed a range of 90° in either directions.
Translation: The maximum translation factor for data augmentation. Our objects can be anywhere on images, so we allowed a range of half the image size of each axis.
Scale: The scaling factor for data augmentation. The scale matters a lot to classify microglia, so it was locked to 1.0.
Vertical Flip Probability: The probability of performing a vertical flip during data augmentation.
Horizontal Flip Probability: The probability of performing a horizontal flip during data augmentation.
Mosaic Augmentation: The factor for mosaic data augmentation. To get more data, we create mosaics of inputs images to create
new images
.
3. Setup
Required to use YOLOv5m because there was not enough learning capacity avec YOLOv5s
Settings |
Description |
---|---|
data_folder |
Parent folder of the |
qc_folder |
Parent folder of the quality-control images (also |
inputs_name |
Name of the folder containing the inputs images and the QC inputs. |
annotations_name |
Name of the folder containing the labels for the training and QC. |
models_path |
Root of the folder in which models will be stored. |
working_directory |
Directory in which the scripts creates its temporary data. Can be deleted after training. |
model_name_prefix |
Prefix that will be given to the folders containing newly created models. |
reset_local_data |
Should the local set of images (== the data in the working directory) be reseted at every training. Recommended. |
validation_percentage |
Percentage of the provided data that will be used for the validation step. |
batch_size |
Number of images processed at the same time while training. Should be as high as your memory can handle. |
epochs |
Number of times that the whole data will be seen during training. |
classes_names |
List of class names that should be predicted by the model being trained. |
optimizer |
Optimizer used for the gradient descent. |
deterministic |
Should the inference be deterministic (one input always give the same output). Works by using a random seed if False. |
cos_lr |
Usually, the learning rate dicreases as the epochs go. If True, it will rather follow a sinusoidal curve, starting on a maxima (hence the cosine) |
label_smoothing |
Should the probability map of classes be smoothed (blurred) before building bounding boxes. |
dropout |
Percentage of neurons randomly disabled at each epoch to improve the generalization. |
4. Bind classification to segmentation
By the end of the classification process, we have a set of bounding-boxes deduced by the model. Each box has a class (garbage, amoeboid, intermediate or homeostatic).
At the previous step, we built masks representing microglia.
However, at this point, there is no relation between the segmentation and the classification. We need to bind each item from the mask to a class.
- To do that, a system of vote was implemented.
Each object starts with a set of N bins, with N being the number of classes. Along the process, each bin will count the number of votes for each class.
For each bounding-box predicted by YOLO, we search the biggest object inside it, and designate it as the target.
- In the target, the bin corresponding to the bounding-box’s class will receive P×S votes with:
P: The number of the target’s pixels in the bounding-box.
S: The certainty score of the bounding-box.
At the end, we take the major vote for each object.
An object with no vote is automatically declared garbage.