Yolov5车牌识别实验

最新推荐文章于 2024-05-09 18:59:42 发布

Ross_Han

最新推荐文章于 2024-05-09 18:59:42 发布

阅读量414

点赞数

分类专栏： AI图像识别文章标签： YOLO 机器学习

本文链接：https://blog.csdn.net/qq_37424837/article/details/133721794

版权

AI图像识别专栏收录该内容

2 篇文章 0 订阅

订阅专栏

Yolov5 车牌识别实验

Yolov5 车牌识别实验

Yolov5 车牌识别实验

模型

模型来源、数据集来源

采用 ultralytics/yolov5: YOLOv5 模型

数据集来源于中科大的CCPD2019数据集。该数据集主要采集于合肥市停车场，采集时间为上午7:30到晚上10:00，停车场采集人员手持Android POS机对停车场的车辆拍照进行数据采集。所拍摄的车牌照片涉及多种复杂环境，包括模糊、倾斜、雨天、雪天等。CCPD2019数据集包含了25万多幅中国城市车牌图像和车牌检测与识别信息的标注。其中以安徽省车牌为主，但因为此次实验不涉及车牌内容识别，故数据集的不完整并不影响实验结果。

CCPD2019 数据集处理

数据集的处理通过脚本 devide.py 将部分数据集图片分为训练集(train)和验证集(validation)。采用对CCPD2019数据集的各文件的文件名中车牌的位置信息进行提取，通过 label_train.py 和 label_val.py 将文件名中的车牌定位点(4个)提取，并通过脚本写入xml文件中，又通过 trans.py 将xml文件转格式为Yolov5可识别的txt文件，导入Yolov5模型中。

实验

依赖库的安装

从github上clone的Yolov5_master中包含有requirements.txt，在python下安装此依赖：

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple --ignore-installed

注：在anaconda下的python虚拟环境，pip3指令以pip指令代替。

下载预训练模型

此次实验选择了Yolov5s，需要的参数更少，在4060laptop显卡上表现良好，训练和验证的速度很快，实验时间15min左右。

yaml文件

通过yaml文件向训练模型导入数据集，文件内容如下：

# train and val datasets (image directory or *.txt file with image paths)  
path: D:\Program\Yolov5\yolo_card\images
train: images\train
val: images\val

# number of classes  
nc: 1

# class names  
names: ['license_card']

code

devide.py

import os
import random

import shutil
from shutil import copy2
trainfiles = os.listdir(r"D:\Program\Yolov5\ccpd_base") # images address
num_train = len(trainfiles)
print("num_train: " + str(num_train))
index_list = list(range(num_train))
print(index_list)
random.shuffle(index_list)  # 打乱顺序
num = 0
trainDir = r"D:\Program\Yolov5\ccpd\train"
validDir = r"D:\Program\Yolov5\ccpd\validation"
for i in index_list:
    fileName = os.path.join(r"D:\Program\Yolov5\ccpd_base", trainfiles[i])  # 图片地址 = 文件夹名 + 文件名
    if num < num_train*0.7:
        print(str(fileName))
        copy2(fileName, trainDir)
    else:
        print(str(fileName))
        copy2(fileName, validDir)
    num += 1

label_train.py

import os
import re
import cv2
root_path = r'D:\Program\Yolov5\yolo_card\images/'
file_name = os.listdir(root_path)
class_name = 'plate'
xml_dir =r'D:\Program\Yolov5\yolo_card\labels/'  # 保存xml的目录
a= 0
for i in file_name:
    if os.path.exists((xml_dir + i)):
        print('已经存在文件%s'%(xml_dir + i))
    else:os.mkdir((xml_dir + i))
    file_name1 = os.listdir(os.path.join(root_path,i))
    for ii in file_name1:
        if os.path.exists(xml_dir + i + '/'+ ii.split('.')[0]+'.xml'):
            continue
        #print(root_path+i+'/' + ii)
        img = cv2.imread((root_path+i+'/'+ii))
        print((root_path+i+'/'+ii))
        if img is None:
            print('文件%s读取失败'%ii)
            continue

        height = img.shape[0]
        width = img.shape[1]
        depth = img.shape[2]
        point = ii.split('.')[0].split('-')[3]
        #Xmin = point.split('_')[2].split('&')[0]
        num  = re.findall('\d+\d*',point)  #  正则表达式 从字符串中提取数值
        Xmin =min(num[0::2])  #　list[start:stop:step]
        Ymin = min(num[1::2])
        Xmax = max(num[0::2])
        Ymax = max(num[1::2])
        #print(ii.split('&'))
        fname = ii.split('&')[0] + '&amp;'+ii.split('&')[1]+ '&amp;'+ii.split('&')[2]+ '&amp;'+ii.split('&')[3]+ '&amp;'+ii.split('&')[4]+ '&amp;'+ii.split('&')[5]+ '&amp;'+ii.split('&')[1]+ '&amp;'+ii.split('&')[6]
        xml_str = "<annotation>\n\t\
                <folder>"+ i+ "</folder>\n\t\
                <filename>" + fname + "</filename>\n\t\
                " + "<path>" + root_path + i + '/'+fname+ "</path>\n\t\
                <source>\n\t\t\
                <database>Unknown</database>\n\t\
                </source>\n\t\
                <size>\n\t\t\
                <width>" + str(width) + "</width>\n\t\t\
                <height>" + str(height) + "</height>\n\t\t\
                <depth>" + str(depth) + "</depth>\n\t\
                </size>\n\t\
                <segmented>0</segmented>"
        obj_str = "\n\t\
                    <object>\n\t\t\
                    <name>" + class_name + "</name>\n\t\t\
                    <pose>Unspecified</pose>\n\t\t\
                    <truncated>0</truncated>\n\t\t\
                    <difficult>0</difficult>\n\t\t\
                    <bndbox>\n\t\t\t\
                    <xmin>" + str(Xmin) + "</xmin>\n\t\t\t\
                    <ymin>" + str(Ymin) + "</ymin>\n\t\t\t\
                    <xmax>" + str(Xmax) + "</xmax>\n\t\t\t\
                    <ymax>" + str(Ymax) + "</ymax>\n\t\t\
                    </bndbox>\n\t\
                    </object>"
        xml_str += obj_str
        xml_str +="\n</annotation>\n"
        with open(xml_dir + i + '/'+ ii.split('.')[0]+'.xml','w') as f:
            f.write(xml_str)
            a += 1
            print('成功读写文件夹%s第%d'%(i,a))
            f.close()
