mirror of
https://github.com/we0091234/yolov7_plate.git
synced 2025-09-26 21:01:13 +08:00
95 lines
4.2 KiB
Python
95 lines
4.2 KiB
Python
from detect_rec_plate_macao import detect_Recognition_plate,attempt_load,init_model,allFilePath,cv_imread,draw_result,four_point_transform,order_points
|
|
import os
|
|
import argparse
|
|
import torch
|
|
import cv2
|
|
import time
|
|
import shutil
|
|
import numpy as np
|
|
import json
|
|
import re
|
|
pattern_str = "([京津沪渝冀豫云辽黑湘皖鲁新苏浙赣鄂桂甘晋蒙陕吉闽贵粤青藏川宁琼]" \
|
|
"{1}(([A-HJ-Z]{1}[A-HJ-NP-Z0-9]{5})|([A-HJ-Z]{1}(([DF]{1}[A-HJ-NP-Z0-9]{1}[0-9]{4})|([0-9]{5}[DF]" \
|
|
"{1})))|([A-HJ-Z]{1}[A-D0-9]{1}[0-9]{3}警)))|([0-9]{6}使)|((([沪粤川云桂鄂陕蒙藏黑辽渝]{1}A)|鲁B|闽D|蒙E|蒙H)" \
|
|
"[0-9]{4}领)|(WJ[京津沪渝冀豫云辽黑湘皖鲁新苏浙赣鄂桂甘晋蒙陕吉闽贵粤青藏川宁琼·•]{1}[0-9]{4}[TDSHBXJ0-9]{1})" \
|
|
"|([VKHBSLJNGCE]{1}[A-DJ-PR-TVY]{1}[0-9]{5})"
|
|
|
|
|
|
def is_car_number(pattern, string):
|
|
if re.findall(pattern, string):
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--detect_model', nargs='+', type=str, default=r'runs/train/yolov711/weights/best.pt', help='model.pt path(s)')
|
|
parser.add_argument('--rec_model', type=str, default=r'weights/plate_rec.pth', help='model.pt path(s)')
|
|
parser.add_argument('--source', type=str, default=r'/mnt/Gpan/Mydata/pytorchPorject/datasets/macao_plate/download/', help='source') # file/folder, 0 for webcam
|
|
# parser.add_argument('--source', type=str, default=r'test2', help='source')
|
|
# parser.add_argument('--img-size', nargs= '+', type=int, default=640, help='inference size (pixels)')
|
|
parser.add_argument('--img_size', type=int, default=640, help='inference size (pixels)')
|
|
parser.add_argument('--output', type=str, default=r'/mnt/Gpan/Mydata/pytorchPorject/datasets/macao_plate/result/', help='source')
|
|
parser.add_argument('--kpt-label', type=int, default=4, help='number of keypoints')
|
|
device =torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
# device = torch.device("cpu")
|
|
opt = parser.parse_args()
|
|
print(opt)
|
|
model = attempt_load(opt.detect_model, map_location=device)
|
|
# torch.save()
|
|
plate_rec_model=init_model(device,opt.rec_model)
|
|
if not os.path.exists(opt.output):
|
|
os.mkdir(opt.output)
|
|
file_list=[]
|
|
index_1=0
|
|
index_small=0
|
|
# error_path =r"E:\study\plate\data\@shaixuan\train\pic"
|
|
allFilePath(opt.source,file_list)
|
|
time_b = time.time()
|
|
|
|
for pic_ in file_list:
|
|
try:
|
|
image_name = os.path.basename(pic_)
|
|
# ori_plate = image_name.split("_")[0]
|
|
# ori_plate = image_name.split(".")[0]
|
|
flag = 1
|
|
index_1+=1
|
|
label_dict_str=""
|
|
lable_dict={}
|
|
# label_dict={}
|
|
print(index_small,index_1,pic_)
|
|
img = cv_imread(pic_)
|
|
if img is None:
|
|
continue
|
|
if img.shape[-1]==4:
|
|
img=cv2.cvtColor(img,cv2.COLOR_BGRA2BGR)
|
|
# img = my_letter_box(img)
|
|
count=0
|
|
dict_list=detect_Recognition_plate(model, img, device,plate_rec_model,opt.img_size)
|
|
for result_ in dict_list:
|
|
if not result_:
|
|
continue
|
|
index_small+=1
|
|
landmarks=result_['landmarks']
|
|
landmarks_np = np.array(landmarks).reshape(-1,2)
|
|
img_roi = four_point_transform(img,landmarks_np)
|
|
plate_no= result_['plate_no']
|
|
# if len(plate_no)<6 or not is_car_number(pattern_str,plate_no):
|
|
if len(plate_no)<6:
|
|
continue
|
|
height = result_['roi_height']
|
|
# if height<48:
|
|
# continue
|
|
pic_name =plate_no+"_"+str(index_small)+".jpg"
|
|
label = result_["label"]
|
|
|
|
|
|
roi_save_folder = os.path.join(opt.output,str(label))
|
|
if not os.path.exists(roi_save_folder):
|
|
os.mkdir(roi_save_folder)
|
|
roi_path = os.path.join(roi_save_folder,pic_name)
|
|
cv2.imencode('.jpg',img_roi)[1].tofile(roi_path)
|
|
except:
|
|
print("error!")
|
|
|
|
|