0%

YOLOv5目标识别使用方法

预备

训练数据集准备

官方推荐使用roboflow(Roboflow: Give your software the power to see objects in images and video)这个网站进行数据集的标注,该网站需要翻墙才能访问

roboflow使用

  1. 注册一个账号
  2. 添加Workspace

image-20220809211233745

  1. 创建Project

image-20220809211302483

  1. Project中上传图片

image-20220809211342256

image-20220809211627912

上传完成后点击右上角绿色的Finish Uploading

image-20220809211658335

  1. 分配任务

由于roboflow支持多人合作,所以要将图片标注任务分配给某个人,如果是自己一个人做,分配给自己即可。

image-20220809211731031

  1. 开始标注

点击左侧Annotate,选择ANNOTATING中的任务开始标注。

image-20220809211914894

image-20220809212031832

图片中可以标注多个类别。

注意,图中所有你需要分类的类别都要做好标注,例如上图中有多个rock图像,那么它们都应当被标注出来。

  1. 生成数据集

标注完成的图片会被放入Annotated目录下,点击右上角的Add to Dataset可以生成数据集。你可以设置测试集、验证集、训练集的比例。

image-20220809212337706

  1. 导出

点击左侧Generate准备导出数据集,可以自行决定是否需要预处理和数据增强。YOLO自带图片处理,训练集的图片可以是不同的大小,所以可以不用图片预处理。

image-20220809212512551

点击Versions,点击右上角的Export导出,选择YOLOv5。

image-20220809212601963

image-20220809212636940

训练模型

训练时使用开发者提供的预训练模型,预训练模型分为多个大小。越小的模型运算速度越快,精度越低;越大的模型运算速度越慢,精度越高。如果没有下载预训练模型,程序会自动下载yolov5s.pt

预训练模型同样也在github仓库中可以下载到。

image-20220809212925936

在项目目录下,有train.py文件,直接用它就可以训练了,train.py中参数接受代码为:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
def parse_opt(known=False):
parser = argparse.ArgumentParser()
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=300)
parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs, -1 for autobatch')
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='--cache images in "ram" (default) or "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')

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

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

参数很多,有兴趣可以一一阅读并尝试,在这里只给出最基础的。

建议新建一个批处理文件run_train.bat,这样每次运行不需要重新写参数,也看得更舒服。

1
2
3
4
5
6
python -m train
--batch-size 8 ^
--epochs 10 ^
--data dataset/data.yml ^
--weights yolov5s.pt ^
--name mymodel ^

data参数是数据集的配置文件,这个data.yml在roboflow导出时会自动生成,在这里个文件中你要注意trainval二者的路径,如果不是把数据集直接放在yolov5的根目录下,那么你需要修改这两个路径让它们指向正确的位置。

weights参数是加载的模型参数,如果你下载了预训练模型,那么就把预训练模型的地址填进去。这一项可以缺省,如果缺省会自动下载yolov5s.pt。注意,当你完成训练后,如果将来想继续在原有的基础上训练模型,那么要把weights改成你想要继续训练的那个模型的地址

name参数是训练结果的保存文件夹,默认保存到名为exp的文件夹中,你可以自定义它的名字。训练出来的模型将被保存在name文件夹下的weights文件夹中,分别是best.pt和last.pt——训练时效果最好的模型和最后一次更新的模型。

模型使用

如果你需要使用你训练好的模型进行预测,这里不推荐使用工程目录下的eval文件,我们可以自己实现。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
import numpy as np
import torch
def detect(self, img0):
'''
img0: raw img
return: pred [[x1, y1, x2, y2, p, cls_id], [...]] or None 如果什么都没检测出来
(x1, y1): 方框左上角坐标
(x2, y2): 方框右下角坐标
p: 目标属于这个分类的概率(神经网络认为目标有多像这个分类)
cls_id: 类别id,类别id对应的类别标签可以在data.yml配置文件中查看
'''
# 1. 将图像处理成适配网络的尺寸
from utils.augmentations import letterbox
# letterbox是yolov5的图像处理函数,yolov5可以输出各种尺寸不同的图像,它会自动调整
img1 = letterbox(img0)[0]
img1 = img1.transpose((2, 0, 1))[::-1] # HWC -> CHW, BGR -> RGB
img1 = np.ascontiguousarray(img1) # 转换为内存中连续存储的数组,加快运行速度

img = torch.from_numpy(img1)
img = img.float()
img /= 255
img = img.unsqueeze(0)

# 2. 预测
pred = model(img)[0]
# pred[0]是预测结果
# pred[1]是其他的参数

# 3. 去除重叠
from utils.general import non_max_suppression
# non_max_suppression是yolov5的内置函数,他用来消除部分重叠
# 对于同一个目标来说,你可以用好多个框将其划出,使用这个函数可以消除这些「明明是同一个目标却用多个方框框出来」的现象
pred = non_max_suppression(pred,
conf_thres=0.4, # 只标注可能性大于该值的目标
iou_thres=0.2 # iou越大越能接受重叠
)


from utils.general import scale_coords
# 4. 放大还原
for det in pred:
if det is not None and len(det): # det不为空而且长度不为0 # 一张图像中的所有目标进行放大
# scale_coords是yolov5中的内置函数,因为在输入神经网络时要求特定的尺寸,所以图片时经过缩放的,同时图片中相应的方框坐标也是0-1之间的浮点数
# 通过这个函数可以将神经网络输出的预测值映射为原图像上对应的像素坐标
det[:, 0:4] = scale_coords(
img.shape[2:], # 缩小后的图像大小
det[:, 0:4], # 预测的目标的坐标
img0.shape # 原始图像大小
).round()# 四舍五入

if pred[0] is not None:
return pred[0].cpu().detach().numpy() # x1, y1, x2, y2, p, cls_id
else:
return None

以上的代码与eval中的处理流程基本一致,优点是可以直接获得预测值,后续操作更加灵活。

后续你可以在图片上画上方框和标签:

1
2
3
4
5
6
7
8
9
10
11
pred = self.detect(img) # 这个img是你输入的img,注意img格式为HWC
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
for p in pred:
img = cv2.rectangle(img, (int(p[0]), int(p[1])), (int(p[2]), int(p[3])), (255,255,255), 10)
img = cv2.putText(img, classes_list[int(p[5])] + " " + str(int(p[4].item() * 100)) + "%", (int(p[0])+np.random.randint(20),int(p[1])), # classes_list是根据data.yml定义的类别索引
cv2.FONT_HERSHEY_SIMPLEX ,
1,
(240,237,133),
3,
cv2.LINE_AA)
cv2.imwrite("img.jpg", img) # 图片会被保存到本地

效果如下:

19

如果是要做视频内容的预测也是一样的,把每一帧作为图像输入,获得预测结果后在每一帧图像上绘制方框,并将处理后的图像显示到UI界面上即可。

参考资料

yolov5/README.md at master · ultralytics/yolov5 (github.com)
yolov5预测 - 简书 (jianshu.com)
Detecting objects with YOLOv5, OpenCV, Python and C++ | by Luiz doleron | MLearning.ai | Medium