【3分钟速览】YOLOv5扑克牌点数识别器技术解析

前言

其实年初的时候，我也跟着别人的源码，用实现过扑克牌的目标检测。虽然也通过博文的方式记录了，但是那个项目使用的 TF 版本比较旧，自身对 TF 并不熟。后期如果说要升级或修改估计够呛，知道最近看到 YOLO 对目标检测的实现方法，发现这个更方便快捷。

于是决定用 YOLOv5.6 来重新实现一下扑克牌的点数识别，本文章主要是个人记录，也面向刚刚接触 YOLO 的同学阅读。所以以下会从数据标注，归一化处理到模型训练的实战内容，也包括各种踩坑经历和处理方法，最后对 detect.py 的重写，完成目标对象的坐标和识别标记输出。

YOLO 简介

YOLO（You Only Look Once）是一种基于深度神经网络的对象识别和定位算法，其最大的特点是运行速度很快，可以用于实时系统。现在 YOLO 已经发展到 v8 版本，每个版本都在原有基础上改进演化的。选择 V5 是因为后面几个新版发行并不久，v5 相关资料和文章较多，踩坑起来也好搜索。其中每个版本里又有按场景划分几个特定场景版本，比如：

YOLOv5s：最小版本，速度最快，但检测性能最低。
YOLOv5m：中等版本，速度比 s 慢一些，但检测性能更好。
YOLOv5l：大型版本，速度比 m 慢，但检测性能更好。
YOLOv5x：最大版本，速度最慢，但是检测性能较好。
YOLOv5n6：中等版本，速度比 m 快一些，但检测性能更好。

总之，根据具体需求和应用场景，可以选择合适的 YOLOv5 版本以满足性能、速度和准确性的要求，我这里以 YOLOv5s 为例。然后在项目开始前呢，我们可以先去 https://github.com/ultralytics/yolov5 下载 demo 代码和权重文件，比如 yolov5s.pt，下载完后就要准备数据集，也就是标记的图片。

数据集

图片的话，看你要识别的目标是啥，可以拍摄整理或去网上下载。为了让训练的模型稍微有点准确，较好准备一两百张，我这里就直接从原来 TF 项目那里拿过来。一共有 363 张扑克牌图片并且都已经用 VOC 格式标注好了，不过为了大家了解，这里还是会介绍如何标注，以及标注工具的基本使用。

LabekImg 使用

操作界面

快捷使用

1 、W 是进行标注 2 、A 是切换到上一张图片 3 、D 是切换到下一张图片

LabelImg 保存格式

1. PascalVOC 默认，xml格式 2. YOLO text可以直接拿来训练不用转换

实战

创建目录

在根目录下新建 datasets 用于放置训练的数据集，由于用同一 yolo 可以创建多个检测目标，所以在 datasets 先创建一个项目名目录。然后再在项目名下可创建 VOC 和 YOLO 两种格式的目录，虽然 labelImg 标注后可以直接保存 YOLO 格式也就是归一化后的 text，但是毕竟三百多张图的标注，我这里依然使用以前 TF 的 VOC 的标注数据。如果需要这些数据集的，可以留言，到时我再放在评论区。

依赖安装

依赖直接安装项目根目录下的 .txt 就可以，但是要严格按照里面的版本。因为我用的是以前 conda 创建的环境，可以包都装过，所以一开始就看到已经有的就没装，导致后面运行的时候很多问题都是包的版本问题。所以较好的方法就是用 conda 创建一个新环境，然后直接 pip -r .txt 是较好的。

归一化和数据集划分

如果 LabelImg 标注后保存的是 YOLO 格式，归一化过程就可以略过了。前面我也说了是用的以前扑克牌的 VOC 格式，所以需要对 xml 里节点数据抽取坐标和标记然后转换成 yolo 的 txt 格式。以下代码是我从其他博主那薅过来的，主要就是归一化和划分训练数据集和校验数据集。特别主要是 classes 部分要修改成自己的目标分类，其他的可以不用改动，当然只是划分的话可以适当删减。

