YOLOv5 usage notes

Friendly participation: @Cotton cotton

YOLO is a target detection method, which is characterized by fast detection and high accuracy. The author regards the target detection task as a regression problem of target region prediction and category prediction. This method uses a single neural network to directly predict the boundary and category probability of items, and realizes end-to-end item detection. At the same time, the detection speed of this method is very fast, and the basic version can achieve real-time detection of 45 frames / s; FastYOLO can reach 155 frames / s. Compared with the current best system, the positioning error of YOLO target area is greater, but the false positive of background prediction is better than the current best method.

Knowledge is limited. Once again, we do not deeply explore the network structure and principle of YOLO, but only record how to use it.

Objective: to detect the characters in the video

1, Install & configure environment

YOLOv5 Download

I use anaconda for environment configuration

conda create --name Yolov5 python=3.8

After creating the new environment, there is a requirement in the installation path of YOLOv5 Txt file, which lists all the package s required for YOLOv5. Open cmd and enter

pip install -r requirements.txt

The installation can be completed automatically. The installation may fail halfway. It is probably a network problem. Just reinstall it.
Since the pytorch it automatically downloads is the cpu version, we need to delete the downloaded pytorch first

pip uninstall torch

Then go Official website Manually download the gpu version of pytorch. After the test can use gpu, the configuration of the environment is completed.


2, Get data

Here I want to extract several frames from the video as data for annotation. The frame extraction code is as follows. In order to facilitate future use, I directly encapsulated it into class:
package required: cv2, os. There should be no need to install it separately. It will be installed when YOLOv5 is configured

