├── LICENSE ├── README.md ├── code ├── carPlateIdentity.py ├── char.pth ├── charNeuralNet.py ├── plate.pth └── plateNeuralNet.py └── images.zip /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 zxbsmk 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # 基于 PyTorch 和 OpenCV 的入门级车牌识别项目 2 | 3 | 这是中山大学智能工程学院大二的图像处理课程的第三次大作业。任务是处理车牌号码识别。 4 | 5 | 为了完成该作业，我参考了网上诸多资料，最终选择了如今的解决方案。即使用 OpenCV 中的传统图像方法对输入图像进行预处理，然后将车牌过滤和字符识别交给神经网络处理。 6 | 7 | 用到的两个神经网络都是由我自己使用 PyTorch 构建的，结构上借鉴了 LeNet5，十分简单，对于这种简单的任务可以说是绰绰有余，只需稍微训练即可取得不错的效果。 8 | 9 | 与其他的 repo 不同，这里我还给出了预训练好的网络模型 plate.pth 和 char.pth，便于大家复现出我的项目。 10 | 11 | # 环境配置 12 | 13 | Python 3.6 14 | 15 | PyTorch 1.6.0+cu101 16 | 17 | OpenCV 4.4.0 18 | 19 | # 模型训练 20 | 21 | 需要注意的是，我的模型之前是在服务器上加载和运行的，模型结果是保存在第八张显卡上的。如果你的运行环境中显卡数目**少于八张**或者说第八张显卡的显存小于2G，则无法直接使用我的预训练模型。 22 | 23 | 如果需要自行训练模型，首先需要解压数据集 images.zip 使得其与 code 在同一目录下，然后还需要对 plateNeuralNet.py 和 charNeuralNet.py 进行一些修改，即 24 | 25 | * 将文件头的 torch.cuda.set_device(7) 改为 torch.cuda.set_device(0) 26 | * 将主函数中的 model = torch.load(train_model_path) 注释掉 27 | * 恢复 model = char_cnn_net() 和 model = plate_cnn_net() 以及 model.apply(weights_init) 28 | 29 | 接下来只需要运行以下代码即可开始训练： 30 | 31 | ```bash 32 | python3 plateNeuralNet.py 33 | python3 charNeuralNet.py 34 | ``` 35 | 36 | 由于训练需要一定时间，也可以用以下命令将进程挂到后台运行： 37 | 38 | ```bash 39 | nohup python3 plateNeuralNet.py 1>plate.txt & 40 | nohup python3 charNeuralNet.py 1>char.txt & 41 | cat plate.txt 42 | cat char.txt 43 | ``` 44 | 45 | 使用 cat 命令即可实时地查看模型训练的进度。 46 | 47 | # 项目运行 48 | 49 | 完成模型的训练后，替换掉原有的预训练模型。检查好数据所在的路径是否有误，然后就可以用以下命令开始运行项目： 50 | 51 | ```bash 52 | python3 carPlateIdentity.py 53 | ``` 54 | 55 | 如果需要测试特定图像，可以放进 images/test/ 目录下。 -------------------------------------------------------------------------------- /code/carPlateIdentity.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import os 3 | import sys 4 | import numpy as np 5 | import torch 6 | from torchvision import transforms 7 | from plateNeuralNet import * 8 | from charNeuralNet import * 9 | import random 10 | 11 | torch.cuda.set_device(7) 12 | 13 | def setup_seed(seed): 14 | torch.manual_seed(seed) 15 | torch.cuda.manual_seed_all(seed) 16 | np.random.seed(seed) 17 | random.seed(seed) 18 | torch.backends.cudnn.deterministic = True 19 | 20 | setup_seed(233) 21 | 22 | char_table = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 23 | 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', '川', '鄂', '赣', '甘', '贵', 24 | '桂', '黑', '沪', '冀', '津', '京', '吉', '辽', '鲁', '蒙', '闽', '宁', '青', '琼', '陕', '苏', '晋', 25 | '皖', '湘', '新', '豫', '渝', '粤', '云', '藏', '浙'] 26 | 27 | def hist_image(img): # 灰度级转灰度图 28 | assert img.ndim==2 29 | hist = [0 for i in range(256)] 30 | img_h,img_w = img.shape[0],img.shape[1] 31 | 32 | for row in range(img_h): 33 | for col in range(img_w): 34 | hist[img[row,col]] += 1 35 | p = [hist[n]/(img_w*img_h) for n in range(256)] 36 | p1 = np.cumsum(p) 37 | for row in range(img_h): 38 | for col in range(img_w): 39 | v = img[row,col] 40 | img[row,col] = p1[v]*255 41 | return img 42 | 43 | def find_board_area(img): # 通过检测行和列的亮点数目来提取矩形 44 | assert img.ndim==2 45 | img_h,img_w = img.shape[0],img.shape[1] 46 | top,bottom,left,right = 0,img_h,0,img_w 47 | flag = False 48 | h_proj = [0 for i in range(img_h)] 49 | v_proj = [0 for i in range(img_w)] 50 | 51 | for row in range(round(img_h*0.5),round(img_h*0.8),3): 52 | for col in range(img_w): 53 | if img[row,col]==255: 54 | h_proj[row] += 1 55 | if flag==False and h_proj[row]>12: 56 | flag = True 57 | top = row 58 | if flag==True and row>top+8 and h_proj[row]<12: 59 | bottom = row 60 | flag = False 61 | 62 | for col in range(round(img_w*0.3),img_w,1): 63 | for row in range(top,bottom,1): 64 | if img[row,col]==255: 65 | v_proj[col] += 1 66 | if flag==False and (v_proj[col]>10 or v_proj[col]-v_proj[col-1]>5): 67 | left = col 68 | break 69 | return left,top,120,bottom-top-10 70 | 71 | def verify_scale(rotate_rect): 72 | error = 0.4 73 | aspect = 4#4.7272 74 | min_area = 10*(10*aspect) 75 | max_area = 150*(150*aspect) 76 | min_aspect = aspect*(1-error) 77 | max_aspect = aspect*(1+error) 78 | theta = 30 79 | 80 | # 宽或高为0，不满足矩形直接返回False 81 | if rotate_rect[1][0]==0 or rotate_rect[1][1]==0: 82 | return False 83 | 84 | r = rotate_rect[1][0]/rotate_rect[1][1] 85 | r = max(r,1/r) 86 | area = rotate_rect[1][0]*rotate_rect[1][1] 87 | if area>min_area and areamin_aspect and r= -90 and rotate_rect[2] < -(90 - theta)) or 90 | (rotate_rect[1][1] < rotate_rect[1][0] and rotate_rect[2] > -theta and rotate_rect[2] <= 0)): 91 | return True 92 | return False 93 | 94 | def img_Transform(car_rect,image): 95 | img_h,img_w = image.shape[:2] 96 | rect_w,rect_h = car_rect[1][0],car_rect[1][1] 97 | angle = car_rect[2] 98 | 99 | return_flag = False 100 | if car_rect[2]==0: 101 | return_flag = True 102 | if car_rect[2]==-90 and rect_w point[0]: 117 | left_point = point 118 | if low_point[1] > point[1]: 119 | low_point = point 120 | if heigth_point[1] < point[1]: 121 | heigth_point = point 122 | if right_point[0] < point[0]: 123 | right_point = point 124 | 125 | if left_point[1] <= right_point[1]: # 正角度 126 | new_right_point = [right_point[0], heigth_point[1]] 127 | pts1 = np.float32([left_point, heigth_point, right_point]) 128 | pts2 = np.float32([left_point, heigth_point, new_right_point]) # 字符只是高度需要改变 129 | M = cv2.getAffineTransform(pts1, pts2) 130 | dst = cv2.warpAffine(image, M, (round(img_w*2), round(img_h*2))) 131 | car_img = dst[int(left_point[1]):int(heigth_point[1]), int(left_point[0]):int(new_right_point[0])] 132 | 133 | elif left_point[1] > right_point[1]: # 负角度 134 | new_left_point = [left_point[0], heigth_point[1]] 135 | pts1 = np.float32([left_point, heigth_point, right_point]) 136 | pts2 = np.float32([new_left_point, heigth_point, right_point]) # 字符只是高度需要改变 137 | M = cv2.getAffineTransform(pts1, pts2) 138 | dst = cv2.warpAffine(image, M, (round(img_w*2), round(img_h*2))) 139 | car_img = dst[int(right_point[1]):int(heigth_point[1]), int(new_left_point[0]):int(right_point[0])] 140 | 141 | return car_img 142 | 143 | def pre_process(orig_img): 144 | 145 | gray_img = cv2.cvtColor(orig_img, cv2.COLOR_BGR2GRAY) 146 | cv2.imwrite('gray_img.jpg', gray_img) 147 | 148 | blur_img = cv2.blur(gray_img, (3, 3)) 149 | cv2.imwrite('blur.jpg', blur_img) 150 | 151 | sobel_img = cv2.Sobel(blur_img, cv2.CV_16S, 1, 0, ksize=3) 152 | sobel_img = cv2.convertScaleAbs(sobel_img) 153 | cv2.imwrite('sobel.jpg', sobel_img) 154 | 155 | hsv_img = cv2.cvtColor(orig_img, cv2.COLOR_BGR2HSV) 156 | 157 | h, s, v = hsv_img[:, :, 0], hsv_img[:, :, 1], hsv_img[:, :, 2] 158 | # 黄色色调区间[26，34],蓝色色调区间:[100,124] 159 | blue_img = (((h > 26) & (h < 34)) | ((h > 100) & (h < 124))) & (s > 70) & (v > 70) 160 | blue_img = blue_img.astype('float32') 161 | cv2.imwrite('hsv.jpg', blue_img) 162 | 163 | mix_img = np.multiply(sobel_img, blue_img) 164 | cv2.imwrite('mix.jpg', mix_img) 165 | 166 | mix_img = mix_img.astype(np.uint8) 167 | 168 | ret, binary_img = cv2.threshold(mix_img, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU) 169 | cv2.imwrite('binary.jpg',binary_img) 170 | 171 | kernel = cv2.getStructuringElement(cv2.MORPH_RECT,(21,5)) 172 | close_img = cv2.morphologyEx(binary_img, cv2.MORPH_CLOSE, kernel) 173 | cv2.imwrite('close.jpg', close_img) 174 | 175 | return close_img 176 | 177 | # 给候选车牌区域做漫水填充算法，一方面补全上一步求轮廓可能存在轮廓歪曲的问题， 178 | # 另一方面也可以将非车牌区排除掉 179 | def verify_color(rotate_rect,src_image): 180 | img_h,img_w = src_image.shape[:2] 181 | mask = np.zeros(shape=[img_h+2,img_w+2],dtype=np.uint8) 182 | connectivity = 4 183 | loDiff,upDiff = 30,30 184 | new_value = 255 185 | flags = connectivity 186 | flags |= cv2.FLOODFILL_FIXED_RANGE 187 | flags |= new_value << 8 188 | flags |= cv2.FLOODFILL_MASK_ONLY 189 | 190 | rand_seed_num = 5000 191 | valid_seed_num = 200 192 | adjust_param = 0.1 193 | box_points = cv2.boxPoints(rotate_rect) 194 | box_points_x = [n[0] for n in box_points] 195 | box_points_x.sort(reverse=False) 196 | adjust_x = int((box_points_x[2]-box_points_x[1])*adjust_param) 197 | col_range = [box_points_x[1]+adjust_x,box_points_x[2]-adjust_x] 198 | box_points_y = [n[1] for n in box_points] 199 | box_points_y.sort(reverse=False) 200 | adjust_y = int((box_points_y[2]-box_points_y[1])*adjust_param) 201 | row_range = [box_points_y[1]+adjust_y, box_points_y[2]-adjust_y] 202 | 203 | if (col_range[1]-col_range[0])/(box_points_x[3]-box_points_x[0])<0.4\ 204 | or (row_range[1]-row_range[0])/(box_points_y[3]-box_points_y[0])<0.4: 205 | points_row = [] 206 | points_col = [] 207 | for i in range(2): 208 | pt1,pt2 = box_points[i],box_points[i+2] 209 | x_adjust,y_adjust = int(adjust_param*(abs(pt1[0]-pt2[0]))),int(adjust_param*(abs(pt1[1]-pt2[1]))) 210 | if (pt1[0] <= pt2[0]): 211 | pt1[0], pt2[0] = pt1[0] + x_adjust, pt2[0] - x_adjust 212 | else: 213 | pt1[0], pt2[0] = pt1[0] - x_adjust, pt2[0] + x_adjust 214 | if (pt1[1] <= pt2[1]): 215 | pt1[1], pt2[1] = pt1[1] + adjust_y, pt2[1] - adjust_y 216 | else: 217 | pt1[1], pt2[1] = pt1[1] - y_adjust, pt2[1] + y_adjust 218 | temp_list_x = [int(x) for x in np.linspace(pt1[0],pt2[0],int(rand_seed_num /2))] 219 | temp_list_y = [int(y) for y in np.linspace(pt1[1],pt2[1],int(rand_seed_num /2))] 220 | points_col.extend(temp_list_x) 221 | points_row.extend(temp_list_y) 222 | else: 223 | points_row = np.random.randint(row_range[0],row_range[1],size=rand_seed_num) 224 | points_col = np.linspace(col_range[0],col_range[1],num=rand_seed_num).astype(np.int) 225 | 226 | points_row = np.array(points_row) 227 | points_col = np.array(points_col) 228 | hsv_img = cv2.cvtColor(src_image, cv2.COLOR_BGR2HSV) 229 | h,s,v = hsv_img[:,:,0],hsv_img[:,:,1],hsv_img[:,:,2] 230 | 231 | flood_img = src_image.copy() 232 | seed_cnt = 0 233 | for i in range(rand_seed_num): 234 | rand_index = np.random.choice(rand_seed_num,1,replace=False) 235 | row,col = points_row[rand_index],points_col[rand_index] 236 | 237 | if (((h[row,col]>26)&(h[row,col]<34))|((h[row,col]>100)&(h[row,col]<124)))&(s[row,col]>70)&(v[row,col]>70): 238 | cv2.floodFill(src_image, mask, (col,row), (255, 255, 255), (loDiff,) * 3, (upDiff,) * 3, flags) 239 | cv2.circle(flood_img,center=(col,row),radius=2,color=(0,0,255),thickness=2) 240 | seed_cnt += 1 241 | if seed_cnt >= valid_seed_num: 242 | break 243 | 244 | show_seed = np.random.uniform(1, 100, 1).astype(np.uint16) 245 | cv2.imwrite('floodfill.jpg',flood_img) 246 | cv2.imwrite('flood_mask.jpg',mask) 247 | 248 | mask_points = [] 249 | for row in range(1,img_h+1): 250 | for col in range(1,img_w+1): 251 | if mask[row,col] != 0: 252 | mask_points.append((col-1,row-1)) 253 | mask_rotateRect = cv2.minAreaRect(np.array(mask_points)) 254 | if verify_scale(mask_rotateRect): 255 | return True,mask_rotateRect 256 | else: 257 | return False,mask_rotateRect 258 | 259 | # 车牌定位 260 | def locate_carPlate(orig_img,pred_image): 261 | carPlate_list = [] 262 | temp1_orig_img = orig_img.copy() #调试用 263 | temp2_orig_img = orig_img.copy() #调试用 264 | contours, heriachy = cv2.findContours(pred_image,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE) 265 | for i,contour in enumerate(contours): 266 | cv2.drawContours(temp1_orig_img, contours, i, (0, 255, 255), 2) 267 | # 获取轮廓最小外接矩形，返回值rotate_rect 268 | rotate_rect = cv2.minAreaRect(contour) 269 | # 根据矩形面积大小和长宽比判断是否是车牌 270 | if verify_scale(rotate_rect): 271 | ret,rotate_rect2 = verify_color(rotate_rect,temp2_orig_img) 272 | if ret == False: 273 | continue 274 | # 车牌位置矫正 275 | car_plate = img_Transform(rotate_rect2, temp2_orig_img) 276 | car_plate = cv2.resize(car_plate,(car_plate_w,car_plate_h)) #调整尺寸为后面CNN车牌识别做准备 277 | #========================调试看效果========================# 278 | box = cv2.boxPoints(rotate_rect2) 279 | for k in range(4): 280 | n1,n2 = k%4,(k+1)%4 281 | cv2.line(temp1_orig_img, (box[n1][0], box[n1][1]),(box[n2][0], box[n2][1]), (255,0,0), 2) 282 | cv2.imwrite('opencv.jpg', car_plate) 283 | #========================调试看效果========================# 284 | carPlate_list.append(car_plate) 285 | 286 | cv2.imwrite('contour.jpg', temp1_orig_img) 287 | return carPlate_list 288 | 289 | # 左右切割 290 | def horizontal_cut_chars(plate): 291 | char_addr_list = [] 292 | area_left,area_right,char_left,char_right= 0,0,0,0 293 | img_w = plate.shape[1] 294 | 295 | # 获取车牌每列边缘像素点个数 296 | def getColSum(img,col): 297 | sum = 0 298 | for i in range(img.shape[0]): 299 | sum += round(img[i,col]/255) 300 | return sum; 301 | 302 | sum = 0 303 | for col in range(img_w): 304 | sum += getColSum(plate,col) 305 | # 每列边缘像素点必须超过均值的60%才能判断属于字符区域 306 | col_limit = 0#round(0.5*sum/img_w) 307 | # 每个字符宽度也进行限制 308 | charWid_limit = [round(img_w/12),round(img_w/5)] 309 | is_char_flag = False 310 | 311 | for i in range(img_w): 312 | colValue = getColSum(plate,i) 313 | if colValue > col_limit: 314 | if is_char_flag == False: 315 | area_right = round((i+char_right)/2) 316 | area_width = area_right-area_left 317 | char_width = char_right-char_left 318 | if (area_width>charWid_limit[0]) and (area_width charWid_limit[0]) and (area_width < charWid_limit[1]): 333 | char_addr_list.append((area_left, area_right, char_width)) 334 | return char_addr_list 335 | 336 | def get_chars(car_plate): 337 | img_h,img_w = car_plate.shape[:2] 338 | h_proj_list = [] # 水平投影长度列表 339 | h_temp_len,v_temp_len = 0,0 340 | h_startIndex,h_end_index = 0,0 # 水平投影记索引 341 | h_proj_limit = [0.2,0.8] # 车牌在水平方向得轮廓长度少于20%或多余80%过滤掉 342 | char_imgs = [] 343 | 344 | # 将二值化的车牌水平投影到Y轴，计算投影后的连续长度，连续投影长度可能不止一段 345 | h_count = [0 for i in range(img_h)] 346 | for row in range(img_h): 347 | temp_cnt = 0 348 | for col in range(img_w): 349 | if car_plate[row,col] == 255: 350 | temp_cnt += 1 351 | h_count[row] = temp_cnt 352 | if temp_cnt/img_wh_proj_limit[1]: 353 | if h_temp_len != 0: 354 | h_end_index = row-1 355 | h_proj_list.append((h_startIndex,h_end_index)) 356 | h_temp_len = 0 357 | continue 358 | if temp_cnt > 0: 359 | if h_temp_len == 0: 360 | h_startIndex = row 361 | h_temp_len = 1 362 | else: 363 | h_temp_len += 1 364 | else: 365 | if h_temp_len > 0: 366 | h_end_index = row-1 367 | h_proj_list.append((h_startIndex,h_end_index)) 368 | h_temp_len = 0 369 | 370 | # 手动结束最后得水平投影长度累加 371 | if h_temp_len != 0: 372 | h_end_index = img_h-1 373 | h_proj_list.append((h_startIndex, h_end_index)) 374 | # 选出最长的投影，该投影长度占整个截取车牌高度的比值必须大于0.5 375 | h_maxIndex,h_maxHeight = 0,0 376 | for i,(start,end) in enumerate(h_proj_list): 377 | if h_maxHeight < (end-start): 378 | h_maxHeight = (end-start) 379 | h_maxIndex = i 380 | if h_maxHeight/img_h < 0.5: 381 | return char_imgs 382 | chars_top,chars_bottom = h_proj_list[h_maxIndex][0],h_proj_list[h_maxIndex][1] 383 | 384 | plates = car_plate[chars_top:chars_bottom+1,:] 385 | cv2.imwrite('car.jpg', car_plate) 386 | cv2.imwrite('plate.jpg', plates) 387 | char_addr_list = horizontal_cut_chars(plates) 388 | 389 | for i,addr in enumerate(char_addr_list): 390 | char_img = car_plate[chars_top:chars_bottom+1,addr[0]:addr[1]] 391 | char_img = cv2.resize(char_img,(char_w,char_h)) 392 | char_imgs.append(char_img) 393 | return char_imgs 394 | 395 | def extract_char(car_plate): 396 | gray_plate = cv2.cvtColor(car_plate, cv2.COLOR_BGR2GRAY) 397 | ret,binary_plate = cv2.threshold(gray_plate, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU) 398 | char_img_list = get_chars(binary_plate) 399 | return char_img_list 400 | 401 | def cnn_select_carPlate(plate_list, model_path): 402 | if len(plate_list) == 0: 403 | return False, plate_list 404 | model = torch.load(model_path) 405 | model = model.cuda() 406 | 407 | idx, val = 0, -1 408 | tf = transforms.ToTensor() 409 | for i in range(len(plate_list)): 410 | with torch.no_grad(): 411 | input = tf(np.array(plate_list[i])).unsqueeze(0).cuda() 412 | output = torch.sigmoid(model(input)) 413 | 414 | if output > val: 415 | idx, val = i, output 416 | return True, plate_list[idx] 417 | 418 | def cnn_recongnize_char(img_list, model_path): 419 | model = torch.load(model_path) 420 | model = model.cuda() 421 | text_list = [] 422 | 423 | if len(img_list) == 0: 424 | return text_list 425 | 426 | tf = transforms.ToTensor() 427 | for img in img_list: 428 | input = tf(np.array(img)).unsqueeze(0).cuda() 429 | # 数字、字母、汉字，从67维向量找到概率最大的作为预测结果 430 | with torch.no_grad(): 431 | output = model(input) 432 | _, preds = torch.topk(output, 1) 433 | text_list.append(char_table[preds]) 434 | 435 | return text_list 436 | 437 | def list_all_files(root): 438 | files = [] 439 | list = os.listdir(root) 440 | for i in range(len(list)): 441 | element = os.path.join(root, list[i]) 442 | if os.path.isdir(element): 443 | files.extend(list_all_files(element)) 444 | elif os.path.isfile(element): 445 | files.append(element) 446 | return files 447 | 448 | if __name__ == '__main__': 449 | car_plate_w,car_plate_h = 136,36 450 | char_w,char_h = 20, 20 451 | plate_model_path = "plate.pth" 452 | char_model_path = "char.pth" 453 | root = '../images/test/' # 测试图片路径 454 | files = list_all_files(root) 455 | files.sort() 456 | 457 | for file in files: 458 | print(file) 459 | img = cv2.imread(file) 460 | if len(img) < 2: 461 | continue 462 | 463 | pred_img = pre_process(img) 464 | car_plate_list = locate_carPlate(img, pred_img) 465 | ret, car_plate = cnn_select_carPlate(car_plate_list,plate_model_path) 466 | if ret == False: 467 | print("未检测到车牌") 468 | continue 469 | cv2.imwrite('cnn_plate.jpg',car_plate) 470 | 471 | char_img_list = extract_char(car_plate) 472 | for idx in range(len(char_img_list)): 473 | img_name = 'char-' + str(idx) + '.jpg' 474 | cv2.imwrite(img_name, char_img_list[idx]) 475 | 476 | text = cnn_recongnize_char(char_img_list,char_model_path) 477 | print(text) 478 | 479 | cv2.waitKey(0) -------------------------------------------------------------------------------- /code/char.pth: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mligg23/CarPlateIdentity/d1cb956ac91d1415e3d267d845603e2ada718e89/code/char.pth -------------------------------------------------------------------------------- /code/charNeuralNet.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | from torch.utils.data import dataloader 4 | import tqdm 5 | import numpy as np 6 | import cv2 7 | import torch 8 | import torch.nn as nn 9 | from torch.autograd import Variable 10 | import torch.utils.data as data 11 | from torch.utils.data import DataLoader 12 | from torchvision import transforms 13 | 14 | torch.cuda.set_device(7) 15 | batch_size = 1 16 | 17 | numbers = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9'] 18 | alphbets = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 19 | 'U', 'V', 'W', 'X', 'Y', 'Z'] 20 | chinese = ['zh_cuan', 'zh_e', 'zh_gan', 'zh_gan1', 'zh_gui', 'zh_gui1', 'zh_hei', 'zh_hu', 'zh_ji', 'zh_jin', 21 | 'zh_jing', 'zh_jl', 'zh_liao', 'zh_lu', 'zh_meng', 'zh_min', 'zh_ning', 'zh_qing', 'zh_qiong', 22 | 'zh_shan', 'zh_su', 'zh_sx', 'zh_wan', 'zh_xiang', 'zh_xin', 'zh_yu', 'zh_yu1', 'zh_yue', 'zh_yun', 23 | 'zh_zang', 'zh_zhe'] 24 | 25 | 26 | class char_cnn_net(nn.Module): 27 | def __init__(self): 28 | super().__init__() 29 | 30 | self.conv = nn.Sequential( 31 | nn.Conv2d(1,64,3,1,1), 32 | nn.PReLU(), 33 | nn.Conv2d(64,16,3,1,1), 34 | nn.PReLU(), 35 | nn.Conv2d(16,4,3,1,1), 36 | nn.PReLU() 37 | ) 38 | 39 | self.fc = nn.Sequential( 40 | nn.Linear(1600, 512), 41 | nn.PReLU(), 42 | nn.Linear(512, 256), 43 | nn.PReLU(), 44 | nn.Linear(256,67) 45 | ) 46 | 47 | def forward(self, x): 48 | y = self.conv(x).reshape(batch_size, -1,) 49 | # print(y.shape) 50 | return self.fc(y) 51 | 52 | class CharPic(data.Dataset): 53 | 54 | def list_all_files(self, root): 55 | files = [] 56 | list = os.listdir(root) 57 | for i in range(len(list)): 58 | element = os.path.join(root, list[i]) 59 | if os.path.isdir(element): 60 | files.extend(self.list_all_files(element)) 61 | elif os.path.isfile(element): 62 | files.append(element) 63 | return files 64 | 65 | def __init__(self, root): 66 | super().__init__() 67 | if not os.path.exists(root): 68 | raise ValueError('没有找到文件夹') 69 | files = self.list_all_files(root) 70 | 71 | self.X = [] 72 | self.y = [] 73 | self.dataset = numbers + alphbets + chinese 74 | 75 | for file in files: 76 | src_img = cv2.imread(file, cv2.COLOR_BGR2GRAY) 77 | if src_img.ndim == 3: 78 | continue 79 | resize_img = cv2.resize(src_img, (20, 20)) 80 | self.X.append(resize_img) 81 | 82 | dir = os.path.dirname(file) 83 | dir_name = os.path.split(dir)[-1] 84 | 85 | # vector_y = [0 for i in range(len(self.dataset))] 86 | index_y = self.dataset.index(dir_name) 87 | # vector_y[index_y] = 1 88 | self.y.append([index_y]) 89 | 90 | self.X = np.array(self.X) 91 | self.y = np.array(self.y) 92 | 93 | def __getitem__(self, index): 94 | tf = transforms.ToTensor() 95 | # print(torch.Tensor(self.y[index]).shape) 96 | return tf(self.X[index]), torch.LongTensor(self.y[index]) 97 | 98 | def __len__(self) -> int: 99 | return len(self.X) 100 | 101 | def weights_init(m): 102 | classname = m.__class__.__name__ 103 | if classname.find('Conv') != -1: 104 | m.weight.data.normal_(mean=0.0, std=0.1) 105 | m.bias.data.fill_(0) 106 | 107 | def train(epoch, lr): 108 | model.train() 109 | 110 | criterion = nn.CrossEntropyLoss() 111 | loss_history = [] 112 | 113 | for batch_idx, (input, target) in enumerate(train_loader): 114 | input, target = input.cuda(), target.cuda() 115 | input, target = Variable(input), Variable(target).reshape(batch_size, ) 116 | 117 | optimizer = torch.optim.Adam(model.parameters(), lr=lr) 118 | optimizer.zero_grad() 119 | 120 | output = model(input) 121 | 122 | loss = criterion(output, target) 123 | loss.backward() 124 | 125 | if loss_history and loss_history[-1] < loss.data: 126 | lr *= 0.95 127 | loss_history.append(loss.data) 128 | 129 | optimizer.step() 130 | 131 | if batch_idx % 12000 == 0: 132 | print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( 133 | epoch, batch_idx * len(input), len(train_loader.dataset), 134 | 100. * batch_idx / len(train_loader), loss.data)) 135 | 136 | def get_accuracy(model, train_model_path): 137 | tot = len(train_loader.dataset) 138 | right = 0 139 | 140 | with torch.no_grad(): 141 | for (input, target) in train_loader: 142 | input, target = input.cuda(), target.cuda() 143 | output = model(input) 144 | 145 | for idx in range(len(output)): 146 | _, predict = torch.topk(output[idx], 1) 147 | if predict == target[idx]: 148 | right += 1 149 | 150 | acc = right / tot 151 | print('accuracy : %.3f' % acc) 152 | 153 | global best_acc 154 | if acc > best_acc: 155 | best_acc = acc 156 | torch.save(model, train_model_path) 157 | 158 | 159 | if __name__ == '__main__': 160 | data_dir = '../images/cnn_char_train' 161 | train_model_path = 'char.pth' 162 | 163 | # model = char_cnn_net() 164 | model = torch.load(train_model_path) 165 | model = model.cuda() 166 | # model.apply(weights_init) 167 | 168 | print("Generate Model.") 169 | 170 | batch_size = 1 171 | dataset = CharPic(data_dir) 172 | train_loader = DataLoader(dataset=dataset, shuffle=True, batch_size=batch_size, 173 | num_workers=14, pin_memory=True, drop_last=True) 174 | 175 | global best_acc 176 | best_acc = 0.0 177 | for epoch in range(800): 178 | lr = 0.001 179 | train(epoch, lr) 180 | get_accuracy(model, train_model_path) 181 | 182 | torch.save(model, train_model_path) 183 | 184 | print("Finish Training") 185 | -------------------------------------------------------------------------------- /code/plate.pth: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mligg23/CarPlateIdentity/d1cb956ac91d1415e3d267d845603e2ada718e89/code/plate.pth -------------------------------------------------------------------------------- /code/plateNeuralNet.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | from torch.utils.data import dataloader 4 | import tqdm 5 | import numpy as np 6 | import cv2 7 | import torch 8 | import torch.nn as nn 9 | from torch.autograd import Variable 10 | import torch.utils.data as data 11 | from torch.utils.data import DataLoader 12 | from torchvision import transforms 13 | 14 | torch.cuda.set_device(7) 15 | batch_size = 1 16 | 17 | class plate_cnn_net(nn.Module): 18 | def __init__(self): 19 | super().__init__() 20 | 21 | self.conv = nn.Sequential( 22 | nn.Conv2d(3,64,3,1,1), 23 | nn.PReLU(), 24 | nn.Conv2d(64,128,3,2,1), 25 | nn.PReLU(), 26 | nn.Conv2d(128,128,3,2,1), 27 | nn.PReLU(), 28 | nn.Conv2d(128,64,3,2,1), 29 | nn.PReLU(), 30 | nn.Conv2d(64,16,3,1,1), 31 | nn.PReLU(), 32 | nn.Conv2d(16,4,3,1,1), 33 | nn.PReLU() 34 | ) 35 | 36 | self.fc = nn.Sequential( 37 | nn.Linear(340, 256), 38 | nn.PReLU(), 39 | nn.Linear(256, 128), 40 | nn.PReLU(), 41 | nn.Linear(128, 32), 42 | nn.PReLU(), 43 | nn.Linear(32, 8), 44 | nn.PReLU(), 45 | nn.Linear(8, 1) 46 | ) 47 | 48 | def forward(self, x): 49 | y = self.conv(x).reshape(batch_size, -1,) 50 | # print(y.shape) 51 | return self.fc(y) 52 | 53 | class PlatePic(data.Dataset): 54 | 55 | def list_all_files(self, root): 56 | files = [] 57 | list = os.listdir(root) 58 | for i in range(len(list)): 59 | element = os.path.join(root, list[i]) 60 | if os.path.isdir(element): 61 | files.extend(self.list_all_files(element)) 62 | elif os.path.isfile(element): 63 | files.append(element) 64 | return files 65 | 66 | def __init__(self, root): 67 | super().__init__() 68 | if not os.path.exists(root): 69 | raise ValueError('没有找到文件夹') 70 | self.files = self.list_all_files(root) 71 | 72 | self.X = [] 73 | self.y = [] 74 | self.labels = [os.path.split(os.path.dirname(file))[-1] for file in self.files] 75 | 76 | for i, file in enumerate(self.files): 77 | src_img = cv2.imread(file) 78 | if src_img.ndim != 3: 79 | continue 80 | resize_img = cv2.resize(src_img, (136, 36)) 81 | self.X.append(resize_img) 82 | self.y.append([0 if self.labels[i] == 'no' else 1]) 83 | 84 | self.X = np.array(self.X) 85 | self.y = np.array(self.y) 86 | 87 | def __getitem__(self, index): 88 | tf = transforms.ToTensor() 89 | # print(torch.Tensor(self.y[index]).shape) 90 | return tf(self.X[index]), torch.FloatTensor(self.y[index]) 91 | 92 | def __len__(self) -> int: 93 | return len(self.X) 94 | 95 | def weights_init(m): 96 | classname = m.__class__.__name__ 97 | if classname.find('Conv') != -1: 98 | m.weight.data.normal_(mean=0.0, std=0.1) 99 | m.bias.data.fill_(0) 100 | 101 | def train(epoch, lr): 102 | model.train() 103 | 104 | criterion = nn.BCEWithLogitsLoss() 105 | loss_history = [] 106 | 107 | for batch_idx, (input, target) in enumerate(train_loader): 108 | input, target = input.cuda(), target.cuda() 109 | input, target = Variable(input), Variable(target) 110 | 111 | optimizer = torch.optim.Adam(model.parameters(), lr=lr) 112 | optimizer.zero_grad() 113 | 114 | output = model(input) 115 | loss = criterion(output, target) 116 | loss.backward() 117 | 118 | if loss_history and loss_history[-1] < loss.data: 119 | lr *= 0.7 120 | loss_history.append(loss.data) 121 | 122 | optimizer.step() 123 | 124 | if batch_idx % 2000 == 0: 125 | print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( 126 | epoch, batch_idx * len(input), len(train_loader.dataset), 127 | 100. * batch_idx / len(train_loader), loss.data)) 128 | 129 | def get_accuracy(model, train_model_path): 130 | tot = len(train_loader.dataset) 131 | right = 0 132 | 133 | with torch.no_grad(): 134 | for (input, target) in train_loader: 135 | input, target = input.cuda(), target.cuda() 136 | output = model(input) 137 | 138 | for idx in range(len(output)): 139 | if (output[idx] > 0.5 and target[idx] > 0.5) or \ 140 | (output[idx] < 0.5 and target[idx] < 0.5): 141 | right += 1 142 | 143 | acc = right / tot 144 | print('accuracy : %.3f' % acc) 145 | 146 | global best_acc 147 | if acc > best_acc: 148 | best_acc = acc 149 | torch.save(model, train_model_path) 150 | 151 | if __name__ == '__main__': 152 | data_dir = '../images/cnn_plate_train' 153 | train_model_path = 'plate.pth' 154 | 155 | # model = plate_cnn_net() 156 | model = torch.load(train_model_path) 157 | model = model.cuda() 158 | # model.apply(weights_init) 159 | 160 | print("Generate Model.") 161 | 162 | batch_size = 1 163 | dataset = PlatePic(data_dir) 164 | train_loader = DataLoader(dataset=dataset, shuffle=True, batch_size=batch_size, 165 | num_workers=14, pin_memory=True, drop_last=True) 166 | 167 | global best_acc 168 | best_acc = 0.0 169 | for epoch in range(0, 30): 170 | lr = 0.001 171 | train(epoch, lr) 172 | get_accuracy(model, train_model_path) 173 | 174 | torch.save(model, train_model_path) 175 | 176 | print("Finish Training") 177 | -------------------------------------------------------------------------------- /images.zip: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mligg23/CarPlateIdentity/d1cb956ac91d1415e3d267d845603e2ada718e89/images.zip --------------------------------------------------------------------------------