import xml.etree.ElementTree as ETimport pickleimport os from os import listdir, getcwdfrom os . path import joinimport random from shutil import copyfileclasses = [ "nine" , "ten" , "jack" , "queen" , "king" , "ace" ]TRAIN_RATIO = 80 def clear_hidden_files( path ): dir_list = os .listdir( path ) abspath = "" for i in dir_list: abspath = os . path .join( os . path .abspath( path ), i) if os . path .isfile(abspath): if i.startswith( "._" ): os . remove (abspath) else : clear_hidden_files(abspath)def convert(size, box): dw = 1. / size[ 0 ] dh = 1. / size[ 1 ] x = (box[ 0 ] + box[ 1 ]) / 2.0 y = (box[ 2 ] + box[ 3 ]) / 2.0 w = box[ 1 ] - box[ 0 ] h = box[ 3 ] - box[ 2 ] x = x * dw w = w * dw y = y * dh h = h * dh return (x, y, w, h)def convert_annotation(image_id): in_file = open ( Poker/VOC/Annotations/%s.xml % image_id) out_file = open ( Poker/VOC/Labels/%s.txt % image_id, w ) tree = ET.parse(in_file) root = tree.getroot() size = root. find ( size ) w = int(size. find ( width ).text) h = int(size. find ( height ).text) for obj in root.iter( object ): difficult = obj. find ( difficult ).text cls = obj. find ( ame ).text if cls not in classes or int(difficult) == 1 : continue cls_id = classes.index(cls) xmlbox = obj. find ( ndbox ) b = (float(xmlbox. find ( xmin ).text), float(xmlbox. find ( xmax ).text), float(xmlbox. find ( ymin ).text), float(xmlbox. find ( ymax ).text)) bb = convert((w, h), b) out_file. write (str(cls_id) + " " + " " .join([str(a) for a in bb]) + ) in_file. close () out_file. close ()# 开始执行 ##wd = os .getcwd()data_base_dir = os . path .join(wd, "Poker/" ) if not os . path .isdir(data_base_dir): os .mkdir(data_base_dir)work_sapce_dir = os . path .join(data_base_dir, "VOC/" ) if not os . path .isdir(work_sapce_dir): os .mkdir(work_sapce_dir)annotation_dir = os . path .join(work_sapce_dir, "Annotations/" ) if not os . path .isdir(annotation_dir): os .mkdir(annotation_dir)image_dir = os . path .join(work_sapce_dir, "Images/" ) if not os . path .isdir(image_dir): os .mkdir(image_dir)yolo_labels_dir = os . path .join(work_sapce_dir, "Labels/" ) if not os . path .isdir(yolo_labels_dir): os .mkdir(yolo_labels_dir)yolov5_images_dir = os . path .join(data_base_dir, "images/" ) if not os . path .isdir(yolov5_images_dir): os .mkdir(yolov5_images_dir)yolov5_labels_dir = os . path .join(data_base_dir, "labels/" ) if not os . path .isdir(yolov5_labels_dir): os .mkdir(yolov5_labels_dir)yolov5_images_train_dir = os . path .join(yolov5_images_dir, "train/" ) if not os . path .isdir(yolov5_images_train_dir): os .mkdir(yolov5_images_train_dir)yolov5_images_test_dir = os . path .join(yolov5_images_dir, "val/" ) if not os . path .isdir(yolov5_images_test_dir): os .mkdir(yolov5_images_test_dir)yolov5_labels_train_dir = os . path .join(yolov5_labels_dir, "train/" ) if not os . path .isdir(yolov5_labels_train_dir): os .mkdir(yolov5_labels_train_dir)yolov5_labels_test_dir = os . path .join(yolov5_labels_dir, "val/" ) if not os . path .isdir(yolov5_labels_test_dir): os .mkdir(yolov5_labels_test_dir)train_file = open ( os . path .join(wd, "yolov5_train.txt" ), w )test_file = open ( os . path .join(wd, "yolov5_val.txt" ), w )train_file. close ()test_file. close ()train_file = open ( os . path .join(wd, "yolov5_train.txt" ), a )test_file = open ( os . path .join(wd, "yolov5_val.txt" ), a )list_imgs = os .listdir(image_dir) # list image filesprob = random .randint( 1 , 100 ) print ( "Probability: %d" % prob) for i in range( 0 , len (list_imgs)): path = os . path .join(image_dir, list_imgs[i]) if os . path .isfile( path ): image_path = image_dir + list_imgs[i] voc_path = list_imgs[i] (nameWithoutExtention, extention) = os . path .splitext( os . path .basename(image_path)) (voc_nameWithoutExtention, voc_extention) = os . path .splitext( os . path .basename(voc_path)) annotation_name = nameWithoutExtention + .xml annotation_path = os . path .join(annotation_dir, annotation_name) label_name = nameWithoutExtention + .txt label_path = os . path .join(yolo_labels_dir, label_name) prob = random .randint( 1 , 100 ) print ( "Probability: %d" % prob) if (prob TRAIN_RATIO): # train dataset if os . path .exists(annotation_path): train_file. write (image_path + ) convert_annotation(nameWithoutExtention) # convert label copyfile(image_path, yolov5_images_train_dir + voc_path) copyfile(label_path, yolov5_labels_train_dir + label_name) else : # test dataset if os . path .exists(annotation_path): test_file. write (image_path + ) convert_annotation(nameWithoutExtention) # convert label copyfile(image_path, yolov5_images_test_dir + voc_path) copyfile(label_path, yolov5_labels_test_dir + label_name)train_file. close ()test_file. close ()