class Video2Frame:
    def __init__(self, videoPath, savePath, timeInterval=1):
        Initialization, get videoPath All videos under

         @ videoPath: Storage path of video
         @ framesSavePath: Save path of pictures after video segmentation into frames
         @ timeInterval: Save interval
        self.videoPath = videoPath
        self.savePath = savePath
        self.timeInterval = timeInterval
        self.videos = os.listdir(videoPath)

    def getVideoList(self):
        return self.videos

    def getFrameFromVideos(self, start=0, end=None):
        Create a folder for each video and save it frame
         @ start: From which video
         @ end: To which end (barring end)
        length = len(self.videos)
        if end == None:
            end = length - 1
        start = start % length
        end = end % length + 1
        for i in range(start, end):

    def getFrameFromVideo(self, i=0):
        Create a folder for a video and save it frame
         @ i: The first i Videos
        video = self.videos[i]
        folderName = video[:-4]
        if not os.path.exists(folderName):

        vidCap = cv2.VideoCapture(self.videoPath + "\\" + video)
        success, image = vidCap.read()
        cnt = 0
        while success:
            success, image = vidCap.read()
            cnt += 1
            if cnt % self.timeInterval == 0:
                cv2.imencode('.jpg', image)[1].tofile(self.savePath + "\\" + folderName + r"\frame%d.jpg" % cnt)
            # if cnt == 200:
            #     break
        print(folderName + ": ", cnt // self.timeInterval, "images")

if __name__ == "__main__":
    videos_path = r'.\video'
    frames_save_path = r'.\data'
    time_interval = 30
    v2f = Video2Frame(videos_path, frames_save_path, time_interval)

3, Label dataset

Labeling datasets requires a tool, labelImg( Download link ). After downloading and decompressing, the first thing to do is to delete labelimg master \ data \ predefined_ classes. Txt file, otherwise some preset categories will be automatically added when marking.

Then install pyqt5 (pay attention to the YOLOv5 environment)

pip install pyqt5

If an error is reported during installation, you need to install SIP first

pip install SIP

After installation, run cmd in the installation directory of labelImg, switch to the YOLOv5 environment, and enter

python labelImg.py

labelImg can be opened, and then the selected pictures can be marked.

4, train

Once the data is ready, you can start training the model. But before we start training, we need to do some preparations.

4.1 preparations before training

We first need to create the following folders:

data folder, detect folder and labels Cache is optional. Let me explain the functions of these documents one by one.
The data folder is used to put all the data. The detect folder is used to put files for prediction. I put videos or pictures here. Under the train folder are two subfolders, images and labels, which respectively contain pictures and labels for training. The test folder has the same structure as train and is used to store test data. data.yaml is a configuration file. We need to write the following code in the file:

train: ../ApexCheater/train/images
val: ../ApexCheater/test/images

nc: 2
names: ['body', 'head']

train and val are relative to the YOLOv5 installation path. nc indicates the number of categories, and names indicates the names of various categories. Pay attention to the order.

4.2 dividing data sets

Due to the small scale of the data, I divide it according to the ratio of 7:3, and the code is as follows:
package required: os, shutil, random

class separateData:
    def __init__(self, lst, dataPath):
        Initialise the class and print the size of train, test, detect set

          @ lst: Represents the proportion of train, test, detect set
          @ dataPath: The path of the data set

        # Get the total size of the data set
        self.tot = len(os.listdir(dataPath + r"\images"))
        self.dataPath = dataPath
        # Get train size, test size and print them
        self.trainSize = int(self.tot * lst[0])
        self.testSize = int(self.tot * lst[1])
        print("train size: ", self.trainSize)
        print("test size: ", self.testSize)

        # If you need to separate data into detect set
        if len(lst) == 3:
            self.detectSize = int(self.tot * lst[2])
            print("detect size: ", self.detectSize)

    def getRandomList(self, num, selected=[]):
        This function is used to choose some data randomly.

         @ num: The number of the elements of list.
         @ selected: A list represents previously selected data. Default value selected=[]

         Return a list represents the serial number of the data, randomly.
        lst = []
        for i in range(0, num):
            # Loops until r is not in lst and selected. That means r
            # was never chosen in this for loop or previous program.
            while True:
                r = random.randint(0, self.tot - 1)
                if not((r in lst) or (r in selected)):
        return lst

    def copyData(self, i, path):
        There must be two folders under the given path, 'images' and 'labels'
        and this function will copy the images and labels into given folders respectively.

         @ i: The serial number of the data.
         @ path: Where the data will be copied to.
        JPG = self.dataPath + r"\images\apexCheater(" + str(i) + ").jpg"
        TXT = self.dataPath + r"\labels\apexCheater(" + str(i) + ").txt"
        shutil.copy(JPG, path + r"\images")
        shutil.copy(TXT, path + r"\labels")

    def deleteData(self, path):
        Delete the data under the given path.
        There must be two folders under the given path, 'images' and 'labels'.

         @ path: The path you want to delete data under it.
        if path == "":

        jpgPath = path + r"\images"
        txtPath = path + r"\labels"
        for data in os.listdir(jpgPath):
        for data in os.listdir(txtPath):

    def separateTrainTestDetect(self, trainPath, testPath, detectPath="", isDataDelete=False):
        Separate data set into train, test, detect set, or just train and test set.

         @ trainPath: The path of train set.
         @ testPath: The path of test set.
         @ detectPath: The path of detect set. Default value detectPath="".
         @ isDataDelete: Whether to delete original data. Default value isDataDelete=False.
        if isDataDelete:

        testList = self.getRandomList(self.testSize)
        detectList = self.getRandomList(self.detectSize, testList)

        for i in detectList:
            self.copyData(i, detectPath)
        for i in testList:
            self.copyData(i, testPath)
        for i in range(0, self.tot):
            if not((i in testList) or (i in detectList)):
                self.copyData(i, trainPath)

if __name__ == "__main__":
    rootPath = r".\ApexCheater"
    dataPath = rootPath + r"\data"
    trainPath = rootPath + r"\train"
    testPath = rootPath + r"\test"
    detectPath = rootPath + r"\detect"

    sd = separateData([0.7, 0.3, 0], dataPath)
    sd.separateTrainTestDetect(trainPath, testPath, isDataDelete=True)

4.3 start training

I use pycharm because YOLOv5 provides a tutorial on training with jupyter notebook. I won't expand it here. I just remember how to use pycharm for training.

First open the YOLOv5 project, and then open train Py file

Configure the following parameters:

--data ../ApexCheater/data.yaml
--cfg models/yolov5s.yaml
--batch-size 8
--epoch 600

There are many models in the models folder, which can be selected according to the situation. Click Run after configuration.

The solution to the error "the page file is too small to complete the operation":

The reason for this error is that the virtual memory is too small. Just set the virtual memory of Disk C to be larger. However, according to Mian Mian, it should be to expand the virtual memory of the disk corresponding to the python installation location. But the specific situation is: my conda is on disk G and occupies virtual memory of C; The continuous conda is in disk E, and the virtual memory of disk E is occupied. In addition, running the same data and model, the continuous memory occupation is almost four times that of me. We don't know the specific reason and can't verify it.

5, detect

Open detect. In the YOLOv5 project Py file, configure the following parameters:

--weight ./runs/train/exp12/weights/best.pt
--img 640
--conf 0.25
--source ../ApexCheater/detect

Then run the program. Note --img's image size must be consistent with that of train.

6, Broken thoughts

In order to run the model, I changed the memory module with Mianmian. It's cool to play games after changing it!


Keywords: Python Machine Learning Pytorch Deep Learning yolo

Added by petrb on Sun, 20 Feb 2022 18:11:05 +0200