Implementation of the YOLO algorithm (version 1) paper in PyTorch: You Only Look Once: Unified, Real-Time Object Detection.
The YOLO model was trained on Pascal VOC dataset.
The hyperparameters of the model can be found in the config.yaml
file.
src/
metrics/
cbow.py
- implemented CBOW model.iou.py
- implemented Skip-Gram model.mean_average_precision.py
- contains common function used for models.
model/
model.py
- YOLO architecture.
utils
utils_file
- utilities for loading and saving files.utils_list.py
- utilities for operating on lists.
dataset.py
- implementation of the PyTorch Dataset in Pascal VOC format.loss.py
- implementation of the loss function used for training YOLO.main.py
- full training pipeline.non_max_suppression.py
- implementation of the Non-maximum Suppression used during the model inference.train.py
- class implementation for training a model.
test/
- collection of unit testsconfig.yaml
- config file with the main parameters of datasets and the model.
- Clone the repository into the
yolo
folder:
git clone https://github.com/slavafive/YOLO-v1.git yolo
- Download the Pascal VOC dataset. Save it in the
yolo
folder as thedata
directory. - Run the training pipeline:
python yolo/src/main.py --config yolo/config.yaml
- After the training, the model and its checkpoints will be saved in the directory specified in the
model_directory
(artifacts
by default) attribute of theconfig.yaml
file.
A part of the source code and unit tests were used from Aladdin Persson repository.