[target detection] yolov5 Runtong visdrone data set

background

stay YOLOv5 Official 6.1 In the version , I find Data There are so many in the catalogue 5 Data sets , Including UAV aerial photo data set VisDrone And remote sensing data sets xView, This reflects that the official is also secretly making efforts in small target detection .
This blog will record how to use YOLOv5 Run through VisDrone Data sets .

I have uploaded the data set to the online disk , Readers who need it can download ：
https://pan.baidu.com/s/1UNQlZGHZlAZs412tbnpAxg?pwd=8888

Data set structure

The data set is divided into four files ：

• VisDrone2019-DET-train： Training set
• VisDrone2019-DET-val： Verification set
• VisDrone2019-DET-test-dev： Test set ( With label )
• VisDrone2019-DET-test-challenge： Test set ( Without a label ）

Unlabeled test sets are used for VisDrone2021 The challenge , If you don't participate in the competition, you don't need to use .

Dataset processing

and VOC The dataset is the same ,VisDrone The data annotation of is xml form , You need to convert it into YOLOv5 The required txt Format
In the official offering VisDrone.yaml below , Scripts for data processing have been provided , Make simple modifications based on it .
Create in the root directory visdrone2yolo.py：

from utils.general import download, os, Path


def visdrone2yolo(dir):
    from PIL import Image
    from tqdm import tqdm

    def convert_box(size, box):
        # Convert VisDrone box to YOLO xywh box
        dw = 1. / size[0]
        dh = 1. / size[1]
        return (box[0] + box[2] / 2) * dw, (box[1] + box[3] / 2) * dh, box[2] * dw, box[3] * dh

    (dir / 'labels').mkdir(parents=True, exist_ok=True)  # make labels directory
    pbar = tqdm((dir / 'annotations').glob('*.txt'), desc=f'Converting {
      dir}')
    for f in pbar:
        img_size = Image.open((dir / 'images' / f.name).with_suffix('.jpg')).size
        lines = []
        with open(f, 'r') as file:  # read annotation.txt
            for row in [x.split(',') for x in file.read().strip().splitlines()]:
                if row[4] == '0':  # VisDrone 'ignored regions' class 0
                    continue
                cls = int(row[5]) - 1  #  Category No -1
                box = convert_box(img_size, tuple(map(int, row[:4])))
                lines.append(f"{
      cls} {
      ' '.join(f'{
        x:.6f}' for x in box)}\n")
                with open(str(f).replace(os.sep + 'annotations' + os.sep, os.sep + 'labels' + os.sep), 'w') as fl:
                    fl.writelines(lines)  # write label.txt


dir = Path('D:/Desktop/Work/Dataset/VisDrone')  # datasets Under the folder Visdrone2019 Folder Directory 
# Convert
for d in 'VisDrone2019-DET-train', 'VisDrone2019-DET-val', 'VisDrone2019-DET-test-dev':
    visdrone2yolo(dir / d)  # convert VisDrone annotations to YOLO labels

After running , You can see that and... Are generated under each data set images Corresponding labels

Create your own data path file

stay data Create below mydata.yaml( It can also be modified directly VisDrone.yaml)
Type in the following ：

train: D:/Dataset/VisDrone/VisDrone2019-DET-train/images  # train images (relative to 'path') 6471 images
val: D:/Dataset/VisDrone/VisDrone2019-DET-val/images  # val images (relative to 'path') 548 images
test: D:/Dataset/VisDrone/VisDrone2019-DET-test-dev/images   # test images (optional) 1610 images

# Classes
nc: 10  # number of classes
names: [ 'pedestrian', 'people', 'bicycle', 'car', 'van', 'truck', 'tricycle', 'awning-tricycle', 'bus', 'motor' ]

notes ： Replace the path here with your own path .

Start training

stay train.py in , It is mainly necessary to modify several quantities shown in the red box below .
batch-size and workers The two parameters are flexibly selected according to your computer configuration .
It is worth noting that YOLOv5 5.0 The default size of the input image of version is 640, stay YOLOv5 6.1 In the version , The default size is changed to 960, This also reflects the official strategy to deal with small target detection .

After setting up , Just start training .

Video detection

After model training , function detect.py It can detect pictures or videos .
When using , Mainly modify the first three parameters , Meaning and train.py It's similar to .

in addition , For dense small targets , Data labels and confidence levels generated by default during output may block the target , Affect perception .
You can modify the following hide-labels and hide-conf Two attributes , In this way, labels and confidence levels can be hidden .

Enjoy yourself！