上面代码放在哪个位置看你自己，由于有些目录是相对位置，所以如果不想改也跟我一样放在 datasets 中。执行完上面代码后会在 datasets 下生成 images 和 labels 两个目录，而且特别注意的是这两个目录名较好就是也是这样，比如在上面代码改了生成名，在后面训练的时候，会提示找不到 labels 目录。当然也可以修改 YOLO 提供源码的 dataset.py，看个人选择，而且 images 里的 train 和 val 需要配置到 data 的 yaml 中，这个马上就会说。

配置

数据集配置

该默认文件位于项目的 data/coco128.yaml，现在我们进行划分了数据集和自己的目标种类，当然要对其进行重新配置。较好的方法就是复制一份 coco128 然后修改名称，如下面 coco128_pocker.yaml 是我的配置示例。

train : D : 3 code 6 pytorchopencv_demo 05 _yolo_v5. 6 datasetsPokerimagestrain # train images (relative to path ) 128 images val : D : 3 code 6 pytorchopencv_demo 05 _yolo_v5. 6 datasetsPokerimagesval # val images (relative to path ) 128 images# Classes nc : 6 # number of classes names : [ "nine" , "ten" , "jack" , "queen" , "king" , "ace" ] # class names# Download script/URL (optional) download : https : //ultralytics.com/assets/coco128.zip

模型配置

该默认文件位于项目的 model 下，然后选择自己的场景，如我用的是 v5s 同样就复制一份重命名。这个文件主要是定义模型的基本机构和参数，比如分类数目 nc，模型深度倍数，每层通道数倍数，以及目标框的大小和比例的。在刚使用 YOLO 的话，我们就改个 nc 就可以，后面再研究这些参数的用法，下面是 .yaml 的示例。

