目录

一、数据集介绍

二、Data set transformation

2.1 xml转txt

2.2 制作VOC数据集

三、yaml文件修改

3.1 数据集yaml

3.2 模型yaml

四、训练评估

4.1  训练

4.2 评估

一、数据集介绍

根据YOLOV5The lightweight small target detection network designed by study note six,本节将用tibnetThe produced dataset is used for training and testing,This dataset is used to detect aerial drones,You can see that the drone is very small.该数据集的labels文件是用labelmemarked by the softwarexml形式.

<annotation>
	<folder>0829_5JPEGImages</folder>
	<filename>0829_5092.jpg</filename>
	<path>C:\Users\lsq\Desktop\图片\0829_5JPEGImages\0829_5092.jpg</path>
	<source>
		<database>Unknown</database>
	</source>
	<size>
		<width>960</width>
		<height>540</height>
		<depth>3</depth>
	</size>
	<segmented>0</segmented>
	<object>
		<name>uav</name>
		<pose>Unspecified</pose>
		<truncated>0</truncated>
		<difficult>0</difficult>
		<bndbox>
			<xmin>477</xmin>
			<ymin>259</ymin>
			<xmax>499</xmax>
			<ymax>279</ymax>
		</bndbox>
	</object>
</annotation>

二、Data set transformation

2.1 xml转txt

xmlThe annotation format of the file is a box with four dotsx,y范围,而yolov5The format used is the center point of the box plus the width and height,Therefore, the format conversion is required,将xml文件转化为txt文件,代码如下.

import xml.etree.ElementTree as ET
 
import pickle
import os
from os import listdir , getcwd
from os.path import join
import glob
 
classes = ["uav"]
 
def convert(size, box):
 
    dw = 1.0/size[0]
    dh = 1.0/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_name):
    in_file = open('./Annotations/'+image_name[:-3]+'xml') #xml文件路径
    out_file = open('./labels/'+image_name[:-3]+'txt', 'w') #转换后的txt文件存放路径
    f = open('./Annotations/'+image_name[:-3]+'xml')
    xml_text = f.read()
    root = ET.fromstring(xml_text)
    f.close()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)
    for obj in root.iter('object'):
        cls = obj.find('name').text
        if cls not in classes:
            print(cls)
            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')
 
wd = getcwd()
 
if __name__ == '__main__':
 
    for image_path in glob.glob("./JPEGImages/*.jpg"): #每一张图片都对应一个xml文件这里写xml对应的图片的路径
        image_name = image_path.split('/')[-1]
        convert_annotation(image_name)

The converted format is as follows,第一个0代表类别,After that are the coordinates of the center point of the box and the width and height

Be sure to check after the conversiontxt中是否有值,不知道什么原因,Sometimes it converts to a null value

0 0.47890625 0.3597222222222222 0.0296875 0.05277777777777778

2.2 制作VOC数据集

 Select two-thirds of the data astrain,The remaining third asval,The directory of the dataset is as shown above

三、yaml文件修改

3.1 数据集yaml

# YOLOv5  by Ultralytics, GPL-3.0 license
# PASCAL VOC dataset http://host.robots.ox.ac.uk/pascal/VOC by University of Oxford
# Example usage: python train.py --data VOC.yaml
# parent
# ├── yolov5
# └── datasets
#     └── VOC  ← downloads here


# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
train: /home/cxl/ros_yolov5/src/yolov5/data/VOCdevkit/images/train/
val: /home/cxl/ros_yolov5/src/yolov5/data/VOCdevkit/images/val/

# Classes
nc: 1  # number of classes
names: ['uav']  # class names

3.2 模型yaml

 Major modification category,Because it's like a drone,所以nc改为1

# Parameters
nc: 1  # number of classes
depth_multiple: 1.0  # model depth multiple
width_multiple: 1.0  # layer channel multiple
anchors:

  - [2,2, 6,8, 10,14]  #4
  - [10,13, 16,30, 33,23]  # P3/8
  - [30,61, 62,45, 59,119]  # P4/16
  - [116,90, 156,198, 373,326]  # P5/32

四、训练评估

4.1  训练

python train.py --data data/VOC.yaml --cfg models/yolov5s-tiny.yaml --weights weights/yolov5stiny.pt --batch-size 16 --epochs 100

 查看训练过程

tensorboard --logdir=./runs

4.2 评估

 可以看到效果不错,map0.5达到了0.94,loss接近于0

 Save the trained weights as yolov5suav.pt,随后进行测试

 测试

python detect.py --source ./data/images/ --weights weights/yolov5suav.pt --conf 0.4
detect: weights=['weights/yolov5suav.pt