杀鸡用牛刀！通过YOLOv5识别斗地主牌型，做一个AI记牌器，再也不用担心如何算牌赢豆了！

YunLord

会魔术的程序员

前言

最近在网上看到一个很火的开源项目《斗地主老是输？一起用Python做个AI出牌器，欢乐豆蹭蹭涨！》，作者Dragon少年也分享了源代码，确实很好用，但是在扑克牌牌型识别的时候由于使用的是特定模板匹配，所以只适用于某游戏大厅的欢乐斗地主，而无法识别某手游助手的欢乐斗地主。

可以看到两者之间扑克牌牌型确实存在差异，所以通过特定模板匹配肯定无法识别出来。于是考虑通过神经网络进行扑克牌型识别，从而可以针对不同平台的斗地主进行扑克牌牌型识别，通用性更强。

由于该过程分为两个过程，一是定位到扑克牌数字或者字母位置，二是对相应位置上的图像进行牌型分类。经过一系列比较与试验，最终选择了YOLOv5进行扑克牌目标检测与识别。

一、YOLOv5环境配置

本教程所用环境：代码版本V3.0，源码下载地址： GitHub - ultralytics/yolov5: YOLOv5 in PyTorch > ONNX > CoreML > TFLite

通过git clone https://github.com/ultralytics/yolov5.git，将YOLOv5源码下载到本地，然后配置虚拟环境：conda create -n yolov5 python==3.7，在YOLOv5中尽量使用python3.7。创建好虚拟环境，并通过pip install -r requirements.txt安装依赖包。

项目的测试平台为：

操作系统：windows10
IDE：Spyder
python版本：anaconda Pyhon3.7
pytorch版本：torch 1.9.1
cuda版本：10.1
显卡：RTX 2080 Ti

二、制作自己的扑克牌数据

1.采集扑克牌数据

分为三类数据：手牌、底牌、出牌。

由于牌型特征明显，所以仅需要针对每一类采集少量数据进行标注，就能取得较好效果，采集结果如下图所示：

2.通过labelimg进行数据标注

由于第一次安装labelimg时，其中安装的pyqt5库与Spyder发生冲突，导致无法打开Spyder界面，折腾了很久才弄好。

所以建议新建一个虚拟环境，避免发生冲突。

（1）在Anaconda Prompt建立建立一个安装labelImg的python虚拟环境；

conda create -n LabelIMG python==3.7

（2）然后激活该环境

activate LabelIMG

（3）打开标注工具LabelImg

我们要标定数据的时候，一般是已经制定好了要标注物体的类别，这样在打开LabelImg之后，我们只要框定目标之后，选择要对应的标签即可。

labelimg cards predefined_classes.txt

其中cards为存放采集数据的文件夹，predefined_classes.txt为存放标注物体的类别的文本文件。

具体labelimg使用可以参考目标检测使用LabelImg标注VOC数据格式和YOLO数据格式——LabelImg使用详细教程。

PS：注意直接选择存储为YOLO格式，即标注数据最终保存为txt格式。

3. 划分训练集和验证集

数据集标注好之后，存放如下目录格式：

pokercard
├── images
│ ├── train # 训练集图片，这里我只列举几张示例
│ │ ├── 1.png
│ │ ├── 2.png
│ │ └── 3.png
│ └── val # 验证集图片
│ ├── t1.png
│ ├── t2.png
│ └── t3.png
└── labels
├── train # 训练集的标签文件
│ ├── 1.txt
│ ├── 2.txt
│ └── 3.txt
└── val2 # 验证集的标签文件
├── t1.txt
├── t2.txt
└── t3.txt

pokercard：存放数据的目录，该目录与YOLOv5目录同一级

images：目录下存放的是图片，包含训练集和验证集图片

labels：目录下存放的是标签文件，包含训练集和验证集图片对应的标签文件

三、训练扑克牌数据集

1.修改数据和模型配置文件

1.1 修改数据配置文件

原先的配置文件为：./yolov5/data/coco128.yaml。

我们把将文件复制一份，重命名为pokercard.yaml。

然后在pokercard.yaml中需要修改3处内容：

（1）训练集和验证集图片的路径

train: ../pokercard/images/train 
val: ../pokercard/images/val

（2）修改类别数nc

nc: 15  # number of classes

（3）修改类别列表，把类别修改为自己的类别

names: [ 'A','2','3', '4', '5','6', '7','8',  '9', '10','J','Q','K','xjk','djk']  # class names

修改后的pokercard.yaml完整配置内容如下：

# YOLOv5   by Ultralytics Ultralytics, licensed under GNU GPL v3.0
# COCO128 dataset coco128 (first 128 images from COCO train2017)
# Example usage: python train.py --data coco128.yaml
# parent
# ├── yolov5
# └── datasets
#     └── coco128  ← downloads here
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
# path: ../datasets/coco128  # dataset root dir
# path:
train: ../pokercard/images/train # train images (relative to 'path') 128 images
val: ../pokercard/images/val  # val images (relative to 'path') 128 images
test:  # test images (optional)