(pytorch) how to train your dataset with yolov3 (pro test is effective)

In view of the failure of many tutorials on training my data sets with yolov3 algorithm on the network, I decided to write one myself. Train in the latest version of pytorch.

First, we should prepare our own data sets (that is, pictures), and then use our annotation tools to annotate them

1. Labeling tool (labelimg)

Labelimg is an open source data annotation tool, which can annotate three formats.

1. VOC tag format, saved as xml file.

2. yolo label format, saved as txt file.

3. The createml tag format is saved in json format.

I won't say much about downloading. You can download it from the command line or the official website (I installed it from the command line)

pip install labelimg -i https://pypi.tuna.tsinghua.edu.cn/simple

2. Preparation before marking

First, we need to prepare the data set we need to mark. Here, I suggest to create a new folder named voc207, in which a folder named JPEGImages is created to store the image files we need to label; Then create a label file named Annotations to store labels; Finally, create a file called predefined_classes.txt file to store the category name to be marked.

The directory structure of VOC2007 is:

├── VOC2007
│♪ JPEGImages store image files that need to be labeled
│♪ - Annotations ♪ store labeled label files
│├── predefined_classes.txt define all categories to be labeled (this file is optional, but when we define many categories, it's best to create such a txt file to store categories)

2. Label the dataset

First, put the pictures to be marked in the JPEGImages folder. Here are six kinds of pictures. Here are the defect pictures of the aircraft skin I use. They are fadongjihuahen, mengpiwuzi, biaomianyiwu, mengpikongdong, mengpliewen and mengpidiaoqi.

And then {predefined_classes.txt enter the defined category in the txt document; As shown in the figure below.

Then enter the VOC2007 path just created from the terminal


Enter the following command to open labelimg. This command means to open labelimg tool; Open the JPEGImage folder and initialize predefined_classes.txt.

labelimg JPEGImages predefined_classes.txt

Running the above command will open the tool; As follows.  

Let's find out how to use this software by ourselves. I won't talk about it in detail

After the tag is typed, you can go to the Annotations} file to see that the tag file has been saved in this directory.

3. Installing yolov3

Now I'm going to download and compile yolov3. There are many versions of yolov3. I recommend you download my version. There will be various problems in other versions

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

These lines of code will download yolov3 and related dependencies.

From the official tutorial on github, you can know the following line of code for training coco128 data

python train.py --img 640 --batch 16 --epochs 5 --data coco128.yaml --weights yolov3.pt

In this way, we only need to put coco128 after data Just change yaml to our own file format.

coco128. The contents in yaml are as follows

# download command/URL (optional)
download: https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128.zip

# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: ../coco128/images/train2017/
val: ../coco128/images/train2017/

# number of classes
nc: 80

# class names
names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
        'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
        'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
        'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
        'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
        'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
        'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 
        'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 
        'teddy bear', 'hair drier', 'toothbrush']

As you can see, train and val are two paths. These two paths are our dataset related files. Don't worry about download ing.

The following class names are our list of categories, so we just need to create our own train and val paths

Here's how to create

4. Create a custom dataset related profile

Create an ImageSet folder under voc207 and a main folder inside

- VOC2007
            - Annotations (label XML File, manually generated with the corresponding image processing tool)
            - ImageSets (The method of generation is to use sh perhaps MATLAB Language generation)
                - Main
                    - test.txt
                    - train.txt
                    - trainval.txt
                    - val.txt
            - JPEGImages(Original file)
            - labels (xml File corresponding txt file)

As shown in the figure above

The contents in JPEGImages and Annotations folders are mainly constructed by labelimg software, and the txt files in the Main folder can be generated by the following python scripts:

import os
import random
trainval_percent = 0.9
train_percent = 1
xmlfilepath = 'Annotations'
txtsavepath = 'ImageSets\Main'
total_xml = os.listdir(xmlfilepath)
trainval= random.sample(list,tv)
ftrainval = open('ImageSets/Main/trainval.txt', 'w')
ftest = open('ImageSets/Main/test.txt', 'w')
ftrain = open('ImageSets/Main/train.txt', 'w')
fval = open('ImageSets/Main/val.txt', 'w')
for i  in list:
    if i in trainval:
        if i in train:

Next, generate the txt file in the labels folder, VOC_ label. The details of the PY file are as follows:

# -*- coding: utf-8 -*-
Created on Tue Oct  2 11:42:13 2018
 Place this document in VOC2007 Under the same level directory, and then you can run it directly
 Areas to be modified:
1. sets Replace with your own dataset
2. classes Replace with your own category in
3. Place this document in VOC2007 Directory
4. Start running directly

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
sets=[('2007', 'train'), ('2007', 'val'), ('2007', 'test')]  #Replace with your own dataset
classes = ["person"]     #Modify to your own category

def convert(size, box):
    dw = 1./(size[0])
    dh = 1./(size[1])
    x = (box[0] + box[1])/2.0 - 1
    y = (box[2] + box[3])/2.0 - 1
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x*dw
    w = w*dw
    y = y*dh
    h = h*dh
    return (x,y,w,h)

