【红外目标检测】用YOLOv5实现对VisDrone-DroneVehicle的目标检测

数据集介绍

不同时段不同城市区域背景的城市道路、居民区、停车场、高速公路 半红外 无人机视角 15532对 840*712 下载链接 https://github.com/VisDrone/DroneVehicle

数据集处理

我们只做红外部分的检测，将数据集下载修改成如下格式：

说明：在YOLO5源代码中data文件夹下建立VISDroneIR文件夹存放红外数据集，分为images（存放红外图片）、labels（存放转换完的label，暂时为空）、xml（存放原有的xml格式标签）三个文件夹，其中分别再分train、test、val三个文件夹

VISDroneIR（数据集名字）
├── images      
       ├── train          
              ├── xx.jpg
       ├── test          
              ├── xx.jpg     
       ├── val         
              ├── xx.jpg 
├── labels      
       ├── train         
       ├── test              
       ├── val         
├── xml
       ├── train          
              ├── xx.xml
       ├── test          
              ├── xx.xml     
       ├── val         
              ├── xx.xml 

新建visdroneiryolov5.py文件对原本的label进行处理，在labels文件夹下分别生成txt存放yolov5格式标签

# -*- coding: utf-8 -*-
import xml.etree.ElementTree as ET
import os
from os import getcwd

sets = ['train', 'val', 'test']
classes = ["car", "truck", "bus", "van", "freight car"]  


def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = abs((box[0] + box[1]) / 2.0 - 1)
    y = abs((box[2] + box[3]) / 2.0 - 1)
    w = abs(box[1] - box[0])
    h = abs(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,image_set):
    in_file = open('/home/yolov5-master/data/VisDroneIR/xml/%s/%s.xml' % ((image_set),(image_id)), encoding='UTF-8')
    out_file = open('/home/yolov5-master/data/VisDroneIR/labels/%s/%s.txt' % ((image_set),(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'):
        '''if obj.find('difficult').text is not None:
            difficult = obj.find('difficult').text
            if int(difficult) == 1:
                continue'''#数据集里部分标签里没有difficult
        cls = obj.find('name').text
        if cls not in classes:
                continue
        cls_id = classes.index(cls)
        #数据集里有三种形式标点
        if obj.find('polygon') is not None:
            xmlbox = obj.find('polygon')
            xmin = int(min((xmlbox.find('x1').text), (xmlbox.find('x2').text), (xmlbox.find('x3').text), (xmlbox.find('x4').text)))
            xmax = int(max((xmlbox.find('x1').text), (xmlbox.find('x2').text), (xmlbox.find('x3').text), (xmlbox.find('x4').text)))
            ymin = int(min((xmlbox.find('y1').text), (xmlbox.find('y2').text), (xmlbox.find('y3').text), (xmlbox.find('y4').text)))
            ymax = int(max((xmlbox.find('y1').text), (xmlbox.find('y2').text), (xmlbox.find('y3').text), (xmlbox.find('y4').text)))
        if obj.find('bndbox') is not None:
            xmlbox = obj.find('bndbox')
            xmin = int(xmlbox.find('xmin').text)
            ymin = int(xmlbox.find('ymin').text)
            xmax = int(xmlbox.find('xmax').text)
            ymax = int(xmlbox.find('ymax').text)
        if obj.find('point') is not None:
            xmlbox = obj.find('point')
            xmin = int(xmlbox.find('x').text)
            ymin = int(xmlbox.find('y').text)
            xmax = int(xmlbox.find('x').text)
            ymax = int(xmlbox.find('y').text)
        b = (xmin, xmax, ymin, ymax)
        print(b)
        b1, b2, b3, b4 = b
        # 标注越界修正
        if b2 > w:
            b2 = w
        if b4 > h:
            b4 = h
        b = (b1, b2, b3, b4)
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')


wd = getcwd()
for image_set in sets:
    #如果不存在label文件夹即建立
    if not os.path.exists('/home/yolov5-master/data/VisDroneIR/labels/%s/' % (image_set)):
        os.makedirs('/home/yolov5-master/data/VisDroneIR/labels/%s/' % (image_set))
    path = '/home/yolov5-master/data/VisDroneIR/images/%s/' % (image_set)  # 文件夹地址
    files = os.listdir(path)  # 读入文件夹

    for image_id in files:
        index = image_id.rfind('.')
        image_id = image_id[:index]
        list_file = open('/home/yolov5-master/data/VisDroneIR/labels/%s/%s.txt' % ((image_set), (image_id)),'w')
        convert_annotation(image_id,image_set)
    list_file.close()

建立VisDroneIR.yaml文件对路径进行存储

path: /home/Newdisk/wyt/yolov5-master/data/VisDroneIR # dataset root dir
train: images/train # train images (relative to 'path') 
val: images/val # val images (relative to 'path') 
test: images/test # test images (optional)  

# Classes
nc: 5
names:
  0: car
  1: truck
  2: bus
  3: van
  4: freight car

 在官网下载所需预训练模型，并在对应文件里（这里以使用yolov5x为例，即yolov5x.yaml）对其中nc数进行更改

训练  

在train.py中对部分数值进行更改，如有需要可自行对部分参数自行更改（比如epoch、batchsize等）

注意：上一步下载的模型放在下面第一行代码所写的路径里

    parser.add_argument("--weights", type=str, default="/home/Newdisk/wyt/yolov5-master/.torch/models/yolov5x.pt")
    parser.add_argument("--cfg", type=str, default="/home/Newdisk/wyt/yolov5-master/models/yolov5x.yaml", help="model.yaml path")
    parser.add_argument("--data", type=str, default= "/home/Newdisk/wyt/yolov5-master/data/VisDroneIR.yaml", help="dataset.yaml path")

运行train.py 

如果遇到大量图片没有被读入的情况，可以检查txt文件数值是否归一化、归一化后是否出现负数、路径里是否有中文 

验证

在val.py中对部分数值进行更改，如有需要可自行对部分参数自行更改（比如epoch、batchsize等）

    parser.add_argument("--data", type=str, default="/home/yolov5-master/data/VisDroneIR.yaml", help="dataset.yaml path")
    parser.add_argument("--weights",  type=str, default="/home/yolov5-master/runs/train/exp34/weights/best.pt", help="model path(s)")
#这里exp选择最新训练出的就好

运行val.py 

测试 

在val.py中将task中的val改为test

    parser.add_argument("--task", default="test", help="train, val, test, speed or study")

运行val.py 

上述部分直接改代码操作也可直接通过传参进行修改

如需进行数据集划分可参考这篇文章