YOLOv5 with Pseudo Labels

This code is a modified version based on YOLOv5 provided by the repository(https://github.com/ultralytics/yolov5). The purpose of this repository is to augment train data by pseudo labels so this repository contains a generator of pseudo labels not provided by the original code.

Tutorial

If you want to learn about training, testing, customizing, and so on, please refer to the original repository.

• Train Custom Data
• Multi-GPU Training
• PyTorch Hub
• ONNX and TorchScript Export
• Test-Time Augmentation (TTA)
• Model Ensembling
• Model Pruning/Sparsity
• Hyperparameter Evolution
• TensorRT Deployment

Generator of pseudo labels

Before using the generator, prepare a pretrained model. The generator makes pseudo labels and their images based on your model. Next, make "extra" in "images" folder, and "pseudo" folders in "images" and "labels" directories. "extra" directory includes unlabelled images. "pseudo" is a folder to save pseudo labels and their images.

# custom.yaml
# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: ../coco128/images/train2017/train/
val: ../coco128/images/train2017/valid/
pseudo: ../coco128/images/train2017/pseudo/ # Selected images from the extra folder for pseudo labels
extra: ../coco128/images/train2017/extra/ # Prepared images for pseudo labels(candidate group for pseudo labels)

Finally, use the code with a threshold which means the lower bound of class confidence in [0,1] range.

$ python pseudolabel.py --data custom.yaml --pseudo_threshold 0.4 --weights 'yolo5s.pt'                                   

Training with pseudo labels

If --pseudo False, the code is running with no difference from the original code. Otherwise, a model uses both train and pseudo labels for training.

$ python train.py --data custom.yaml --cfg yolov5s.yaml --weights '' --batch-size 64 --pseudo True                                    

To Do

• Modification of the loss function
• Scheduling of a loss penalty for pseudo labels