# Parameters nc: 6 # number of classes depth_multiple: 0.33 # model depth multiple width_multiple: 0.50 # layer channel multiple anchors: - [10,13, 16 ,30, 33 ,23] # P3/8 - [30,61, 62 ,45, 59 ,119] # P4/16 - [116,90, 156 ,198, 373 ,326] # P5/32 # YOLOv5 v6.0 backbone backbone: # [from, number, module, args] [[-1, 1 , Conv, [64, 6 , 2 , 2 ]], # 0-P1/2 [-1, 1 , Conv, [128, 3 , 2 ]], # 1-P2/4 [-1, 3 , C3, [128]], [-1, 1 , Conv, [256, 3 , 2 ]], # 3-P3/8 [-1, 6 , C3, [256]], [-1, 1 , Conv, [512, 3 , 2 ]], # 5-P4/16 [-1, 9 , C3, [512]], [-1, 1 , Conv, [1024, 3 , 2 ]], # 7-P5/32 [-1, 3 , C3, [1024]], [-1, 1 , SPPF, [1024, 5 ]], # 9 ] # YOLOv5 v6.0 head head: [[-1, 1 , Conv, [512, 1 , 1 ]], [-1, 1 , nn.Upsample, [None, 2 , earest ]], [[-1, 6 ], 1 , Concat, [1]], # cat backbone P4 [-1, 3 , C3, [512, False ]], # 13 [-1, 1 , Conv, [256, 1 , 1 ]], [-1, 1 , nn.Upsample, [None, 2 , earest ]], [[-1, 4 ], 1 , Concat, [1]], # cat backbone P3 [-1, 3 , C3, [256, False ]], # 17 (P3/8-small) [-1, 1 , Conv, [256, 3 , 2 ]], [[-1, 14 ], 1 , Concat, [1]], # cat head P4 [-1, 3 , C3, [512, False ]], # 20 (P4/16-medium) [-1, 1 , Conv, [512, 3 , 2 ]], [[-1, 10 ], 1 , Concat, [1]], # cat head P5 [-1, 3 , C3, [1024, False ]], # 23 (P5/32-large) [[17, 20 , 23 ], 1 , Detect, [nc, anchors]], # Detect(P3, P4, P5) ]

权重文件

这个就是我们从 yolo 开源那里下载的 .pt，这里面有 80 个类目的检测，在基础识别时候可以直接用它。而训练自己的目标也需要基于他的权重再进行训练，所以可以在项目中创建一个，然后再将 .pt 放到其中。

训练

基础预测

在训练目标模型前，可以先试试以前的基础检测，这里我们可以随便拿一张图放到 data/images 下，然后执行下面命令，根目录下会生成 runs，而 runs 下的 detect 会根据每次执行的次数生成不同次的标注结果。

python detect.py --weights weights/yolov5s.pt -- source data/images/zidane.jpg

指令介绍

source : 需要进行检测的图片或视频的文件位置 weigths : 指的是训练好的网络模型，用来初始化网络权重cfg：为configuration的缩写，指的是网络结构，一般对应models文件夹下的xxx.yaml文件data：训练数据路径，一般为data文件夹下的xxx.yaml文件epochs：设置训练的轮数（自己电脑上一般建议先小一点，测试一下，看跑一轮要多久）batch-size：每次输出给神经网络的图片数，（需要根据自己电脑性能进行调整）img-size：用于分别设置训练集和测试集的大小。两个数字前者为训练集大小，后者为测试集大小image-weights：测试过程中，图像的那些测试地方不太好，对这些不太好的地方加权重device：训练网络的设备cpu还是gpumulti-scale：训练过程中对图片进行尺度变换 workers : 多线程训练label-smoothing：对标签进行平滑处理，防止过拟合

开始训练

以下是我踩过坑后执行的命令，因为当时我用 CPU 训练一轮要 7 分钟就改成了 GPU，加了 -- 0 后速度快了 4 倍。为 2，这么小是因为我的显卡 750Ti 显存只有 2G，就暂时用了这个数值。而加了 -- 0，也是因为有坑，这个后面再说。训练完成后，会在根目录的 runs/train 下生成每次运行的权重文件和损失准确示例图等，其中下的 best.pt 和 last.pt，分别表示本轮训练较好的权重和最新的权重。

python train.py --weights weights/yolov5s.pt --cfg models/yolov5s_pocker.yaml --data data/coco128_pocker.yaml --epochs 1 --batch-size 2 --multi-scale --device 0 --workers 0

开始检测

训练完成后，我们可以把本轮的 best.pt 放到根目录的下，以下是不同检测方式的命令，可以分别都试一下。

1 . 图片python detect.py --weights weights/best.pt -- data data /coco128_pocker.yaml --source data /images/cam_image16.jpg 2 . 视频python detect.py --weights weights/best.pt -- data data /coco128_pocker.yaml --source data /images/test.mov 3 . 使用cuda测试python detect.py --device 0 --weights weights/best.pt -- data data /coco128_pocker.yaml --source data /images/IMG_2681.JPG

detect.py 重写

