TensorFlow入门教程(30)车牌识别之整合EAST+DenseNet进行车牌识别(六)

本文主要是介绍TensorFlow入门教程(30)车牌识别之整合EAST+DenseNet进行车牌识别(六)，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

#
#作者：韦访
#博客：https://blog.csdn.net/rookie_wei
#微信：1007895847
#添加微信的备注一下是CSDN的
#欢迎大家一起学习
#

1、概述

前面几讲，我们已经分别实现了车牌检测和车牌号文本识别，现在就来将它们整合在一起进行完整的车牌识别。

环境配置：

操作系统：Ubuntu 64位

显卡：GTX 1080ti

Python：Python3.7

TensorFlow：2.3.0

2、车牌检测

首先需要做的是车牌检测，先导入所有要检测的图片，

'''
导入图片数据
'''
def get_images(data_path):files = []  for ext in ['jpg', 'png', 'jpeg']:files.extend(glob.glob(os.path.join(data_path, '*.{}'.format(ext))))    return files

接着导入并加载EAST模型，

'''
创建和加载east模型
'''
def east_create_and_load_model():# 创建模型# load trained modelmodel = EAST_model()ckpt = tf.train.Checkpoint(step=tf.Variable(0), model=model)latest_ckpt = tf.train.latest_checkpoint(FLAGS.east_model_path)# 加载模型if latest_ckpt:ckpt.restore(latest_ckpt)print('global_step : {}, checkpoint is restored!'.format(int(ckpt.step)))return model

接着就是遍历列表，并读取图片数据，

    # 遍历图片列表for filename in image_filenames:        image = cv2.imread(filename)print("filename:", filename)# 根据文件名规则判断文件名中是否包含车牌信息，如果包含的话，解析出车牌号，用于跟模型预测的车牌号对比if len(filename.split("-")) == 7:_,number = get_plate_attribute(filename, image)else:number = "None"

然后，将图片缩放至宽和高都是32的整数倍，并且返回缩放比例，

'''
缩放图片至宽和高都是32的整数倍，且限定最大长度
'''
def east_resize_image(image, max_side_len=2400):h, w, _ = image.shaperesize_w = wresize_h = h# limit the max sideif max(resize_h, resize_w) > max_side_len:ratio = float(max_side_len) / resize_h if resize_h > resize_w else float(max_side_len) / resize_welse:ratio = 1.resize_h = int(resize_h * ratio)resize_w = int(resize_w * ratio)resize_h = resize_h if resize_h % 32 == 0 else (resize_h // 32) * 32resize_w = resize_w if resize_w % 32 == 0 else (resize_w // 32) * 32image = cv2.resize(image, (int(resize_w), int(resize_h)))ratio_h = resize_h / float(h)ratio_w = resize_w / float(w)return image, (ratio_h, ratio_w)

接着，就是预测出score map和geometry map了，

        # 得到score_map和geo_mapscore_map, geo_map = east_model.predict(image_resized[np.newaxis, :, :, :])

得到 score map和geometry map以后，先还原出所有score map预测出的文本框里所有像素点的预测出的文本框，代码如下，

