301 lines
12 KiB
Python
301 lines
12 KiB
Python
import numpy as np
|
||
import cv2, os, argparse
|
||
from glob import glob
|
||
from tqdm import tqdm
|
||
|
||
from plate_number import random_select, generate_plate_number_white, generate_plate_number_army, generate_plate_number_yellow_xue, generate_plate_number_yellow
|
||
from plate_number import generate_plate_number_black_gangao
|
||
from plate_number import generate_plate_number_black_shi_old, generate_plate_number_black_shi_new
|
||
from plate_number import generate_plate_number_black_ling_old, generate_plate_number_black_ling_new
|
||
from plate_number import generate_plate_number_blue
|
||
from plate_number import letters, digits
|
||
|
||
|
||
def get_location_data(length=7, split_id=1, height=140):
|
||
"""
|
||
获取车牌号码在底牌中的位置
|
||
length: 车牌字符数，7或者8，7为普通车牌、8为新能源车牌
|
||
split_id: 分割空隙
|
||
height: 车牌高度，对应单层和双层车牌
|
||
"""
|
||
# 字符位置
|
||
location_xy = np.zeros((length, 4), dtype=np.int32)
|
||
|
||
# 单层车牌高度
|
||
if height == 140:
|
||
# 单层车牌，y轴坐标固定
|
||
location_xy[:, 1] = 25
|
||
location_xy[:, 3] = 115
|
||
# 螺栓间隔
|
||
step_split = 34 if length == 7 else 49
|
||
# 字符间隔
|
||
step_font = 12 if length == 7 else 9
|
||
|
||
# 字符宽度
|
||
width_font = 45
|
||
for i in range(length):
|
||
if i == 0:
|
||
location_xy[i, 0] = 15
|
||
elif i == split_id:
|
||
location_xy[i, 0] = location_xy[i - 1, 2] + step_split
|
||
else:
|
||
location_xy[i, 0] = location_xy[i - 1, 2] + step_font
|
||
# 新能源车牌
|
||
if length == 8 and i > 0:
|
||
width_font = 43
|
||
location_xy[i, 2] = location_xy[i, 0] + width_font
|
||
else:
|
||
# 双层车牌第一层
|
||
location_xy[0, :] = [110, 15, 190, 75]
|
||
location_xy[1, :] = [250, 15, 330, 75]
|
||
|
||
# 第二层
|
||
width_font = 65
|
||
step_font = 15
|
||
for i in range(2, length):
|
||
location_xy[i, 1] = 90
|
||
location_xy[i, 3] = 200
|
||
if i == 2:
|
||
location_xy[i, 0] = 27
|
||
else:
|
||
location_xy[i, 0] = location_xy[i - 1, 2] + step_font
|
||
location_xy[i, 2] = location_xy[i, 0] + width_font
|
||
|
||
return location_xy
|
||
|
||
|
||
# 字符贴上底板
|
||
def copy_to_image_multi(img, font_img, bbox, bg_color, is_red):
|
||
x1, y1, x2, y2 = bbox
|
||
font_img = cv2.resize(font_img, (x2 - x1, y2 - y1))
|
||
img_crop = img[y1: y2, x1: x2, :]
|
||
|
||
if is_red:
|
||
img_crop[font_img < 200, :] = [0, 0, 255]
|
||
elif 'blue' in bg_color or 'black' in bg_color:
|
||
img_crop[font_img < 200, :] = [255, 255, 255]
|
||
else:
|
||
img_crop[font_img < 200, :] = [0, 0, 0]
|
||
return img
|
||
|
||
|
||
class MultiPlateGenerator:
|
||
def __init__(self, adr_plate_model, adr_font):
|
||
# 车牌底板路径
|
||
self.adr_plate_model = adr_plate_model
|
||
# 车牌字符路径
|
||
self.adr_font = adr_font
|
||
|
||
# 车牌字符图片，预存处理
|
||
self.font_imgs = {}
|
||
font_filenames = glob(os.path.join(adr_font, '*jpg'))
|
||
for font_filename in font_filenames:
|
||
font_img = cv2.imdecode(np.fromfile(font_filename, dtype=np.uint8), 0)
|
||
|
||
if '140' in font_filename:
|
||
font_img = cv2.resize(font_img, (45, 90))
|
||
elif '220' in font_filename:
|
||
font_img = cv2.resize(font_img, (65, 110))
|
||
elif font_filename.split('_')[-1].split('.')[0] in letters + digits:
|
||
font_img = cv2.resize(font_img, (43, 90))
|
||
# 仅对"1"膨胀加粗（字为暗色，需反转→膨胀→反转回）
|
||
key = os.path.basename(font_filename).split('.')[0]
|
||
if key.endswith('_1') or key == '1':
|
||
font_img = 255 - font_img
|
||
font_img = cv2.dilate(font_img, np.ones((3, 3), np.uint8), iterations=2)
|
||
font_img = 255 - font_img
|
||
self.font_imgs[key] = font_img
|
||
|
||
# 字符位置
|
||
self.location_xys = {}
|
||
for i in [7, 8]:
|
||
for j in [1, 2, 3, 4]:
|
||
for k in [140, 220]:
|
||
self.location_xys['{}_{}_{}'.format(i, j, k)] = \
|
||
get_location_data(length=i, split_id=j, height=k)
|
||
|
||
# 获取字符位置
|
||
def get_location_multi(self, plate_number, height=140):
|
||
length = len(plate_number)
|
||
if '警' in plate_number:
|
||
split_id = 1
|
||
elif '使' in plate_number:
|
||
# 旧版：使在首位 → 间隙在第4位前；新版：使在末位 → 间隙在第3位前
|
||
split_id = 4 if plate_number[0] == '使' else 3
|
||
else:
|
||
split_id = 2
|
||
return self.location_xys['{}_{}_{}'.format(length, split_id, height)]
|
||
|
||
# 随机生成车牌号码（黄/白/黑 + 大型新能源）
|
||
def generate_plate_number(self):
|
||
# 白牌警车/军车、黄牌教练车/大型车、港澳、使领馆新旧版、大型新能源
|
||
generate_plate_number_funcs = [
|
||
generate_plate_number_white,
|
||
generate_plate_number_army,
|
||
generate_plate_number_yellow_xue,
|
||
generate_plate_number_yellow,
|
||
generate_plate_number_black_gangao,
|
||
generate_plate_number_black_shi_old,
|
||
generate_plate_number_black_shi_new,
|
||
generate_plate_number_black_ling_old,
|
||
generate_plate_number_black_ling_new,
|
||
lambda: generate_plate_number_blue(length=8), # 大型新能源
|
||
]
|
||
plate_number = random_select(generate_plate_number_funcs)()
|
||
|
||
# 按GA36规范确定底板颜色
|
||
if '港' in plate_number or '澳' in plate_number:
|
||
bg_color = 'black'
|
||
elif '领' in plate_number:
|
||
bg_color = 'black'
|
||
elif '使' in plate_number:
|
||
bg_color = 'black_shi'
|
||
elif '警' in plate_number:
|
||
bg_color = 'white'
|
||
elif plate_number[0] in letters:
|
||
bg_color = 'white_army'
|
||
elif '学' in plate_number:
|
||
bg_color = 'yellow'
|
||
else:
|
||
bg_color = 'yellow' # 大型车
|
||
|
||
# 大型新能源（8位）
|
||
if len(plate_number) == 8:
|
||
bg_color = 'green_truck'
|
||
|
||
return plate_number, bg_color, False
|
||
|
||
# 随机生成车牌图片（全部单层，已移除双层挂车）
|
||
def generate_plate(self, enhance=False):
|
||
plate_number, bg_color, _ = self.generate_plate_number()
|
||
height = 140
|
||
|
||
# 获取底板图片
|
||
number_xy = self.get_location_multi(plate_number, height)
|
||
img_plate_model = cv2.imread(os.path.join(self.adr_plate_model, '{}_{}.PNG'.format(bg_color, height)))
|
||
img_plate_model = cv2.resize(img_plate_model, (440 if len(plate_number) == 7 else 480, height))
|
||
|
||
for i in range(len(plate_number)):
|
||
if len(plate_number) == 8:
|
||
# 新能源（green_car / green_truck）
|
||
font_img = self.font_imgs['green_{}'.format(plate_number[i])]
|
||
else:
|
||
if '{}_{}'.format(height, plate_number[i]) in self.font_imgs:
|
||
font_img = self.font_imgs['{}_{}'.format(height, plate_number[i])]
|
||
else:
|
||
# 兜底：220字体缩放到140
|
||
font_img = self.font_imgs.get('220_{}'.format(plate_number[i]),
|
||
self.font_imgs.get('220_up_{}'.format(plate_number[i]),
|
||
self.font_imgs['220_down_{}'.format(plate_number[i])]))
|
||
|
||
# 字符是否红色
|
||
# 军车首字母、警、旧版使馆(使在首位)、旧版领馆(省A·XXXX领)为红色
|
||
is_red = False
|
||
if (i == 0 and plate_number[0] in letters) or plate_number[i] == '警':
|
||
is_red = True
|
||
elif plate_number[i] == '使' and i == 0:
|
||
is_red = True # 旧版使馆：使在首位 → 红字
|
||
elif plate_number[i] == '领' and len(plate_number) >= 2 and plate_number[1] in letters:
|
||
is_red = True # 旧版领馆：省A·XXXX领，第二位是城市字母 → 红字
|
||
elif i == 1 and plate_number[0] in letters and np.random.random(1) > 0.5:
|
||
is_red = True # 军车第二字符随机红色
|
||
|
||
if enhance:
|
||
k = np.random.randint(1, 6)
|
||
kernel = np.ones((k, k), np.uint8)
|
||
if np.random.random(1) > 0.5:
|
||
font_img = np.copy(cv2.erode(font_img, kernel, iterations=1))
|
||
else:
|
||
font_img = np.copy(cv2.dilate(font_img, kernel, iterations=1))
|
||
|
||
# 贴上底板
|
||
img_plate_model = copy_to_image_multi(img_plate_model, font_img,
|
||
number_xy[i, :], bg_color, is_red)
|
||
|
||
img_plate_model = cv2.blur(img_plate_model, (3, 3))
|
||
|
||
return img_plate_model, number_xy, plate_number, bg_color, False
|
||
|
||
def generate_plate_special(self, plate_number, bg_color, is_double=False, enhance=False):
|
||
"""
|
||
生成特定号码、颜色车牌
|
||
:param plate_number: 车牌号码
|
||
:param bg_color: 背景颜色，blue/yellow/white/white_army/black/black_shi/green_car/green_truck
|
||
:param is_double: 是否双层（默认False，已移除双层挂车生成）
|
||
:param enhance: 图像增强
|
||
:return: 车牌图
|
||
"""
|
||
height = 220 if is_double else 140
|
||
|
||
# print(plate_number, height, bg_color, is_double)
|
||
number_xy = self.get_location_multi(plate_number, height)
|
||
img_plate_model = cv2.imread(os.path.join(self.adr_plate_model, '{}_{}.PNG'.format(bg_color, height)))
|
||
img_plate_model = cv2.resize(img_plate_model, (440 if len(plate_number) == 7 else 480, height))
|
||
|
||
for i in range(len(plate_number)):
|
||
if len(plate_number) == 8:
|
||
font_img = self.font_imgs['green_{}'.format(plate_number[i])]
|
||
else:
|
||
if '{}_{}'.format(height, plate_number[i]) in self.font_imgs:
|
||
font_img = self.font_imgs['{}_{}'.format(height, plate_number[i])]
|
||
else:
|
||
if i < 2:
|
||
font_img = self.font_imgs['220_up_{}'.format(plate_number[i])]
|
||
else:
|
||
font_img = self.font_imgs['220_down_{}'.format(plate_number[i])]
|
||
|
||
# 字符是否红色（同 generate_plate 逻辑）
|
||
is_red = False
|
||
if (i == 0 and plate_number[0] in letters) or plate_number[i] == '警':
|
||
is_red = True
|
||
elif plate_number[i] == '使' and i == 0:
|
||
is_red = True # 旧版使馆：使在首位 → 红字
|
||
elif plate_number[i] == '领' and len(plate_number) >= 2 and plate_number[1] in letters:
|
||
is_red = True # 旧版领馆：省A·XXXX领 → 红字
|
||
elif i == 1 and plate_number[0] in letters and np.random.random(1) > 0.5:
|
||
is_red = True
|
||
|
||
if enhance:
|
||
k = np.random.randint(1, 6)
|
||
kernel = np.ones((k, k), np.uint8)
|
||
if np.random.random(1) > 0.5:
|
||
font_img = np.copy(cv2.erode(font_img, kernel, iterations=1))
|
||
else:
|
||
font_img = np.copy(cv2.dilate(font_img, kernel, iterations=1))
|
||
|
||
img_plate_model = copy_to_image_multi(img_plate_model, font_img,
|
||
number_xy[i, :], bg_color, is_red)
|
||
|
||
# is_double = 'double' if is_double else 'single'
|
||
img_plate_model = cv2.blur(img_plate_model, (3, 3))
|
||
|
||
return img_plate_model
|
||
|
||
|
||
def parse_args():
|
||
parser = argparse.ArgumentParser(description='中国车牌生成器')
|
||
parser.add_argument('--number', default=10, type=int, help='生成车牌数量')
|
||
parser.add_argument('--save-adr', default='multi_val', help='车牌保存路径')
|
||
args = parser.parse_args()
|
||
return args
|
||
|
||
|
||
def mkdir(path):
|
||
try:
|
||
os.makedirs(path)
|
||
except:
|
||
pass
|
||
|
||
|
||
if __name__ == '__main__':
|
||
args = parse_args()
|
||
print(args)
|
||
# 随机生成车牌
|
||
print('save in {}'.format(args.save_adr))
|
||
|
||
mkdir(args.save_adr)
|
||
generator = MultiPlateGenerator('plate_model', 'font_model')
|
||
|
||
for i in tqdm(range(args.number)):
|
||
img, number_xy, gt_plate_number, bg_color, _ = generator.generate_plate()
|
||
cv2.imwrite(os.path.join(args.save_adr, '{}_{}.jpg'.format(gt_plate_number, bg_color)), img) |