因为在实际项目中，我们对图片和视频的目标检测，更多的可能是通过前端传递过来的图片或视频，模型给出的是预测的类目名称和目标的位置，而不是直接生成标注的图片和视频。所以我对 demo 提供的 detect 进行了简单的修改，修改结果就是前面放出的图片，不过还是采用命令方式，当然也可以直接改成接口，最后就是如果需要代码的可以评论区留言哈。

import warningswarnings.filterwarnings( "ignore" )import argparsefrom utils.datasets import *from utils.torch_utils import *from utils.augmentations import *from utils.general import *from models.common import DetectMultiBackendfrom utils.plots import Annotator, colors, save_one_boximport timeimport cv2import torchimport randomimport numpy as npdef detect(save_img=False): # 解析配置参数 source , weights, data, imgsz = opt.source, opt.weights, opt.data, opt.img_size # 初始化模型推理硬件 device = select_device(opt.device) model = DetectMultiBackend(weights, device=device, dnn=False, data=data) stride, names, pt, jit, onnx, engine = model.stride, model.names, model.pt, model.jit, model.onnx, model.engine imgsz = check_img_size(imgsz, s=stride) dataset = LoadImages( source , img_size=imgsz, stride=stride, auto=pt) dt, seen = [0.0, 0.0, 0.0], 0 for path, im, im0s, vid_cap, s in dataset: t1 = time_sync() # im = torch.from_numpy(im).to(device) im = torch.from_numpy(im.astype(np.float32)).to(device) im /= 255 # 0 - 255 to 0.0 - 1.0 if len(im.shape) == 3: im = im[None] # expand for batch dim t2 = time_sync() dt[0] += t2 - t1 dt[0] += t2 - t1 # Inference # visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False pred = model(im) t3 = time_sync() dt[1] += t3 - t2 # NMS pred = non_max_suppression(pred, 0.25, 0.45, None, False, max_det=1000) dt[2] += time_sync() - t3 for i, det in enumerate(pred): p, s, im0 = source , , im0s # print(im0_shape：, im0.shape) s += %gx%g % im.shape[2:] # print string gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh if det is not None and len(det): # 推理的图像分辨率转为原图分辨率：Rescale boxes from img_size to im0 size det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round() # Print results for c in det[:, -1].unique(): n = (det[:, -1] == c).sum() # detections per class s += %g %ss, % (n, names[int(c)]) # add to string # Write results output_dict_ = [] for *xyxy, conf, cls in det: x1, y1, x2, y2 = xyxy output_dict_.append(( float (x1), float (y1), float (x2), float (y2))) label = %s %.2f % (names[int(cls)], conf) print ( "---------------------------------------------------------------------" ) print ( "尺寸：" , im0.shape) print ( "坐标：" , ( float (x1), float (y1), float (x2), float (y2))) print ( "标识：" , label) # print("output_dict_ : ", output_dict_) if __name__ == __main__ : parser = argparse.ArgumentParser() parser.add_argument( --weights , type =str, default= weights/yolov5s.pt , help = model.pt path ) parser.add_argument( --data , type =str, default= data/coco128_pocker.yaml , help = dataset.yaml path ) parser.add_argument( --source , type =str, default= "./video/1.mp4" , help = source ) # file/folder, 0 for webcam parser.add_argument( --img-size , type =int, default=640, help = inference size (pixels) ) parser.add_argument( --conf-thres , type = float , default=0.31, help = object confidence threshold ) parser.add_argument( --iou-thres , type = float , default=0.45, help = IOU threshold for NMS ) parser.add_argument( --fourcc , type =str, default= mp4v , help = output video codec (verify ffmpeg support) ) parser.add_argument( --device , default= , help = cuda device, i.e. 0 or 0,1,2,3 or cpu ) parser.add_argument( --classes , nargs= + , type =int, help = filter by class ) parser.add_argument( --agnostic-nms , action= store_true , help = class-agnostic NMS ) parser.add_argument( --augment , default=False, help = augmented inference ) opt = parser.parse_args() print (opt) # 打印输入配置参数 with torch.no_grad(): detect(save_img=True)

前言

YOLO 简介

数据集

实战

相关文章