当前位置:   article > 正文

[python]毕业设计基于yolov8车牌检测识别中文车牌识别检测python源码和实现过程_yolov8识别汉字

yolov8识别汉字

yolov8车牌识别算法,支持12种中文车牌类型

图片测试demo:

直接运行detect_plate.py 或者运行如下命令行

python detect_rec_plate.py --detect_model weights/yolov8s.pt  --rec_model weights/plate_rec_color.pth --image_path imgs --output result

测试文件夹imgs,结果保存再 result 文件夹中

车牌检测训练

车牌识别训练

支持如下:

  •  1.单行蓝牌
  •  2.单行黄牌
  •  3.新能源车牌
  •  4.白色警用车牌
  •  5.教练车牌
  •  6.武警车牌
  •  7.双层黄牌
  •  8.双层白牌
  •  9.使馆车牌
  •  10.港澳粤Z牌
  •  11.双层绿牌
  •  12.民航车牌

识别核心代码:

  1. import torch
  2. import cv2
  3. import numpy as np
  4. import argparse
  5. import copy
  6. import time
  7. import os
  8. from ultralytics.nn.tasks import attempt_load_weights
  9. from plate_recognition.plate_rec import get_plate_result,init_model,cv_imread
  10. from plate_recognition.double_plate_split_merge import get_split_merge
  11. from fonts.cv_puttext import cv2ImgAddText
  12. def allFilePath(rootPath,allFIleList):# 读取文件夹内的文件,放到list
  13. fileList = os.listdir(rootPath)
  14. for temp in fileList:
  15. if os.path.isfile(os.path.join(rootPath,temp)):
  16. allFIleList.append(os.path.join(rootPath,temp))
  17. else:
  18. allFilePath(os.path.join(rootPath,temp),allFIleList)
  19. def four_point_transform(image, pts): #透视变换得到车牌小图
  20. # rect = order_points(pts)
  21. rect = pts.astype('float32')
  22. (tl, tr, br, bl) = rect
  23. widthA = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
  24. widthB = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
  25. maxWidth = max(int(widthA), int(widthB))
  26. heightA = np.sqrt(((tr[0] - br[0]) ** 2) + ((tr[1] - br[1]) ** 2))
  27. heightB = np.sqrt(((tl[0] - bl[0]) ** 2) + ((tl[1] - bl[1]) ** 2))
  28. maxHeight = max(int(heightA), int(heightB))
  29. dst = np.array([
  30. [0, 0],
  31. [maxWidth - 1, 0],
  32. [maxWidth - 1, maxHeight - 1],
  33. [0, maxHeight - 1]], dtype = "float32")
  34. M = cv2.getPerspectiveTransform(rect, dst)
  35. warped = cv2.warpPerspective(image, M, (maxWidth, maxHeight))
  36. return warped
  37. def letter_box(img,size=(640,640)): #yolo 前处理 letter_box操作
  38. h,w,_=img.shape
  39. r=min(size[0]/h,size[1]/w)
  40. new_h,new_w=int(h*r),int(w*r)
  41. new_img = cv2.resize(img,(new_w,new_h))
  42. left= int((size[1]-new_w)/2)
  43. top=int((size[0]-new_h)/2)
  44. right = size[1]-left-new_w
  45. bottom=size[0]-top-new_h
  46. img =cv2.copyMakeBorder(new_img,top,bottom,left,right,cv2.BORDER_CONSTANT,value=(114,114,114))
  47. return img,r,left,top
  48. def load_model(weights, device): #加载yolov8 模型
  49. model = attempt_load_weights(weights,device=device) # load FP32 model
  50. return model
  51. def xywh2xyxy(det): #xywh转化为xyxy
  52. y = det.clone()
  53. y[:,0]=det[:,0]-det[0:,2]/2
  54. y[:,1]=det[:,1]-det[0:,3]/2
  55. y[:,2]=det[:,0]+det[0:,2]/2
  56. y[:,3]=det[:,1]+det[0:,3]/2
  57. return y
  58. def my_nums(dets,iou_thresh): #nms操作
  59. y = dets.clone()
  60. y_box_score = y[:,:5]
  61. index = torch.argsort(y_box_score[:,-1],descending=True)
  62. keep = []
  63. while index.size()[0]>0:
  64. i = index[0].item()
  65. keep.append(i)
  66. x1=torch.maximum(y_box_score[i,0],y_box_score[index[1:],0])
  67. y1=torch.maximum(y_box_score[i,1],y_box_score[index[1:],1])
  68. x2=torch.minimum(y_box_score[i,2],y_box_score[index[1:],2])
  69. y2=torch.minimum(y_box_score[i,3],y_box_score[index[1:],3])
  70. zero_=torch.tensor(0).to(device)
  71. w=torch.maximum(zero_,x2-x1)
  72. h=torch.maximum(zero_,y2-y1)
  73. inter_area = w*h
  74. nuion_area1 =(y_box_score[i,2]-y_box_score[i,0])*(y_box_score[i,3]-y_box_score[i,1]) #计算交集
  75. union_area2 =(y_box_score[index[1:],2]-y_box_score[index[1:],0])*(y_box_score[index[1:],3]-y_box_score[index[1:],1])#计算并集
  76. iou = inter_area/(nuion_area1+union_area2-inter_area)#计算iou
  77. idx = torch.where(iou<=iou_thresh)[0] #保留iou小于iou_thresh的
  78. index=index[idx+1]
  79. return keep
  80. def restore_box(dets,r,left,top): #坐标还原到原图上
  81. dets[:,[0,2]]=dets[:,[0,2]]-left
  82. dets[:,[1,3]]= dets[:,[1,3]]-top
  83. dets[:,:4]/=r
  84. # dets[:,5:13]/=r
  85. return dets
  86. # pass
  87. def post_processing(prediction,conf,iou_thresh,r,left,top): #后处理
  88. prediction = prediction.permute(0,2,1).squeeze(0)
  89. xc = prediction[:, 4:6].amax(1) > conf #过滤掉小于conf的框
  90. x = prediction[xc]
  91. if not len(x):
  92. return []
  93. boxes = x[:,:4] #框
  94. boxes = xywh2xyxy(boxes) #中心点 宽高 变为 左上 右下两个点
  95. score,index = torch.max(x[:,4:6],dim=-1,keepdim=True) #找出得分和所属类别
  96. x = torch.cat((boxes,score,x[:,6:14],index),dim=1) #重新组合
  97. score = x[:,4]
  98. keep =my_nums(x,iou_thresh)
  99. x=x[keep]
  100. x=restore_box(x,r,left,top)
  101. return x
  102. def pre_processing(img,opt,device): #前处理
  103. img, r,left,top= letter_box(img,(opt.img_size,opt.img_size))
  104. # print(img.shape)
  105. img=img[:,:,::-1].transpose((2,0,1)).copy() #bgr2rgb hwc2chw
  106. img = torch.from_numpy(img).to(device)
  107. img = img.float()
  108. img = img/255.0
  109. img =img.unsqueeze(0)
  110. return img ,r,left,top
  111. def det_rec_plate(img,img_ori,detect_model,plate_rec_model):
  112. result_list=[]
  113. img,r,left,top = pre_processing(img,opt,device) #前处理
  114. predict = detect_model(img)[0]
  115. outputs=post_processing(predict,0.3,0.5,r,left,top) #后处理
  116. for output in outputs:
  117. result_dict={}
  118. output = output.squeeze().cpu().numpy().tolist()
  119. rect=output[:4]
  120. rect = [int(x) for x in rect]
  121. label = output[-1]
  122. roi_img = img_ori[rect[1]:rect[3],rect[0]:rect[2]]
  123. # land_marks=np.array(output[5:13],dtype='int64').reshape(4,2)
  124. # roi_img = four_point_transform(img_ori,land_marks) #透视变换得到车牌小图
  125. if int(label): #判断是否是双层车牌,是双牌的话进行分割后然后拼接
  126. roi_img=get_split_merge(roi_img)
  127. plate_number,rec_prob,plate_color,color_conf=get_plate_result(roi_img,device,plate_rec_model,is_color=True)
  128. result_dict['plate_no']=plate_number #车牌号
  129. result_dict['plate_color']=plate_color #车牌颜色
  130. result_dict['rect']=rect #车牌roi区域
  131. result_dict['detect_conf']=output[4] #检测区域得分
  132. # result_dict['landmarks']=land_marks.tolist() #车牌角点坐标
  133. # result_dict['rec_conf']=rec_prob #每个字符的概率
  134. result_dict['roi_height']=roi_img.shape[0] #车牌高度
  135. # result_dict['plate_color']=plate_color
  136. # if is_color:
  137. result_dict['color_conf']=color_conf #颜色得分
  138. result_dict['plate_type']=int(label) #单双层 0单层 1双层
  139. result_list.append(result_dict)
  140. return result_list
  141. def draw_result(orgimg,dict_list,is_color=False): # 车牌结果画出来
  142. result_str =""
  143. for result in dict_list:
  144. rect_area = result['rect']
  145. x,y,w,h = rect_area[0],rect_area[1],rect_area[2]-rect_area[0],rect_area[3]-rect_area[1]
  146. padding_w = 0.05*w
  147. padding_h = 0.11*h
  148. rect_area[0]=max(0,int(x-padding_w))
  149. rect_area[1]=max(0,int(y-padding_h))
  150. rect_area[2]=min(orgimg.shape[1],int(rect_area[2]+padding_w))
  151. rect_area[3]=min(orgimg.shape[0],int(rect_area[3]+padding_h))
  152. height_area = result['roi_height']
  153. # landmarks=result['landmarks']
  154. result_p = result['plate_no']
  155. if result['plate_type']==0:#单层
  156. result_p+=" "+result['plate_color']
  157. else: #双层
  158. result_p+=" "+result['plate_color']+"双层"
  159. result_str+=result_p+" "
  160. # for i in range(4): #关键点
  161. # cv2.circle(orgimg, (int(landmarks[i][0]), int(landmarks[i][1])), 5, clors[i], -1)
  162. cv2.rectangle(orgimg,(rect_area[0],rect_area[1]),(rect_area[2],rect_area[3]),(0,0,255),2) #画框
  163. labelSize = cv2.getTextSize(result_p,cv2.FONT_HERSHEY_SIMPLEX,0.5,1) #获得字体的大小
  164. if rect_area[0]+labelSize[0][0]>orgimg.shape[1]: #防止显示的文字越界
  165. rect_area[0]=int(orgimg.shape[1]-labelSize[0][0])
  166. orgimg=cv2.rectangle(orgimg,(rect_area[0],int(rect_area[1]-round(1.6*labelSize[0][1]))),(int(rect_area[0]+round(1.2*labelSize[0][0])),rect_area[1]+labelSize[1]),(255,255,255),cv2.FILLED)#画文字框,背景白色
  167. if len(result)>=6:
  168. orgimg=cv2ImgAddText(orgimg,result_p,rect_area[0],int(rect_area[1]-round(1.6*labelSize[0][1])),(0,0,0),21)
  169. # orgimg=cv2ImgAddText(orgimg,result_p,rect_area[0]-height_area,rect_area[1]-height_area-10,(0,255,0),height_area)
  170. print(result_str)
  171. return orgimg
  172. if __name__ == "__main__":
  173. parser = argparse.ArgumentParser()
  174. parser.add_argument('--detect_model', nargs='+', type=str, default=r'weights/yolov8s.pt', help='model.pt path(s)') #yolov8检测模型
  175. parser.add_argument('--rec_model', type=str, default=r'weights/plate_rec_color.pth', help='model.pt path(s)')#车牌字符识别模型
  176. parser.add_argument('--image_path', type=str, default=r'imgs', help='source') #待识别图片路径
  177. parser.add_argument('--img_size', type=int, default=640, help='inference size (pixels)') #yolov8 网络模型输入大小
  178. parser.add_argument('--output', type=str, default='result', help='source') #结果保存的文件夹
  179. device =torch.device("cuda" if torch.cuda.is_available() else "cpu")
  180. clors = [(255,0,0),(0,255,0),(0,0,255),(255,255,0),(0,255,255)]
  181. opt = parser.parse_args()
  182. save_path = opt.output
  183. if not os.path.exists(save_path):
  184. os.mkdir(save_path)
  185. detect_model = load_model(opt.detect_model, device) #初始化yolov8识别模型
  186. plate_rec_model=init_model(device,opt.rec_model,is_color=True) #初始化识别模型
  187. #算参数量
  188. total = sum(p.numel() for p in detect_model.parameters())
  189. total_1 = sum(p.numel() for p in plate_rec_model.parameters())
  190. print("yolov8 detect params: %.2fM,rec params: %.2fM" % (total/1e6,total_1/1e6))
  191. detect_model.eval()
  192. # print(detect_model)
  193. file_list = []
  194. allFilePath(opt.image_path,file_list)
  195. count=0
  196. time_all = 0
  197. time_begin=time.time()
  198. for pic_ in file_list:
  199. print(count,pic_,end=" ")
  200. time_b = time.time() #开始时间
  201. img = cv2.imread(pic_)
  202. img_ori = copy.deepcopy(img)
  203. result_list=det_rec_plate(img,img_ori,detect_model,plate_rec_model)
  204. time_e=time.time()
  205. ori_img=draw_result(img,result_list) #将结果画在图上
  206. img_name = os.path.basename(pic_)
  207. save_img_path = os.path.join(save_path,img_name) #图片保存的路径
  208. time_gap = time_e-time_b #计算单个图片识别耗时
  209. if count:
  210. time_all+=time_gap
  211. count+=1
  212. cv2.imwrite(save_img_path,ori_img) #op
  213. # print(result_list)
  214. print(f"sumTime time is {time.time()-time_begin} s, average pic time is {time_all/(len(file_list)-1)}")

完整源码下载地址:https://download.csdn.net/download/FL1768317420/89325363

声明:本文内容由网友自发贡献,转载请注明出处:【wpsshop博客】
推荐阅读
相关标签
  

闽ICP备14008679号