【目标检测算法实现之yolov5】二、使用yolov5在VisDrone2019数据集上训练并测试

0.准备

Yolov5环境配置

参考上一篇博客，如下：将YOLOv5成功部署到远程服务器上。配置YOLOV5的环境
【目标检测算法实现】——YOLOv5部署到远程服务器上

下载预训练权重

此外，还需下载预训练权重，根据自己需要下载相应的预训练权重版本
YOLOv5的官方代码中，给出了四种版本的目标检测网络，分别是YOLOv5s、YOLOv5m、YOLOv5l、YOLOv5x四个模型。
YOLOv5s网络最小，速度最快，AP精度也最低。但如果检测的大目标为主的场景，追求速度，那么这个模型也是一个很好的选择。
其他的三种网络，在此基础上不断的加深加宽网络，AP精度也不断提升，但速度的消耗也在不断增加。
其中四种网络的预训练权重下载链接如下：https://github.com/ultralytics/yolov5/releases/tag/v7.0
页面拉倒最下面的Assets

找到需要的预训练权重

在yolov5-master文件夹下新建一个weights文件夹，放入下载所需要的权重，笔者使用的是yolov5s，因此下载的是yolov5s.pt权重文件

下载好的权重文件放在yolov5-maser/weights下面

VisDrone数据集准备

数据集下载

github链接上下载：官方链接
下载Task1：Object Detectino in Images下面的四个VisDrone-DET dataset数据集

也可以在我的百度网盘链接里下载：
链接：https://pan.baidu.com/s/18TYcaRdJI7jF55c0ZlI3LA?pwd=ef7n
提取码：ef7n

下载好zip文件后，使用winscp将zip文件传输到远程服务器上。
在服务器上进入到zip文件所在的文件夹中使用unzip命令解压zip文件。
如： unzip VisDrone2019-DET-val.zip

数据集处理

新版本的yolov5中已经集成了训练visdrone数据集的配置文件，其中附带了数据集的处理方式，主要是labels的生成，可以在yolov5-master下面新建一个visdrone2yolov5.py文件。

from utils.general import download, os, Path
def visdrone2yolo(dir):from PIL import Imagefrom tqdm import tqdmdef convert_box(size, box):# Convert VisDrone box to YOLO xywh boxdw = 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 directorypbar = tqdm((dir / 'annotations').glob('*.txt'), desc=f'Converting {dir}')for f in pbar:img_size = Image.open((dir / 'images' / f.name).with_suffix('.jpg')).sizelines = []with open(f, 'r') as file:  # read annotation.txtfor row in [x.split(',') for x in file.read().strip().splitlines()]:if row[4] == '0':  # VisDrone 'ignored regions' class 0continuecls = int(row[5]) - 1  # 类别号-1box = 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('/home/huaxiang/yolov5/datasets/VisDrone2019')  # datasets文件夹下Visdrone2019文件夹目录
# Convert
for d in 'VisDrone2019-DET-train', 'VisDrone2019-DET-val', 'VisDrone2019-DET-test-dev':visdrone2yolo(dir / d)  # convert VisDrone annotations to YOLO labels

正确执行代码后，会在’VisDrone2019-DET-train’, ‘VisDrone2019-DET-val’, 'VisDrone2019-DET-test-dev三个文件夹内新生成labels文件夹，用以存放将VisDrone数据集处理成YoloV5格式后的数据标

1.训练

修改配置文件

修改VisDrone.yaml文件

记事本或notepad++打开yolov5/yolov5-master/data文件夹下的VisDrone.yaml文件，将其中path参数修改为VisDrone2019文件夹所在的路径。

修改yolov5s.yaml文件（若训练其他yolo网络修改相应的yaml文件即可）

记事本或notepad++打开yolov5/yolov5-master/models文件夹下的yolov5m.yaml文件，将其中nc参数修改为VisDrone2019数据集训练的类别，即nc：10。

修改train.py文件

在yolov5/yolov5-master下的train.py文件，需要修改几个default参数
主要更改了–weights、–cfg、–data、–epoch以及–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')  # 100parser.add_argument('--batch-size', type=int, default=8, help='total batch size for all GPUs, -1 for autobatch')  # 16parser.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 argumentsparser.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()

开始训练

运行train.py文件，即开始训练。
在训练的过程中会产生过程文件及训练模型，会保存在runs/文件夹中

2.验证

修改配置文件

在yolov5/yolov5-master下的val.py文件，需要修改几个default参数
主要更改了–weights、–data、–epoch以及–batch-size的默认参数

weights更改为自己训练好的权重
data修改为VisDrone数据集
batch-size我改为16

def parse_opt():parser = argparse.ArgumentParser()parser.add_argument('--data', type=str, default='data/VisDrone.yaml', help='dataset.yaml path')parser.add_argument('--weights', nargs='+', type=str, default='runs/train/exp9/weights/best.pt', help='model path(s)')# parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')# parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model path(s)')parser.add_argument('--batch-size', type=int, default=16, help='batch size') # 32parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='inference size (pixels)')parser.add_argument('--conf-thres', type=float, default=0.001, help='confidence threshold')parser.add_argument('--iou-thres', type=float, default=0.6, help='NMS IoU threshold')parser.add_argument('--max-det', type=int, default=300, help='maximum detections per image')parser.add_argument('--task', default='val', help='train, val, test, speed or study')parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)')parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset')parser.add_argument('--augment', action='store_true', help='augmented inference')parser.add_argument('--verbose', action='store_true', help='report mAP by class')parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')parser.add_argument('--save-hybrid', action='store_true', help='save label+prediction hybrid results to *.txt')parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')parser.add_argument('--save-json', action='store_true', help='save a COCO-JSON results file')parser.add_argument('--project', default=ROOT / 'runs/val', 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('--half', action='store_true', help='use FP16 half-precision inference')parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')opt = parser.parse_args()opt.data = check_yaml(opt.data)  # check YAMLopt.save_json |= opt.data.endswith('coco.yaml')opt.save_txt |= opt.save_hybridprint_args(vars(opt))return opt

若不想修改配置文件，可直接通过以下命令运行

python test.py --weights runs/train/exp9/weights/best.pt --data VisDrone.yaml

开始测试

将要运行的文件Script path 改为val.py

直接运行val.py文件即可。
我运行的结果如下：