(本项目所有代码打包至我的资源中,大家可在我的文章底部选择下载)
目录
需求
通过车辆识别技术,识别视频中每个车辆及其车牌号,车辆应进行追踪,避免重复计数量。
实现效果
车牌识别
学习视频
使用 Python、Yolov8 和 EasyOCR 自动识别车牌 计算机视觉教程_哔哩哔哩_bilibili
大致思路
通过 opencv 将视频转换为帧,对帧应用车辆识别模型,并使用 model.track 或者 sort 追踪器进行追踪,给每个车辆一个唯一的 id ,然后进行车牌识别,对每一帧识别到的车牌,通过几何判断是否位于某个车辆中,是则将该车牌分配给该车辆,否则说明车牌识别错误或车辆识别错误,不作考虑。将车牌分配好车辆后,对车牌进行裁剪,将裁剪好的车牌使用 opencv 技术转换为灰度值图片,再设置阈值转换为阈值灰白图像,然后使用 easyocr 或者 Paddleocr 等文字识别技术,对阈值黑白图像进行字符识别,最终得到车牌信息,然后将所有信息,包括车辆、车牌 box 信息、置信度、车辆 id ,车牌号等信息放入 csv 表格中,方面查看测试结果。
代码实现
引入模块:
import numpy as np
from ultralytics import YOLO
import cv2
from util import *
# from sort.sort import *
# 跟踪器对象
# mot_tracker = Sort()
这里由于我的环境问题,下载不了使用 sort 的库,即这几个:
filterpy==1.4.5
scikit-image==0.17.2
lap==0.4.0
所以只能使用 yolo 自带的追踪器进行追踪
定义存储字典并解析视频:
# 存储所有信息
results = {}
# 加载模型
coco_model = YOLO('yolov8n.pt', task='detect')
license_plate_detector = YOLO('best.pt')
# 加载视频
cap = cv2.VideoCapture('./3.mp4')
通过 opencv 的 VideoCapture 函数加载视频,并加载模型,这里的车牌模型我使用的官方模型,车牌模型是自己训练的,大家下载后自取。
识别帧,为了方便测试,只取前十帧:
# 读取视频帧
ret = True
count = -1
vehicles = [2, 3, 5, 7]
while ret:
count += 1
ret, frame = cap.read()
# print(frame)
if ret and count < 10:
# 创建帧编号空字典,方便后面存储数据
results[count] = {}
# detections = coco_model(frame)[0]
detections = coco_model.track(frame, persist=True) # persist=True 会记住上一帧的信息
# detection = detections.boxes.data.tolist()
# print(detection)
detections_ = []
如果检测到车辆,则放入列表中
检测车牌,并将车牌分配给对应的车辆:
# 检测车牌
licence_plates = license_plate_detector(frame)[0]
for licence_plate in licence_plates.boxes.data.tolist():
x1, y1, x2, y2, score ,class_id = licence_plate
# get_car 分配每一个车牌给车辆
xcar_1, ycar_1, x_car2, y_car2, car_id = get_car(licence_plate, track_ids)
import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
# 函数集合
import csv
import string
import easyocr
# 指定模型存储目录
model_storage_directory = "E:\python_works\yolo\my_detect\my_detect\models"
# 定义字符拾取器
reader = easyocr.Reader(['en', 'ch_sim'], model_storage_directory=model_storage_directory)
def get_car(licence_plate, vehicle_track_ids):
"""
将识别车牌分配给已经追踪的车辆
"""
global car_index
x1, y1, x2, y2, score, class_id = licence_plate
foundIt = False
# 遍历所有车辆
for j in range(len(vehicle_track_ids)):
xcar1, ycar_1, xcar_2, ycar_2, car_id = vehicle_track_ids[j][:5]
# 坐标原点为左上角,横轴为x轴,纵轴为y轴
if x1 > xcar1 and y1 > ycar_1 and x2 < xcar_2 and y2 < ycar_2:
car_index = j
foundIt = True
break
if foundIt:
return vehicle_track_ids[car_index][:5]
return -1, -1, -1, -1, -1
分配算法是通过判断车牌 box 是否位于某个车辆 box 中,左上角为坐标原点,水平为 X 轴,竖直为 Y 轴,若分分配失败,则认为车辆识别或车牌识别出现误差。
如果分配成功,则裁剪并处理车牌图像:
# 如果车牌匹配到了车辆
if car_id != -1:
# 裁剪车牌
licence_plate_crop = frame[int(y1):int(y2), int(x1):int(x2), :]
# 转换成灰度图片
licence_plate_crop_gray = cv2.cvtColor(licence_plate_crop, cv2.COLOR_BGR2GRAY)
# 调整阈值,若低于64,则设置值为 255,否则设置成 0 ,输出阈值图像
_,licence_plate_crop_thresh = cv2.threshold(licence_plate_crop_gray, 64, 255, cv2.THRESH_BINARY_INV)
# # 显示灰度图片和阈值黑白图像
# cv2.imshow('original_crop', licence_plate_crop)
# cv2.imshow('threshold', licence_plate_crop_thresh)
#
# cv2.waitKey(0)
然后识别车牌:
# read_licence_plate 识别车牌字符
licence_plate_text, licence_plate_core = read_licence_plate(licence_plate_crop_thresh)
def read_licence_plate(licence_plate):
"""
识别车牌字符串
返回格式化字符串和置信度得分
"""
detections = reader.readtext(licence_plate)
for detection in detections:
bbox, text, score = detection
text = text.upper().replace(' ','')
print(text)
return 0, 0
识别车牌是通过 easyocr 的文字识别进行的,先创建识别器,将阈值黑白图像输入函数,去掉空格后打印车牌字符串。
车牌识别后,将所有信息都存储到 csv 表格中进行查看:
# 存储车牌信息
if licence_plate_text is not None:
results[count][car_id] = {'car':{'bbox':[xcar_1, ycar_1, x_car2, y_car2]},
'licence_plates':{'bbox':[x1, y1, x2, y2],
'text':licence_plate_text,
'bbox_score':score,
'text_score':licence_plate_core}}
# 写入csv
write_csv(results, './test.csv')
def write_csv(results, output_path):
"""
Write the results to a CSV file.
Args:
results (dict): Dictionary containing the results.
output_path (str): Path to the output CSV file.
"""
with open(output_path, 'w', newline='') as f: # 使用newline=''避免在Windows上出现多余的空行
writer = csv.writer(f)
# 写入表头
writer.writerow(['frame_nmr', 'car_id', 'car_bbox', 'license_plate_bbox',
'license_plate_bbox_score', 'license_number', 'license_number_score'])
for frame_nmr in results.keys():
for car_id in results[frame_nmr].keys():
print(results[frame_nmr][car_id])
if 'car' in results[frame_nmr][car_id] and \
'licence_plates' in results[frame_nmr][car_id] and \
'text' in results[frame_nmr][car_id]['licence_plates']:
car_bbox = results[frame_nmr][car_id]['car']['bbox']
license_plate_bbox = results[frame_nmr][car_id]['licence_plates']['bbox']
bbox_score = results[frame_nmr][car_id]['licence_plates']['bbox_score']
text = results[frame_nmr][car_id]['licence_plates']['text']
text_score = results[frame_nmr][car_id]['licence_plates']['text_score']
# 写入数据行
writer.writerow([
frame_nmr,
car_id,
'[{} {} {} {}]'.format(*car_bbox),
'[{} {} {} {}]'.format(*license_plate_bbox),
bbox_score,
text,
text_score
])
最后运行,识别到车牌信息及其对应车辆信息:
可以看到已能基本势必到车牌,并且都是 id 为2的车辆,这说明 id 为 1 的车辆并没有分配到对应的车牌,或者说根本没识别到车牌,这是可接受的,因为我的车牌识别模型是自己训练的,只用了大概 300 张训练集,识别不准确是正常的。然后 id 为 2 的车辆在这几个帧中识别到的车牌字符是不一样的,但是仔细观察发现 “肃R18”、“京凡168”、“就凡768”,识别到的车牌可以认为是很相像的,可以认为是字符识别误差或者图片质量误差,后期可设计算法匹配到正确的车牌号或者继续优化模型,以识别到更精确的车牌号,这里不做过多介绍,后期会再发博客探讨。
这里再提供下源码:
main.py:
import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
import numpy as np
from ultralytics import YOLO
import cv2
from util import *
# 存储所有信息
results = {}
# 加载模型
coco_model = YOLO('yolov8n.pt', task='detect')
license_plate_detector = YOLO('best.pt')
# 加载视频
cap = cv2.VideoCapture('./3.mp4')
# 读取视频帧
ret = True
count = -1
vehicles = [2, 3, 5, 7]
while ret:
count += 1
ret, frame = cap.read()
# print(frame)
if ret and count < 8:
# 创建帧编号空字典,方便后面存储数据
results[count] = {}
# detections = coco_model(frame)[0]
detections = coco_model.track(frame, persist=True) # persist=True 会记住上一帧的信息
# detection = detections.boxes.data.tolist()
# print(detection)
detections_ = []
if len(detections) > 0:
track_ids = detections[0].boxes.data.cpu().numpy().tolist()
for detection in track_ids:
x1, y1, x2, y2, score, class_id = detection[:6]
if int(class_id) in vehicles:
detections_.append([x1, y1, x2, y2, score])
else:
track_ids = []
print("Track IDs:", track_ids)
# for detection in detections.boxes.data.tolist():
# x1, y1, x2, y2, score ,class_id = detection
# # print
# if int(class_id) in vehicles:
# detections_.append([x1, y1, x2, y2, score])
# 汽车跟踪,返回的一个 追踪id 列表,每一个车辆都有 id 及其坐标,即使帧变了,同一个目标的 id 不会变
# track_ids = mot_tracker.update(np.asarray(detections_))
# 检测车牌
licence_plates = license_plate_detector(frame)[0]
for licence_plate in licence_plates.boxes.data.tolist():
x1, y1, x2, y2, score ,class_id = licence_plate
# get_car 分配每一个车牌给车辆
xcar_1, ycar_1, x_car2, y_car2, car_id = get_car(licence_plate, track_ids)
# 如果车牌匹配到了车辆
if car_id != -1:
# 裁剪车牌
licence_plate_crop = frame[int(y1):int(y2), int(x1):int(x2), :]
# 转换成灰度图片
licence_plate_crop_gray = cv2.cvtColor(licence_plate_crop, cv2.COLOR_BGR2GRAY)
# 调整阈值,若低于64,则设置值为 255,否则设置成 0 ,输出阈值图像
_,licence_plate_crop_thresh = cv2.threshold(licence_plate_crop_gray, 64, 255, cv2.THRESH_BINARY_INV)
# # 显示灰度图片和阈值黑白图像
# cv2.imshow('original_crop', licence_plate_crop)
# cv2.imshow('threshold', licence_plate_crop_thresh)
#
# cv2.waitKey(0)
# read_licence_plate 识别车牌字符
licence_plate_text, licence_plate_core = read_licence_plate(licence_plate_crop_thresh)
# 存储车牌信息
if licence_plate_text is not None:
results[count][car_id] = {'car':{'bbox':[xcar_1, ycar_1, x_car2, y_car2]},
'licence_plates':{'bbox':[x1, y1, x2, y2],
'text':licence_plate_text,
'bbox_score':score,
'text_score':licence_plate_core}}
# 写入csv
write_csv(results, './test.csv')
util.py:
import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
# 函数集合
import csv
import string
import easyocr
# 指定模型存储目录
model_storage_directory = "E:\python_works\yolo\my_detect\my_detect\models"
# 定义字符拾取器
reader = easyocr.Reader(['en', 'ch_sim'], model_storage_directory=model_storage_directory)
def get_car(licence_plate, vehicle_track_ids):
"""
将识别车牌分配给已经追踪的车辆
"""
global car_index
x1, y1, x2, y2, score, class_id = licence_plate
foundIt = False
# 遍历所有车辆
for j in range(len(vehicle_track_ids)):
xcar1, ycar_1, xcar_2, ycar_2, car_id = vehicle_track_ids[j][:5]
# 坐标原点为左上角,横轴为x轴,纵轴为y轴
if x1 > xcar1 and y1 > ycar_1 and x2 < xcar_2 and y2 < ycar_2:
car_index = j
foundIt = True
break
if foundIt:
return vehicle_track_ids[car_index][:5]
return -1, -1, -1, -1, -1
def read_licence_plate(licence_plate):
"""
识别车牌字符串
返回格式化字符串和置信度得分
"""
detections = reader.readtext(licence_plate)
if detections:
# 按照置信度排序,取最高分的结果
detections.sort(key=lambda x: x[2], reverse=True)
bbox, text, score = detections[0]
text = text.upper().replace(' ','')
print(text)
return text, score
else:
return None, None
def write_csv(results, output_path):
"""
Write the results to a CSV file.
Args:
results (dict): Dictionary containing the results.
output_path (str): Path to the output CSV file.
"""
with open(output_path, 'w', newline='') as f: # 使用newline=''避免在Windows上出现多余的空行
writer = csv.writer(f)
# 写入表头
writer.writerow(['frame_nmr', 'car_id', 'car_bbox', 'license_plate_bbox',
'license_plate_bbox_score', 'license_number', 'license_number_score'])
for frame_nmr in results.keys():
for car_id in results[frame_nmr].keys():
print(results[frame_nmr][car_id])
if 'car' in results[frame_nmr][car_id] and \
'licence_plates' in results[frame_nmr][car_id] and \
'text' in results[frame_nmr][car_id]['licence_plates']:
car_bbox = results[frame_nmr][car_id]['car']['bbox']
license_plate_bbox = results[frame_nmr][car_id]['licence_plates']['bbox']
bbox_score = results[frame_nmr][car_id]['licence_plates']['bbox_score']
text = results[frame_nmr][car_id]['licence_plates']['text']
text_score = results[frame_nmr][car_id]['licence_plates']['text_score']
# 写入数据行
writer.writerow([
frame_nmr,
car_id,
'[{} {} {} {}]'.format(*car_bbox),
'[{} {} {} {}]'.format(*license_plate_bbox),
bbox_score,
text,
text_score
])
# 示例调用
资源下载
感谢您的观看!!!