[Implementation of target detection algorithm] 2. Use yolov5 to train VisDrone2019 data set

Table of Contents

  • 0. Prepare
    • Yolov5 environment configuration
    • Download pretrained weights
    • VisDrone dataset preparation
      • Dataset download
      • Dataset processing
  • 1. Training
    • Modify the configuration file
      • Modify the VisDrone.yaml file
      • Modify the yolov5s.yaml file (you can modify the corresponding yaml file if training other yolo networks)
      • Modify the train.py file
    • start training

0. Preparation

Yolov5 environment configuration

Refer to the previous blog, as follows: Successfully deploy YOLOv5 to a remote server. Configure the environment of YOLOV5
[Implementation of target detection algorithm] – YOLOv5 is deployed on a remote server

Download pre-trained weights

In addition, you need to download the pre-training weight, and download the corresponding pre-training weight version according to your needs
In the official code of YOLOv5, four versions of the target detection network are given, namely four models of YOLOv5s, YOLOv5m, YOLOv5l, and YOLOv5x.
YOLOv5s has the smallest network, the fastest speed, and the lowest AP accuracy. However, if the detected scene is dominated by large targets and the pursuit of speed, then this model is also a good choice.
The other three networks continue to deepen and widen the network on this basis, and the AP accuracy is also continuously improved, but the speed consumption is also increasing.
The pre-training weight download links of the four networks are as follows: https://github.com/ultralytics/yolov5/releases/tag/v7.0
The page pulls down the bottom Assets

Find the required pre-trained weights

Create a new weights folder under the yolov5-master folder, and put in the weights required for download. I use yolov5s, so I downloaded the yolov5s.pt weight file

The downloaded weight file is placed under yolov5-maser/weights

VisDrone dataset preparation

Dataset download

Download from the github link: official link
Download the four VisDrone-DET datasets under Task1: Object Detectino in Images

It can also be downloaded from my Baidu network disk link:
Link: https://pan.baidu.com/s/18TYcaRdJI7jF55c0ZlI3LA?pwd=ef7n
Extraction code: ef7n

After downloading the zip file, use winscp to transfer the zip file to the remote server.
Go to the folder where the zip file is located on the server and use the unzip command to decompress the zip file.
Such as: unzip VisDrone2019-DET-val.zip

Dataset processing

The configuration file for training the visdrone data set has been integrated in the new version of yolov5, which includes the processing method of the data set, mainly for the generation of labels. You can create a new visdrone2yolov5.py file under yolov5-master.

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 number -1
                box = convert_box(img_size, tuple(map(int, row[:4])))
                lines.append(f"{<!-- -->cls} {<!-- -->' '.join(f'{<!-- -->x:.6f}' for x in box)}\\
")
                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('/home/huaxiang/yolov5/datasets/VisDrone2019') # Visdrone2019 folder directory under the datasets folder
#Convert
for d in 'VisDrone2019-DET-train', 'VisDrone2019-DET-val', 'VisDrone2019-DET-test-dev':
    visdrone2yolo(dir / d) # convert VisDrone annotations to YOLO labels

After the code is executed correctly, the labels folder will be newly generated in the three folders ‘VisDrone2019-DET-train’, ‘VisDrone2019-DET-val’, ‘VisDrone2019-DET-test-dev to store the VisDrone dataset Data label processed into YoloV5 format

1. Training

Modify configuration file

Modify the VisDrone.yaml file

Notepad or notepad + + open the VisDrone.yaml file under the yolov5/yolov5-master/data folder, and modify the path parameter to the path where the VisDrone2019 folder is located.

Modify the yolov5s.yaml file (you can modify the corresponding yaml file if training other yolo networks)

Notepad or notepad + + open the yolov5m.yaml file under the yolov5/yolov5-master/models folder, and modify the nc parameter to the category trained by the VisDrone2019 dataset, ie nc: 10.

Modify the train.py file

In the train.py file under yolov5/yolov5-master, several default parameters need to be modified
Mainly changed the default parameters of –weights, –cfg, –data, –epoch and –batch-size

def parse_opt(known=False):
    parser = argparse. ArgumentParser()
    parser.add_argument('--weights', type=str, default= 'weights/yolov5s.pt', help='initial weights path')
    parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml' , help='model.yaml path')
    parser.add_argument('--data', type=str, default= 'data/VisDrone.yaml', help='dataset.yaml path')

    # parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='initial weights path')
    # parser.add_argument('--cfg', type=str, default='', help='model.yaml path')
    # parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
    parser.add_argument('--hyp', type=str, default=ROOT / 'data/hyps/hyp.scratch-low.yaml', help='hyperparameters path')
    parser.add_argument('--epochs', type=int, default=10, help='total training epochs') # 100
    parser.add_argument('--batch-size', type=int, default=8, help='total batch size for all GPUs, -1 for autobatch') # 16
    parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='train, val image size (pixels) ')
    parser.add_argument('--rect', action='store_true', help='rectangular training')
    parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
    parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
    parser.add_argument('--noval', action='store_true', help='only validate final epoch')
    parser.add_argument('--noautoanchor', action='store_true', help='disable AutoAnchor')
    parser.add_argument('--noplots', action='store_true', help='save no plot files')
    parser.add_argument('--evolve', type=int, nargs='?', const=300, help='evolve hyperparameters for x generations')
    parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
    parser.add_argument('--cache', type=str, nargs='?', const='ram', help='image --cache ram/disk')
    parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--multi-scale', action='store_true', help='vary img-size + /- 50%%')
    parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class')
    parser.add_argument('--optimizer', type=str, choices=['SGD', 'Adam', 'AdamW'], default='SGD', help=' optimizer')
    parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
    parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)')
    parser.add_argument('--project', default=ROOT / 'runs/train', help='save to project/name')
    parser.add_argument('--name', default='exp', help='save to project/name')
    parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
    parser.add_argument('--quad', action='store_true', help='quad dataloader')
    parser.add_argument('--cos-lr', action='store_true', help='cosine LR scheduler')
    parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')
    parser.add_argument('--patience', type=int, default=100, help='EarlyStopping patience (epochs without improvement)')
    parser.add_argument('--freeze', nargs=' + ', type=int, default=[0], help='Freeze layers: backbone=10, first3=0 1 2')
    parser.add_argument('--save-period', type=int, default=-1, help='Save checkpoint every x epochs (disabled if < 1)')
    parser.add_argument('--seed', type=int, default=0, help='Global training seed')
    parser.add_argument('--local_rank', type=int, default=-1, help='Automatic DDP Multi-GPU argument, do not modify')

    # Logger arguments
    parser.add_argument('--entity', default=None, help='Entity')
    parser.add_argument('--upload_dataset', nargs='?', const=True, default=False, help='Upload data, "val" option')
    parser.add_argument('--bbox_interval', type=int, default=-1, help='Set bounding-box image logging interval')
    parser.add_argument('--artifact_alias', type=str, default='latest', help='Version of dataset artifact to use')

    return parser.parse_known_args()[0] if known else parser.parse_args()

Start training

Run the train.py file to start training.
During the training process, process files and training models will be generated, which will be saved in the runs/ folder