目录
根据YOLOV5学习笔记六所设计的轻量化小目标检测网络,本节将用tibnet制作的数据集进行训练测试,该数据集是用来检测空中无人机的,可以看到无人机十分的小。该数据集的labels文件是用labelme软件进行标注的xml形式。

- <annotation>
- <folder>0829_5JPEGImagesfolder>
- <filename>0829_5092.jpgfilename>
- <path>C:\Users\lsq\Desktop\图片\0829_5JPEGImages\0829_5092.jpgpath>
- <source>
- <database>Unknowndatabase>
- source>
- <size>
- <width>960width>
- <height>540height>
- <depth>3depth>
- size>
- <segmented>0segmented>
- <object>
- <name>uavname>
- <pose>Unspecifiedpose>
- <truncated>0truncated>
- <difficult>0difficult>
- <bndbox>
- <xmin>477xmin>
- <ymin>259ymin>
- <xmax>499xmax>
- <ymax>279ymax>
- bndbox>
- object>
- annotation>
xml文件的标注格式是一个框的四个点的x,y范围,而yolov5使用的格式是框的中心点加上宽高,所以需要进行格式的转化,将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)
转化后的格式如下,第一个0代表类别,之后是框的中心点坐标和宽高
转化完后一定要检查一下txt中是否有值,不知道什么原因,有时会转化为空值
0 0.47890625 0.3597222222222222 0.0296875 0.05277777777777778

选取三分之二的数据作为train,剩下的三分之一作为val,数据集的目录如上图
- # 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
主要修改类别,因为就无人机一类,所以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
python train.py --data data/VOC.yaml --cfg models/yolov5s-tiny.yaml --weights weights/yolov5stiny.pt --batch-size 16 --epochs 100

查看训练过程
tensorboard --logdir=./runs

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

将训练好的权重保存为yolov5suav.pt,随后进行测试
测试
- python detect.py --source ./data/images/ --weights weights/yolov5suav.pt --conf 0.4
- detect: weights=['weights/yolov5suav.pt

