Home

Plastic-detection wiki

This wiki will help you to activate my yolov5 custom model!

Download dataset and source code

Go to the releases to download my dataset and source code.

Arrange the file

Arrange file just like this:

Plastic-detection
|── dataset
|   |── train
|   └── data.yaml
|── PlasticDetection.ipynb
└── test_video.mp4

Prepare to train model

Open your command prompt and go to the file Plastic-detection. If you succeed to open dir, type as below:

> git clone https://github.com/ultralytics/yolov5.git
> cd yolov5
> pip install -r requirements.txt
> cd ..

If you finished downloading, open python code editor and type as below:

from sklearn.model_selection import train_test_split

train_img_list, val_img_list = train_test_split(img_list, test_size=0.2, random_state=2000)
print(len(train_img_list), len(val_img_list))

with open('./dataset/train.txt', 'w') as f:
    f.write('\n'.join(train_img_list) + '\n')

with open('/dataset/val.txt', 'w') as f:
    f.write('\n'.join(val_img_list) + '\n')

import yaml

with open('./dataset/data.yaml', 'r') as f:
    data = yaml.load(f, Loader=yaml.Loader)

print(data)

data['train'] = './dataset/train.txt'
data['val'] = './dataset/val.txt'

with open('./dataset/data.yaml', 'w') as f:
    yaml.dump(data, f)

print(data)

Start training

Open command prompt same dir as previous one and type as below: (you can change epochs value)

> python train.py --img 416 --batch 16 --epochs 300 --data ./dataset/data.yaml --cfg .yolov5/models/yolov5n.yaml --weights yolov5n.pt --name plastic_yolov5n_results