The main goal of this repo is to prepare your data for training using YOLOv5 💡.
In order to train your custom dataset using yolov5, you must create the following hierarchy ✏️:
├─dataset
│ ├─images
│ │ ├─train
│ │ ├─val
│ │ ├─test
│ ├─labels
│ │ ├─train
│ │ ├─val
│ │ ├─test
Please install requirements using:
pip install -R requirements.txt
First of all, you have to edit the dataset class. Please make sure to keep the same variable names ⛔. The __getitem__
in the dataset must return two things 🚬:
- Path to image ⚡
- Bounding boxes, which is a list that has
k
list. Each one of thek
is a list that has [cls_id
,x
,y
,w
,h
] ⚡
After that, you have to create three instances from this class, for train, val and test 💣
Now you can just create instance from the prepare class and use it:
from src.prepare import Prepare
from src.dataset import Dataset
traindataset = Dataset(**kwargs)
valdataset = Dataset(**kwargs)
testdataset = Dataset(**kwargs)
path = "path..."
prep = Prepare(train_dataset=traindataset, val_dataset=valdataset, test_dataset=testdataset, path=path)
prep.create()
This will create the above hierarchy. Then you can just copy the dataset
folder to the yolov5 repo when you clone it. Don't forget to add the yaml
file to tell yolo model the path to train, val and test, also the number of classes and the name of the classes. Here is an example:
train: ../dataset/images/train/
val: ../dataset/images/val/
test: ../dataset/images/test/
# number of classes
nc: 4
# class names
names: ['pedestrian', 'car', 'traffic light', 'traffic sign']