使用yolov5训练数据集并利用模型实现游戏脚本

环境篇

在正式训练之前，需要运行detect.py文件（按照默认设置）。如果运行出错，说明你的环境没有配置好。需要保证以下三点：1.按照requirments.txt文件的要求去安装库；2.tensorflow的版本与安装的cuda、cudnn的版本匹配；3.安装nvidia驱动程序。

对于第一点，可以直接在命令行执行：

pip install -r requirements.txt

通过pip list查看安装的库是否是需要的；如果不是，需要通过pip uninstall卸载之后再重新安装。

对于第二点，可以在tensorflow的管网查看：https://tensorflow.google.cn/install/source_windows?hl=zh-cn（注意这里是window系统，其他系统可在左边菜单栏选择）

tensorflow：直接pip安装对应的版本。
cuda下载链接：https://developer.nvidia.com/cuda-toolkit-archive
cudnn下载链接：https://developer.nvidia.com/cudnn-archive

对于第三点，尽量选择与自己电脑显卡匹配的驱动程序，不匹配会导致开机黑屏。
驱动程序下载链接：https://www.nvidia.cn/geforce/drivers/

训练篇

训练自定义数据集（官网）：https://docs.ultralytics.com/yolov5/tutorials/train_custom_data/#22-create-labels

实际上，不需要完全按照官方的目录结构去设置。训练需要的数据主要有两个：图片和标注文件。因此只需要在yolov5的同级目录下新建datasets文件夹（包含images和labels）即可。

参数说明

• weights：权重文件（模型就是基于这个文件运行的）
• data：配置文件（yaml文件）
• epochs：训练次数（总共要训练多少轮）
• batch-size：所有gpu的总批处理大小（设置不合理会报错）
• imgsz：调整图片大小（不要动，640x640 是 YOLOv5 的默认值）
• device：使用的设备，是cpu还是gpu（如果是gpu还要传入编号）
• project：保存训练数据的根目录
• name：存放训练结果的目录
• exist-ok：是否要新建一个exp来保存训练结果
• source：图片目录
• line-thickness：矩形边框粗细
• hide-labels：是否隐藏labels
• hide-conf：是否隐藏confidence

​注意：detect.py文件会在source对应的目录下找image，而不是在coco128.yaml文件中设置的目录下找image。

OK，在了解上述参数后，我们开始训练模型。
第一步，创建上述的datasets文件夹（包含images和labels）；
第二步，修改coco128.yaml文件（当然也可以自己新建一个yaml文件，只要修改data参数即可）：

1. 修改path、train、val

path:…/datasets
train: images
val: images

这里是否存在疑惑？不用怀疑，train和val都是images。

2. 修改name（标注时使用的label名）

Names:
0: HighLight

第三步，在train.py文件的parse_opt函数中修改参数（也可以在执行train.py文件的时候再传），主要修改weights、data、epochs、batch-size这几个参数。

第四步，执行train.py文件（开始训练）

python yolov5/train.py

这里有必要插入一个很重要的知识点——相对路径。就以上述的…/datasets为例，很多人会错误的认为…/datasets就是coco128.yaml文件上级目录下的datasets文件夹，这么认为也不完全错；我们要看使用coco128.yaml文件的文件是哪个，这里是train.py文件，那么…/datasets对应的就应该是train.py文件上级目录下的datasets文件夹。

其实相对路径的主要问题并不在于以上描述，而在于导包的时候。导包使用的就是相对路径，如果你去导入一个<导入了自定义包>的包，显示import错误，就是由于相对路径导致的，如果你不想处理，可以将两个文件放在同一个目录下；处理的方式也不难，设置根目录，以根目录为基础来导包。这里不好演示，有兴趣的小伙伴可以自己百度一下。

关于labels，官方的解释是：自动去images的同级目录找。也就是说，在images同级目录中有一个labels即可。但也有可能报错，这个时候需要在yolov5/utils/dataloders.py文件找到img2label_paths函数（yolov5的底层就是通过这个函数找到labels文件的），实现也很简单，就是将images路径中的‘images’换成‘labels’。

标注平台：https://www.makesense.ai/
关于make sense的具体操作这里不多赘述，注意最后导出的格式是yolo format。

另外，这里补充一下batch-size设置不合理会出现的问题：
1.报错：‘页面太小’
参考：https://blog.csdn.net/weixin_51697369/article/details/120101292
2.import错误（当时忘记截图，模糊记得是import错误）
将batch-size改小之后就没有再报这个错误了。

脚本篇

‘大家来找茬’应该都玩过吧？这个游戏的玩法很简单——找出两张图片的不同点。

那么要怎么通过代码来实现上述操作呢？实际上也非常简单，只需要一个函数：ImageChops.difference(leftImage, rightImage)，该函数会根据色差返回两张图片的‘差异图片’（色差越大的区域亮度就会越高）。这里的‘差异图片’就是我们要用来训练模型的原始数据，然后再通过训练好的模型来定位不同点。

Over…

关于窗口操作可以查看：https://blog.csdn.net/qq_57155669/article/details/140765599

改写dectect.py文件

原始的detect.py文件会返回一张带有标注信息的图片，但是我们想要的并不是这张图片，而是矩形区域的中心坐标。因此就需要改写detect.py文件。

要改写detect.py文件肯定要过一遍源码，好在源码中都给了注释，根据注释我们很快就能定位到生成矩形区域的位置（Process predictions）。

# Process predictions
for i, det in enumerate(pred):  # per image
    seen += 1
    if webcam:  # batch_size >= 1
        p, im0, frame = path[i], im0s[i].copy(), dataset.count
        s += f'{i}: '
    else:
        p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)

    p = Path(p)  # to Path
    save_path = str(save_dir / p.name)  # im.jpg
    txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}')  # im.txt
    s += '%gx%g ' % im.shape[2:]  # print string
    gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
    imc = im0.copy() if save_crop else im0  # for save_crop
    annotator = Annotator(im0, line_width=line_thickness, example=str(names))
    if len(det):
        # Rescale boxes from img_size to im0 size
        det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()

        # Print results
        for c in det[:, 5].unique():
            n = (det[:, 5] == c).sum()  # detections per class
            s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string

        # Write results
        for *xyxy, conf, cls in reversed(det):
            if save_txt:  # Write to file
                xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh
                line = (cls, *xywh, conf) if save_conf else (cls, *xywh)  # label format
                with open(f'{txt_path}.txt', 'a') as f:
                    f.write(('%g ' * len(line)).rstrip() % line + '\n')

            if save_img or save_crop or view_img:  # Add bbox to image
                c = int(cls)  # integer class
                label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
                annotator.box_label(xyxy, label, color=colors(c, True))
            if save_crop:
                save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)

annotator.box_label(xyxy, label, color=colors(c, True))

上述代码就是在图片中绘制矩形框，显然xyxy就是矩形框对应的坐标。

查看上面的代码可以发现有这样一条语句：

xywh=(xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()

根据上图，xy就是我们需要的中心坐标，然后我们要做的就是return。

针对我们的需求，detect.py文件中大部分内容都可以去除，这里就不做展示了，有兴趣的小伙伴可以查看脚本源码。

然后再将训练好的权重文件best.pt和detect.py文件及其依赖文件打包，通过调用detect.py文件中的run函数就可以使用模型生成目标数据了。