【YOLOv5】实现扑克牌的点数识别

前言

       其实年初的时候，我也跟着别人的源码，用 Tensoflow 实现过扑克牌的目标检测。虽然也通过博文的方式记录了，但是那个项目使用的 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 格式标注好了，不过为了大家了解，这里还是会介绍如何标注，以及标注工具 labelImg 的基本使用。 

LabelImg 使用

操作界面

快捷使用

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

LabelImg 保存格式

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

实战

创建目录

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

依赖安装

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

归一化和数据集划分

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

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
import random
from shutil import copyfile

classes = ["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('name').text
        if cls not in classes or int(difficult) == 1:
            continue

        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        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]) + '\n')

    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 files
prob = 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 + '\n')
            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 + '\n')
            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()