print('end')

label_val.py

import os
import re
import cv2
root_path = r"D:\Program\Yolov5\yolo_card\images\val"
file_name = os.listdir(root_path)
class_name = 'plate'
xml_dir = r"D:\Program\Yolov5\yolo_card\labels\val"   # xml address
a = 0
for i in file_name:
    if os.path.exists((xml_dir + i)):
        print("already exist file%s"%(xml_dir + i))
    else:
        os.mkdir((xml_dir + i))
    file_name = os.listdir(os.path.join(root_path, i))
    for ii in file_name:
        if os.path.exists(xml_dir + i + "/" + ii.split(".")[0] + ".xml"):
            continue
            img = cv2.imread((root_path + i + "/" + ii))
            print((root_path + i + "/" + ii))
            if img is None:
                print("file%s read fault"%ii)
                continue

            height = img.shape[0]
            width = img.shape[1]
            depth = omg.shape[2]
            point = ii.split(".")[0].split("-")[3]
            num = re.findall('\d+\d*',ponit)
            Xmin = min(nun[0::2])
            Ymin = min(num[1::2])
            Xmax = max(num[0::2])
            Ymax = max(num[1::2])
            fname = ii.split('&')[0] + '&amp;' + ii.split('&')[1] + '&amp;' + ii.split('&')[2] + '&amp;' + \
                    ii.split('&')[3] + '&amp;' + ii.split('&')[4] + '&amp;' + ii.split('&')[5] + '&amp;' + \
                    ii.split('&')[1] + '&amp;' + ii.split('&')[6]
            xml_str = "<annotation>\n\t\
                            <folder>" + i + "</folder>\n\t\
                            <filename>" + fname + "</filename>\n\t\
                            " + "<path>" + root_path + i + '/' + fname + "</path>\n\t\
                            <source>\n\t\t\
                            <database>Unknown</database>\n\t\
                            </source>\n\t\
                            <size>\n\t\t\
                            <width>" + str(width) + "</width>\n\t\t\
                            <height>" + str(height) + "</height>\n\t\t\
                            <depth>" + str(depth) + "</depth>\n\t\
                            </size>\n\t\
                            <segmented>0</segmented>"
            obj_str = "\n\t\
                                <object>\n\t\t\
                                <name>" + class_name + "</name>\n\t\t\
                                <pose>Unspecified</pose>\n\t\t\
                                <truncated>0</truncated>\n\t\t\
                                <difficult>0</difficult>\n\t\t\
                                <bndbox>\n\t\t\t\
                                <xmin>" + str(Xmin) + "</xmin>\n\t\t\t\
                                <ymin>" + str(Ymin) + "</ymin>\n\t\t\t\
                                <xmax>" + str(Xmax) + "</xmax>\n\t\t\t\
                                <ymax>" + str(Ymax) + "</ymax>\n\t\t\
                                </bndbox>\n\t\
                                </object>"
            xml_str += obj_str
            xml_str += "\n</annotation>\n"
            with open(xml_dir + i + '/' + ii.split('.')[0] + '.xml', 'w') as f:
                f.write(xml_str)
                a += 1
                print('成功读写文件夹%s第%d' % (i, a))
                f.close()
        print('end')

trans.py

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


def convert(size, box):
    # size=(width, height)  b=(xmin, xmax, ymin, ymax)
    # x_center = (xmax+xmin)/2        y_center = (ymax+ymin)/2
    # x = x_center / width            y = y_center / height
    # w = (xmax-xmin) / width         h = (ymax-ymin) / height

    x_center = (box[0] + box[1]) / 2.0
    y_center = (box[2] + box[3]) / 2.0
    x = x_center / size[0]
    y = y_center / size[1]

    w = (box[1] - box[0]) / size[0]
    h = (box[3] - box[2]) / size[1]

    # print(x, y, w, h)
    return (x, y, w, h)


def convert_annotation(xml_files_path, save_txt_files_path, classes):
    xml_files = os.listdir(xml_files_path)
    # print(xml_files)
    for xml_name in xml_files:
        # print(xml_name)
        xml_file = os.path.join(xml_files_path, xml_name)
        out_txt_path = os.path.join(save_txt_files_path, xml_name.split('.')[0] + '.txt')
        out_txt_f = open(out_txt_path, 'w')
        tree = ET.parse(xml_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))
            # b=(xmin, xmax, ymin, ymax)
            # print(w, h, b)
            bb = convert((w, h), b)
            out_txt_f.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')


if __name__ == "__main__":
    classes = ['plate']
    xml_files1 = (r'D:/Program/Yolov5/yolo_card/labels/val_annotations')

    save_txt_files1 = (r'D:/Program/Yolov5/yolo_card/labels/val')

    convert_annotation(xml_files1, save_txt_files1, classes)