def convert_annotation(year, image_id):
    in_file = open('VOC%s/Annotations/%s.xml'%(year, image_id))  #Put the dataset in the current directory
    out_file = open('VOC%s/labels/%s.txt'%(year, image_id), 'w')
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)
    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult)==1:
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
        bb = convert((w,h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
wd = getcwd()
for year, image_set in sets:
    if not os.path.exists('VOC%s/labels/'%(year)):
    image_ids = open('VOC%s/ImageSets/Main/%s.txt'%(year, image_set)).read().strip().split()
    list_file = open('%s_%s.txt'%(year, image_set), 'w')
    for image_id in image_ids:
        list_file.write('VOC%s/JPEGImages/%s.jpg\n'%(year, image_id))
        convert_annotation(year, image_id)
voc_ labels. The txt file generated by py needs to be reorganized and divided into training set and verification set

Let's talk about the code division later

We continue to create a new folder in the VOC2007 folder

         - VOC2007
            - Annotations (label XML File, with the corresponding labelimg Generated)
            - ImageSets (The method of generation is to use python generate)
                - Main
                    - test.txt
                    - train.txt
                    - trainval.txt
                    - val.txt
            - JPEGImages(Original file)
            - labels (xml File corresponding txt file)
            - Allempty (Used to store unqualified pictures)
            - images (be used for pytorch Version of the picture is saved)
            	- train2014
            		- 001.jpg
            		- 002.jpg
            	- val2014
            		- 100.jpg
            		- 101.jpg
            - label (be used for pytorch Label save for version)
            	- train2014
            		- 001.txt
            		- 002.txt
            	- val2014
            		- 100.txt
            		- 101.txt

The txt file in the labels we generated before and the files generated by code division are in the label folder

We create the python script file make_for_yolov3_torch.py is placed in the same level folder of voc207

import os, shutil

The following conditions need to be met:
1. stay JPEGImages Prepare pictures in
2. stay labels Ready in labels
3. Create the following file directory:
    - images
        - train2014
        - val2014
    - labels(because voc There are in the format labels Folder, so rename it to label)
        - train2014
        - val2014

def make_for_torch_yolov3(dir_image, 
    if not os.path.exists(dir1_train):
    if not os.path.exists(dir1_val):
    if not os.path.exists(dir2_train):
    if not os.path.exists(dir2_val):

    with open(main_trainval, "r") as f1:
        for line in f1:
            # print(os.path.join(dir_image, line[:-1]+".jpg"), os.path.join(dir1_train, line[:-1]+".jpg"))
            shutil.copy(os.path.join(dir_image, line[:-1]+".jpg"),
                        os.path.join(dir1_train, line[:-1]+".jpg"))
            shutil.copy(os.path.join(dir_label, line[:-1]+".txt"), 
                        os.path.join(dir2_train, line[:-1]+".txt"))

    with open(main_test, "r") as f2:
        for line in f2:
            shutil.copy(os.path.join(dir_image, line[:-1]+".jpg"), 
                        os.path.join(dir1_val, line[:-1]+".jpg"))
            shutil.copy(os.path.join(dir_label, line[:-1]+".txt"), 
                        os.path.join(dir2_val, line[:-1]+".txt"))

if __name__ == "__main__":
    this pytorch Data set organization for version
    - images
        - train2014 # dir1_train
        - val2014 # dir1_val
    - labels
        - train2014 # dir2_train
        - val2014 # dir2_val
    trainval.txt, test.txt By create_main.py Constructed

    dir_image = r"C:\Users\pprp\Desktop\VOCdevkit\VOC2007\JPEGImages"
    dir_label = r"C:\Users\pprp\Desktop\VOCdevkit\VOC2007\labels"

    dir1_train = r"C:\Users\pprp\Desktop\VOCdevkit\VOC2007\images\train2014"
    dir1_val = r"C:\Users\pprp\Desktop\VOCdevkit\VOC2007\images\val2014"

    dir2_train = r"C:\Users\pprp\Desktop\VOCdevkit\VOC2007\label\train2014"
    dir2_val = r"C:\Users\pprp\Desktop\VOCdevkit\VOC2007\label\val2014"

    main_trainval = r"C:\Users\pprp\Desktop\VOCdevkit\VOC2007\ImageSets\Main\trainval.txt"
    main_test = r"C:\Users\pprp\Desktop\VOCdevkit\VOC2007\ImageSets\Main\test.txt"


The file path in the code needs to be modified by ourselves. Then we copy the train2014 folder and val2014 folder under the images file to the data images under yolov3, and put the train2014 and val2014 under the label under the data labels of yolov3, as shown in the following figure

Then create a custom configuration file in the cfg folder of yolov3, and name it yolov3.0 for the time being Yaml, imitation coco 128 The format in yaml. Just change the tran and val paths to our path, comment out the download address, and change the category to our own category.


So far, the data set conforming to yolov3 format has been constructed and can be trained with official code

python train.py --img 640 --batch 16 --epochs 50 --data cfg/yolov3.yaml --weights yolov3.pt

You may need to install a wandb before training

pip install wandb

Then you can register a wandb account. This is a visual website, which can easily see your training results. It works well

Weights & Biases – Developer tools for ML


Keywords: Pytorch Deep Learning Object Detection

Added by chmpdog on Thu, 16 Dec 2021 20:49:28 +0200