'''
通过geometry map和origin得到文本框
思路：先根据旋转角度是的正负，以及像素点相对最小外接矩形的距离，在原点处还原出矩形。再根据旋转角度得到旋转矩阵，然后求出5个关键点（四个顶点加像素点）旋转以后的坐标，最后再根据旋转后的像素点坐标和实际的像素点坐标的相对位置，将旋转后的矩形平移，即可得到文本框坐标。
'''
def get_polys_by_geometry(origin, geometry, is_positive_angle):# 得到距离d = geometry[:, :4]# 得到角度angle = geometry[:, 4]    if is_positive_angle:# 角度大于0的情况origin = origin[angle >= 0]d = d[angle >= 0]angle = angle[angle >= 0]else:# 角度小于0的情况origin = origin[angle < 0]d = d[angle < 0]angle = angle[angle < 0]if origin.shape[0] > 0:        # top-0, right-1, bottom-2, left-3d_top = d[:,0]d_right = d[:,1]d_bottom = d[:,2]d_left = d[:,3]d_w = d_left + d_rightd_h = d_top + d_bottomp_zeros = np.zeros(d_top.shape[0])if is_positive_angle:# 角度大于0，则以p3顶点为原点，得到平行x轴的矩形p0 = [p_zeros, -d_h]p1 = [d_w, -d_h]p2 = [d_w, p_zeros]p3 = [p_zeros, p_zeros]p_orgin = [d_left, -d_bottom] # 这个是像素点相对矩形的位置# 旋转矩阵rotate_matrix_x = np.array([np.cos(angle), np.sin(angle)]).transpose((1, 0))rotate_matrix_y = np.array([-np.sin(angle), np.cos(angle)]).transpose((1, 0))else:# 角度大于0，则以p2顶点为原点，得到平行x轴的矩形p0 = [-d_w, -d_h]p1 = [p_zeros, -d_h]p2 = [p_zeros, p_zeros]p3 = [-d_w, p_zeros]p_orgin = [-d_right, -d_bottom] # 这个是像素点相对矩形的位置# 旋转矩阵rotate_matrix_x = np.array([np.cos(-angle), -np.sin(-angle)]).transpose((1, 0))rotate_matrix_y = np.array([np.sin(-angle), np.cos(-angle)]).transpose((1, 0))# 旋转矩阵，因为总共要旋转5个点，所以要repeat出5个旋转矩阵rotate_matrix_x = np.repeat(rotate_matrix_x, 5, axis=1).reshape(-1, 2, 5).transpose((0, 2, 1))  # N*5*2rotate_matrix_y = np.repeat(rotate_matrix_y, 5, axis=1).reshape(-1, 2, 5).transpose((0, 2, 1))# 将上面得到的5个点经过旋转矩阵得到旋转以后的点的坐标p = np.asarray(np.concatenate([p0, p1, p2, p3, p_orgin])).transpose((1, 0)).reshape((-1, 5, 2)) # N*5*2p_rotate_x = np.sum(rotate_matrix_x * p, axis=2)[:, :, np.newaxis]  # N*5*1p_rotate_y = np.sum(rotate_matrix_y * p, axis=2)[:, :, np.newaxis]  # N*5*1# 这里得到的是旋转后的坐标p_rotate = np.concatenate([p_rotate_x, p_rotate_y], axis=2)  # N*5*2# 根据旋转后的像素点坐标和原像素点坐标的相对位置，平移整个矩形p3_in_origin = origin - p_rotate[:, 4, :]new_p0 = p_rotate[:, 0, :] + p3_in_originnew_p1 = p_rotate[:, 1, :] + p3_in_originnew_p2 = p_rotate[:, 2, :] + p3_in_originnew_p3 = p_rotate[:, 3, :] + p3_in_origin# 得到最终的文本框new_p_0 = np.concatenate([new_p0[:, np.newaxis, :], new_p1[:, np.newaxis, :],new_p2[:, np.newaxis, :], new_p3[:, np.newaxis, :]], axis=1)  # N*5*2else:new_p_0 = np.zeros((0, 4, 2))return new_p_0'''
根据score map和geometry map得到所有相关像素点的文本框
'''
def restore_rectangle(origin, geometry):return np.concatenate([get_polys_by_geometry(origin, geometry, is_positive_angle=True), get_polys_by_geometry(origin, geometry, is_positive_angle=False)])

得到这些文本框以后，再通过lanms得到最终的文本框，

'''
east预测车牌坐标
'''
def east_detect(score_map, geo_map, timer, score_map_thresh=0.8, box_thresh=0.1, nms_thres=0.2):np.set_printoptions(threshold=np.inf)if len(score_map.shape) == 4:score_map = score_map[0, :, :, 0]geo_map = geo_map[0, :, :, ]# filter the score map# 过滤得分小于score_map_thresh的区域xy_text = np.argwhere(score_map > score_map_thresh)# sort the text boxes via the y axisxy_text = xy_text[np.argsort(xy_text[:, 0])]# restorestart = time.time()text_box_restored = restore_rectangle(xy_text[:, ::-1]*4, geo_map[xy_text[:, 0], xy_text[:, 1], :]) # N*4*2# print('{} text boxes before nms'.format(text_box_restored.shape[0]))boxes = np.zeros((text_box_restored.shape[0], 9), dtype=np.float32)boxes[:, :8] = text_box_restored.reshape((-1, 8))boxes[:, 8] = score_map[xy_text[:, 0], xy_text[:, 1]]timer['restore'] = time.time() - start# nms partstart = time.time()# boxes = nms_locality.nms_locality(boxes.astype(np.float64), nms_thres)boxes = lanms.merge_quadrangle_n9(boxes.astype('float32'), nms_thres)timer['nms'] = time.time() - startif boxes.shape[0] == 0:return None, timer# here we filter some low score boxes by the average score map, this is different from the orginal paperfor i, box in enumerate(boxes):mask = np.zeros_like(score_map, dtype=np.uint8)cv2.fillPoly(mask, box[:8].reshape((-1, 4, 2)).astype(np.int32) // 4, 1)boxes[i, 8] = cv2.mean(score_map, mask)[0]boxes = boxes[boxes[:, 8] > box_thresh]return boxes, timer

得到文本框坐标以后，就要把车牌单独截取出来了，先遍历所有文本框，过滤太小的文本框，

        # 如果有车牌坐标，则去识别车牌号if boxes is not None:# 先将车牌坐标恢复成跟缩放前图片对应的坐标boxes = boxes[:, :8].reshape((-1, 4, 2))boxes[:, :, 0] /= ratio_wboxes[:, :, 1] /= ratio_h# 遍历所有识别出来的车牌坐标for box in boxes:# 重新排序顶点box = east_sort_poly(box.astype(np.int32))# 过滤太小的坐标或者小于0的坐标if np.linalg.norm(box[0] - box[1]) < 10 or np.linalg.norm(box[3]-box[0]) < 10 or np.min(box) < 0:continue

然后，获取文本框的最小外接矩形，再重新排序顶点，接着旋转图片以使最小外接矩形水平，然后就可以截取出车牌了。代码如下，

def get_min_area_rect(poly):rect = cv2.minAreaRect(poly)box = cv2.boxPoints(rect)return box, rect[2]'''
对矩形的顶点进行排序，排序后的结果是，p0-左上角，p1-右上角，p2-右下角，p3-左下角
矩形与水平轴的夹角，即为p2_p3边到x轴的夹角，以逆时针为正，顺时针为负
'''
def sort_poly_and_get_angle(poly, image=None):    # 先找到矩形最底部的顶点p_lowest = np.argmax(poly[:, 1])if np.count_nonzero(poly[:, 1] == poly[p_lowest, 1]) == 2:# 如果矩形底部的边刚好与x轴平行# 这种情况下，x坐标加y坐标之和最小的顶点就是左上角的顶点，即p0p0_index = np.argmin(np.sum(poly, axis=1))p1_index = (p0_index + 1) % 4p2_index = (p0_index + 2) % 4p3_index = (p0_index + 3) % 4return poly[[p0_index, p1_index, p2_index, p3_index]], 0.else:        # 如果矩形底部与x轴有夹角# 找到最底部顶点的右边的第一个顶点p_lowest_right = (p_lowest - 1) % 4    # 求最底部顶点与其右边第一个顶点组成的边与x轴的夹角                angle = np.arctan(-(poly[p_lowest][1] - poly[p_lowest_right][1])/(poly[p_lowest][0] - poly[p_lowest_right][0] + 1e-5))        # 下面的代码其实自己画个图就很好理解了if angle > np.pi/4:    # 如果这个夹角大于45度，那么，最底部的顶点为p2顶点        p2_index = p_lowestp1_index = (p2_index - 1) % 4p0_index = (p2_index - 2) % 4p3_index = (p2_index + 1) % 4# 这种情况下，p2-p3边与x轴的夹角就为-(np.pi/2 - angle)return poly[[p0_index, p1_index, p2_index, p3_index]], -(np.pi/2 - angle)else:# 如果这个夹角小于等于45度，那么，最底部的顶点为p3顶点p3_index = p_lowestp0_index = (p3_index + 1) % 4p1_index = (p3_index + 2) % 4p2_index = (p3_index + 3) % 4return poly[[p0_index, p1_index, p2_index, p3_index]], angledef get_plate_image(image, poly):# DEBUG = Trueif DEBUG:image = draw_line(image, poly)box, angle = get_min_area_rect(poly)if DEBUG:image = draw_line(image, box, (0,255,0))  (p0, p1, p2, p3), angle = sort_poly_and_get_angle(box)image = crop_plate(image, angle, p0, p1, p2, p3)# cv2.imshow("get_min_area_rect", image)# cv2.waitKey(0)return image# 将车牌裁剪出来，因为车牌有可能是斜的，所以要先将图片旋转到车牌的矩形框为水平时，再裁剪
def crop_plate(image, angle, p0, p1, p2, p3):# DEBUG = Trueangle = -angle * (180 / math.pi)# print("angle:", angle)    h, w, _ = image.shaperotate_mat = cv2.getRotationMatrix2D((w / 2, h / 2), angle, 1)  # 按angle角度旋转图像h_new = int(w * np.fabs(np.sin(np.radians(angle))) + h * np.fabs(np.cos(np.radians(angle))))w_new = int(h * np.fabs(np.sin(np.radians(angle))) + w * np.fabs(np.cos(np.radians(angle))))rotate_mat[0, 2] += (w_new - w) / 2rotate_mat[1, 2] += (h_new - h) / 2rotated_image = cv2.warpAffine(image, rotate_mat, (w_new, h_new), borderValue=(255, 255, 255))# 旋转后图像的四点坐标[[p1[0]], [p1[1]]] = np.dot(rotate_mat, np.array([[p1[0]], [p1[1]], [1]]))[[p3[0]], [p3[1]]] = np.dot(rotate_mat, np.array([[p3[0]], [p3[1]], [1]]))[[p2[0]], [p2[1]]] = np.dot(rotate_mat, np.array([[p2[0]], [p2[1]], [1]]))[[p0[0]], [p0[1]]] = np.dot(rotate_mat, np.array([[p0[0]], [p0[1]], [1]]))if DEBUG:cv2.circle(rotated_image, tuple(p0), 10, (0,255,0), 4)cv2.circle(rotated_image, tuple(p1), 10, (0,0,255), 4)cv2.imshow('image',  image)cv2.imshow('rotateImg2',  rotated_image)cv2.waitKey(0)crop_image = rotated_image[int(p0[1]):int(p3[1]), int(p0[0]):int(p1[0])]return crop_image

3、车牌号识别

得到车牌以后，就要对车牌进行字符识别了，先导入DenseNet模型，

'''
创建和加载densenet模型
'''
def densenet_create_and_load_model():_,char_list = get_char_vector(FLAGS.char_filename)num_classes = len(char_list) + 1model,_,_ = densenet(FLAGS, num_classes)h5_path = os.path.join("./densenet/models", "plate-"+FLAGS.densenet_model_net+".h5") if os.path.exists(h5_path):        model.load_weights(h5_path)return model, char_list, num_classes

然后将车牌图片缩放至固定高度，代码如下，

def pading_plate_width(image):h, w, _ = image.shapenew_w = np.where(FLAGS.input_size_w > w+20, FLAGS.input_size_w, w+20)new_image = np.zeros((h, new_w, 3), dtype=np.uint8)start_w = 10new_image[:h,start_w:w+start_w,:] = image[:h,:w,:]return new_image'''
随机旋转，这里不旋转一下识别效果反而降低，可能是训练的时候大部分都旋转的原因
'''
def random_rotate(images):h, w, _ = images.shaperandom_angle =  np.random.randint(-15,15)random_scale =  np.random.randint(8,10) / 10.0mat_rotate = cv2.getRotationMatrix2D(center=(w*0.5, h*0.5), angle=random_angle, scale=random_scale)images = cv2.warpAffine(images, mat_rotate, (w, h))return images'''
将图片缩放至固定高度，且保持原来的宽高比
'''
def densenet_resize_image(image):image = random_rotate(image)h, w, _ = image.shape    new_w = np.around(w / (h/FLAGS.input_size_h)).astype(np.int32)    image = cv2.resize(image, (new_w, FLAGS.input_size_h))image = pading_plate_width(image)# cv2.imshow("image", image)    image = image[np.newaxis,:,:,:]return image

接着直接将图片输入到模型中，得到预测结果，

                # densenet模型识别车牌号y_pred = densenet_model.predict(resized_plate_image)

对上面的预测结果还不能直接用，还要先删除掉ctc的blank字符和重复字符，再将其转成文字，

'''
解析densenet的预测结果，将blank和重复字符去掉
'''
def densenet_decode(pred, char_list, num_classes):plate_char_list = []pred_text = pred.argmax(axis=2)[0]# print("pred_text:", pred_text, " char_list len:", len(char_list))for i in range(len(pred_text)):if pred_text[i] != num_classes - 1 and ((not (i > 0 and pred_text[i] == pred_text[i - 1])) or (i > 1 and pred_text[i] == pred_text[i - 2])):plate_char_list.append(char_list[pred_text[i]])return u''.join(plate_char_list)

这样，我们就得到了车牌号了，为了方便查看，将识别结果显示出来，

'''
显示预测的车牌号
'''
def show_plate_number(image, number, coor):h, w, _ = image.shape    font = ImageFont.truetype(FLAGS.fontpath, 25)draw = ImageDraw.Draw(Image.fromarray(image))    textsize = draw.textsize(number, font=font)end_x = coor[0] + textsize[0] end_x = np.where(end_x > w, w, end_x)            image[coor[1]:coor[1]+textsize[1], coor[0]:end_x] = 0img = Image.fromarray(image)draw = ImageDraw.Draw(img)draw.text(coor, number, font = font, fill = (0,255,255,0))image = np.array(img)  return image