Table of contents
labelimg makes its own training data set
Convert your own training data set into training set and test set
YOLOv5 trains its own data set
Data, pre-training weight preparation
labelimg makes its own training data set
labelimg install
For installation in the window system, enter the following command in Anaconda Prompt:
pip install labelimg -i https://pypi.tuna.tsinghua.edu.cn/simpledata preparation
Create a new VOC2007, and create 3 new files JPEGImages, Annotations, and predefined_classes.txt (not required) under this file. It is best to name them in the same way to facilitate subsequent operations;
JPEGImages: Store images that need to be labeled;
Annotations: Store annotation labels File;
predefined_classes.txt: Define all the categories you want to mark (when there are many categories defined, it is best to create such a txt file to store the categories);
data tag
Enter in Anaconda Prompt: labelimg to open the labeling tool
: Open a single picture to be marked;
: the folder to label the picture;
: This button can indicate that the label we marked is in voc format, click to change it to yolo or createML format.
: Select the target window to draw;
We choose the VOC format and place the marked files under the Annotations file
Convert your own training data set into training set and test set
The format of the data set resource tags for general target detection is VOC: xml format, and the file format required for yolov5 training is yolo: txt format, you need to convert the tag file in xml format to txt file. At the same time, when training your own yolov5 detection model, the data set needs to be divided into a training set and a verification set. File format: Create a new folder VOCdevkit, put the above VOC2007 file into this folder; run the code as follows - the file format is the same, no need to change
TRAIN_RATIO = 80: The ratio of training set to verification set is 80%, you can change it yourself
import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
import random
from shutil import copyfile
# 自己标注的类型
classes=["mouse"]
#训练集与验证集比例80%，自己可改
TRAIN_RATIO = 80
def clear_hidden_files(path):
dir_list = os.listdir(path)
for i in dir_list:
abspath = os.path.join(os.path.abspath(path), i)
if os.path.isfile(abspath):
if i.startswith("._"):
os.remove(abspath)
else:
clear_hidden_files(abspath)
def convert(size, box):
dw = 1. / size[0]
dh = 1. / size[1]
x = (box[0] + box[1]) / 2.0
y = (box[2] + box[3]) / 2.0
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(image_id):
in_file = open('VOCdevkit/VOC2007/Annotations/%s.xml' % image_id)
out_file = open('VOCdevkit/VOC2007/YOLOLabels/%s.txt' % image_id, 'w')
tree = ET.parse(in_file)
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:
continue
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')
in_file.close()
out_file.close()
wd = os.getcwd()
wd = os.getcwd()
data_base_dir = os.path.join(wd, "VOCdevkit/")
if not os.path.isdir(data_base_dir):
os.mkdir(data_base_dir)
work_sapce_dir = os.path.join(data_base_dir, "VOC2007/")
if not os.path.isdir(work_sapce_dir):
os.mkdir(work_sapce_dir)
annotation_dir = os.path.join(work_sapce_dir, "Annotations/")
if not os.path.isdir(annotation_dir):
os.mkdir(annotation_dir)
clear_hidden_files(annotation_dir)
image_dir = os.path.join(work_sapce_dir, "JPEGImages/")
if not os.path.isdir(image_dir):
os.mkdir(image_dir)
clear_hidden_files(image_dir)
yolo_labels_dir = os.path.join(work_sapce_dir, "YOLOLabels/")
if not os.path.isdir(yolo_labels_dir):
os.mkdir(yolo_labels_dir)
clear_hidden_files(yolo_labels_dir)
yolov5_images_dir = os.path.join(data_base_dir, "images/")
if not os.path.isdir(yolov5_images_dir):
os.mkdir(yolov5_images_dir)
clear_hidden_files(yolov5_images_dir)
yolov5_labels_dir = os.path.join(data_base_dir, "labels/")
if not os.path.isdir(yolov5_labels_dir):
os.mkdir(yolov5_labels_dir)
clear_hidden_files(yolov5_labels_dir)
yolov5_images_train_dir = os.path.join(yolov5_images_dir, "train/")
if not os.path.isdir(yolov5_images_train_dir):
os.mkdir(yolov5_images_train_dir)
clear_hidden_files(yolov5_images_train_dir)
yolov5_images_test_dir = os.path.join(yolov5_images_dir, "val/")
if not os.path.isdir(yolov5_images_test_dir):
os.mkdir(yolov5_images_test_dir)
clear_hidden_files(yolov5_images_test_dir)
yolov5_labels_train_dir = os.path.join(yolov5_labels_dir, "train/")
if not os.path.isdir(yolov5_labels_train_dir):
os.mkdir(yolov5_labels_train_dir)
clear_hidden_files(yolov5_labels_train_dir)
yolov5_labels_test_dir = os.path.join(yolov5_labels_dir, "val/")
if not os.path.isdir(yolov5_labels_test_dir):
os.mkdir(yolov5_labels_test_dir)
clear_hidden_files(yolov5_labels_test_dir)
train_file = open(os.path.join(wd, "yolov5_train.txt"), 'w')
test_file = open(os.path.join(wd, "yolov5_val.txt"), 'w')
train_file.close()
test_file.close()
train_file = open(os.path.join(wd, "yolov5_train.txt"), 'a')
test_file = open(os.path.join(wd, "yolov5_val.txt"), 'a')
list_imgs = os.listdir(image_dir) # list image files
prob = random.randint(1, 100)
print("Probability: %d" % prob)
for i in range(0, len(list_imgs)):
path = os.path.join(image_dir, list_imgs[i])
if os.path.isfile(path):
image_path = image_dir + list_imgs[i]
voc_path = list_imgs[i]
(nameWithoutExtention, extention) = os.path.splitext(os.path.basename(image_path))
(voc_nameWithoutExtention, voc_extention) = os.path.splitext(os.path.basename(voc_path))
annotation_name = nameWithoutExtention + '.xml'
annotation_path = os.path.join(annotation_dir, annotation_name)
label_name = nameWithoutExtention + '.txt'
label_path = os.path.join(yolo_labels_dir, label_name)
prob = random.randint(1, 100)
print("Probability: %d" % prob)
if (prob < TRAIN_RATIO): # train dataset
if os.path.exists(annotation_path):
train_file.write(image_path + '\n')
convert_annotation(nameWithoutExtention) # convert label
copyfile(image_path, yolov5_images_train_dir + voc_path)
copyfile(label_path, yolov5_labels_train_dir + label_name)
else: # test dataset
if os.path.exists(annotation_path):
test_file.write(image_path + '\n')
convert_annotation(nameWithoutExtention) # convert label
copyfile(image_path, yolov5_images_test_dir + voc_path)
copyfile(label_path, yolov5_labels_test_dir + label_name)
train_file.close()
test_file.close()
YOLOv5 trains its own data set
project download
First open yolov5's github official website GitHub - ultralytics/yolov5 at v5.0 : the opened official website interface is as follows, glenn-jocher open source yolov5 project
Open the project directory as follows:
data : It is mainly a configuration file that stores some hyperparameters (.yaml file: it is used to configure the path of training set, test set, verification set, and also includes the number and name of the types of target detection); if you train your own data set , then you need to modify the .yaml file. However, it is not recommended to put your own data set under this path, but it is recommended to put the data set under the same directory of the yolov5 project.
models : Mainly some configuration files and functions for network construction, including four different versions of the project, namely s, m, l, and x. The size of these several versions. Their detection measures are from fast to slow, but the accuracy is from low to high. If you train your own data set, you need to modify the corresponding yaml file to train your own model.
utils : Functions for placing tools, including loss functions, metrics functions, plots functions, etc.
weights : Place the trained weight parameters.
detect.py : Use the trained weight parameters for target detection, which can detect images, videos and cameras.
train.py : Functions to train your own dataset.
test.py : A function to test the results of training.
requirements.txt : This is a text file that contains some versions of the environment-dependent packages that use the yolov5 project. You can use this text to import the corresponding version of the package.
Data, pre-training weight preparation
1. Modify the data configuration file: one is the corresponding yaml file under the data directory, and the other is the corresponding yaml file under the model directory file. Modify the voc.yaml file in the data directory, copy the file, and rename the copied file, preferably related to the project, mouse.yaml.
Note: Note download, otherwise an error will be reported; names: the category is English, and Chinese is not recognized; it is best to put an absolute path in the path, so it is not easy to report an error;
2. Modify the model configuration file: the corresponding parameters in the yolov5s.yaml file in the models directory. Make a copy of the yolov5s.yaml file and rename it to yolov5_mouse.yaml.
To open the yolov5_mouse.yaml file, you only need to modify the numbers in the figure
3. Pre-training weight download URL: https://github.com/ultralytics/yolov5/releases , put yolo5s.pt in the weights folder;
model training
The main parameter analysis of the train.py file model is as follows:
- weights: the path address of the initialized weight file
- cfg: the path address of the model yaml file
- data: the path address of the data yaml file
- hyp: hyperparameter file path address
- epochs: training rounds
- batch-size: how much to feed batch files
- img-size: input image size
- rect: Whether to use rectangle training, the default is False
- resume: Then interrupt the last result of training and continue training
- nosave: Do not save the model, default False
- notest: no test, default False
- noautoanchor: Do not automatically adjust the anchor, the default is False
- Volve: Whether to perform hyperparameter evolution, the default is False
- bucket: Google cloud disk bucket, generally not used
- cache-images: Whether to cache images in memory in advance to speed up training, default False
- image-weights: use weighted image selection for training
- device: training equipment, cpu; 0 (represents a gpu device cuda:0); 0,1,2,3 (multiple gpu devices)
- multi-scale: Whether to perform multi-scale training, the default is False
- single-cls: Whether the data set has only one category, the default is False
- adam: whether to use the adam optimizer
- ync-bn: Whether to use cross-card synchronization BN, used in DDP mode
- local_rank: DDP parameter, do not modify
- workers: the maximum number of working cores
- project: the storage location of the training model
- name: the name of the directory where the model is saved
- exist-ok: Whether the model directory exists, create it if it does not exist
These parameters must be modified parameters:
parser.add_argument('--weights', type=str, default='weights/yolov5s.pt', help='initial weights path')
parser.add_argument('--cfg', type=str, default='models/yolov5s_mouse.yaml', help='models/yolov5s_gps.yaml')
parser.add_argument('--data', type=str, default='data/mouse.yaml', help='data.yaml path')epochs, batch-size can be set according to its own data, run train.py for model training;
reasoning test
After the data is trained, a run folder will be generated in the main directory, and two weight files will be generated in the run/train/exp/weights directory, one is the weight file of the last round, and the other is the best weight file
detect.py in the main directory, the parameter analysis is as follows:
- weights: the path address of the weight
- source: test data, which can be a picture/video path, or '0' (the computer comes with a camera), or a video stream such as rtsp
- output: the save path of the picture/video after network prediction
- img-size: network input image size
- conf-thres: confidence threshold
- iou-thres: do the iou threshold of nms
- device: whether to use GPU or CPU for inference
- view-img: Whether to display the picture/video after prediction, the default is False
- save-txt: Whether to save the predicted box coordinates in a txt file, the default is False
- classes: Set to keep only a certain category, in the form of 0 or 0 2 3
- agnostic-nms: Whether to remove the boxes between different categories when performing nms, the default is False
- Augment: Perform multi-scale, flip and other operations (TTA) reasoning during reasoning
- update: If it is True, the strip_optimizer operation will be performed on all models, and the optimizer and other information in the pt file will be removed. The default is False
- project: The results of reasoning are saved in the runs/detect directory
Modify these two places of the main function and run detect.py;
parser.add_argument('--weights', nargs='+', type=str, default='runs/train/exp/weights/best.pt', help='model.pt path(s)')
parser.add_argument('--source', type=str, default='127.jpg', help='source') 
Some error reports during the running process can be found in my other blog