Image classification using yolov5

insert image description here

Article directory

before the start

You should clone this repository first

git clone https://github.com/ultralytics/yolov5 # clone

After the download is complete, enter the cloned warehouse directory

cd yolov5

download dependencies

pip install -r requirements.txt # install

Dataset download

Here I have prepared a data set. For the convenience of downloading, there are not many data sets in the data set.At the end I will share a few data set download addresses：

Garbage classification dataset download
Extraction code: nr5i

After unzipping, you will see these folders:

Please add a picture description
View part of the content at random.

I decompressed here to amydataTable of contents. This is not required, but you will need to be able to find your dataset directory later.

New configuration file

The meaning of each file is roughly as follows:

  0: cardboard #纸板
  1: glass #玻璃
  2: metal #金属
  3: paper #纸
  4: plastic #塑料
  5: trash #垃圾

executive training

Before starting, please download yolov5s-cls.ptthe model in advance, remember this location, because you will need it to start training below.

Click to download the yolov5s-cls.pt model

python classify/train.py --model yolov5s-cls.pt --data mydata --epochs 5 --img 224 --batch 128

model selection

It is used when performing training --model yolov5s-cls.pt. This is a model. You can refer to the specific selection in the figure below. If we choose yolov5x, you can use it --model yolov5x-cls.pt. It is recommended to use it yolov5s. Unless you require very high accuracy, otherwise you will need to spend a very long time and Sufficient hardware support to train it.

insert image description here

training completed

That best.ptis the trained model, it is in runs/the directory

insert image description here

Test the model to make predictions

Choose an image to test:

insert image description here

python classify/predict.py --weights runs/train-cls/exp9/weights/best.pt --source metal4.jpg

insert image description here

Congratulations, you have successfully trained a simple classification model.

Custom model download

If you don't want to train the model from scratch, you can download this trained model for the previous stepTest the model to make predictions。