添加注释
This commit is contained in:
15
plate_recognition/double_plate_split_merge.py
Normal file
15
plate_recognition/double_plate_split_merge.py
Normal file
@ -0,0 +1,15 @@
|
||||
import os
|
||||
import cv2
|
||||
import numpy as np
|
||||
def get_split_merge(img):
|
||||
h,w,c = img.shape
|
||||
img_upper = img[0:int(5/12*h),:]
|
||||
img_lower = img[int(1/3*h):,:]
|
||||
img_upper = cv2.resize(img_upper,(img_lower.shape[1],img_lower.shape[0]))
|
||||
new_img = np.hstack((img_upper,img_lower))
|
||||
return new_img
|
||||
|
||||
if __name__=="__main__":
|
||||
img = cv2.imread("double_plate/tmp8078.png")
|
||||
new_img =get_split_merge(img)
|
||||
cv2.imwrite("double_plate/new.jpg",new_img)
|
||||
203
plate_recognition/plateNet.py
Normal file
203
plate_recognition/plateNet.py
Normal file
@ -0,0 +1,203 @@
|
||||
import torch.nn as nn
|
||||
import torch
|
||||
|
||||
|
||||
class myNet_ocr(nn.Module):
|
||||
def __init__(self,cfg=None,num_classes=78,export=False):
|
||||
super(myNet_ocr, self).__init__()
|
||||
if cfg is None:
|
||||
cfg =[32,32,64,64,'M',128,128,'M',196,196,'M',256,256]
|
||||
# cfg =[32,32,'M',64,64,'M',128,128,'M',256,256]
|
||||
self.feature = self.make_layers(cfg, True)
|
||||
self.export = export
|
||||
# self.classifier = nn.Linear(cfg[-1], num_classes)
|
||||
# self.loc = nn.MaxPool2d((2, 2), (5, 1), (0, 1),ceil_mode=True)
|
||||
# self.loc = nn.AvgPool2d((2, 2), (5, 2), (0, 1),ceil_mode=False)
|
||||
self.loc = nn.MaxPool2d((5, 2), (1, 1),(0,1),ceil_mode=False)
|
||||
self.newCnn=nn.Conv2d(cfg[-1],num_classes,1,1)
|
||||
# self.newBn=nn.BatchNorm2d(num_classes)
|
||||
def make_layers(self, cfg, batch_norm=False):
|
||||
layers = []
|
||||
in_channels = 3
|
||||
for i in range(len(cfg)):
|
||||
if i == 0:
|
||||
conv2d =nn.Conv2d(in_channels, cfg[i], kernel_size=5,stride =1)
|
||||
if batch_norm:
|
||||
layers += [conv2d, nn.BatchNorm2d(cfg[i]), nn.ReLU(inplace=True)]
|
||||
else:
|
||||
layers += [conv2d, nn.ReLU(inplace=True)]
|
||||
in_channels = cfg[i]
|
||||
else :
|
||||
if cfg[i] == 'M':
|
||||
layers += [nn.MaxPool2d(kernel_size=3, stride=2,ceil_mode=True)]
|
||||
else:
|
||||
conv2d = nn.Conv2d(in_channels, cfg[i], kernel_size=3, padding=(1,1),stride =1)
|
||||
if batch_norm:
|
||||
layers += [conv2d, nn.BatchNorm2d(cfg[i]), nn.ReLU(inplace=True)]
|
||||
else:
|
||||
layers += [conv2d, nn.ReLU(inplace=True)]
|
||||
in_channels = cfg[i]
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.feature(x)
|
||||
x=self.loc(x)
|
||||
x=self.newCnn(x)
|
||||
# x=self.newBn(x)
|
||||
if self.export:
|
||||
conv = x.squeeze(2) # b *512 * width
|
||||
conv = conv.transpose(2,1) # [w, b, c]
|
||||
# conv =conv.argmax(dim=2)
|
||||
return conv
|
||||
else:
|
||||
b, c, h, w = x.size()
|
||||
assert h == 1, "the height of conv must be 1"
|
||||
conv = x.squeeze(2) # b *512 * width
|
||||
conv = conv.permute(2, 0, 1) # [w, b, c]
|
||||
# output = F.log_softmax(self.rnn(conv), dim=2)
|
||||
output = torch.softmax(conv, dim=2)
|
||||
return output
|
||||
|
||||
myCfg = [32,'M',64,'M',96,'M',128,'M',256]
|
||||
class myNet(nn.Module):
|
||||
def __init__(self,cfg=None,num_classes=3):
|
||||
super(myNet, self).__init__()
|
||||
if cfg is None:
|
||||
cfg = myCfg
|
||||
self.feature = self.make_layers(cfg, True)
|
||||
self.classifier = nn.Linear(cfg[-1], num_classes)
|
||||
def make_layers(self, cfg, batch_norm=False):
|
||||
layers = []
|
||||
in_channels = 3
|
||||
for i in range(len(cfg)):
|
||||
if i == 0:
|
||||
conv2d =nn.Conv2d(in_channels, cfg[i], kernel_size=5,stride =1)
|
||||
if batch_norm:
|
||||
layers += [conv2d, nn.BatchNorm2d(cfg[i]), nn.ReLU(inplace=True)]
|
||||
else:
|
||||
layers += [conv2d, nn.ReLU(inplace=True)]
|
||||
in_channels = cfg[i]
|
||||
else :
|
||||
if cfg[i] == 'M':
|
||||
layers += [nn.MaxPool2d(kernel_size=3, stride=2,ceil_mode=True)]
|
||||
else:
|
||||
conv2d = nn.Conv2d(in_channels, cfg[i], kernel_size=3, padding=1,stride =1)
|
||||
if batch_norm:
|
||||
layers += [conv2d, nn.BatchNorm2d(cfg[i]), nn.ReLU(inplace=True)]
|
||||
else:
|
||||
layers += [conv2d, nn.ReLU(inplace=True)]
|
||||
in_channels = cfg[i]
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.feature(x)
|
||||
x = nn.AvgPool2d(kernel_size=3, stride=1)(x)
|
||||
x = x.view(x.size(0), -1)
|
||||
y = self.classifier(x)
|
||||
return y
|
||||
|
||||
|
||||
class MyNet_color(nn.Module):
|
||||
def __init__(self, class_num=6):
|
||||
super(MyNet_color, self).__init__()
|
||||
self.class_num = class_num
|
||||
self.backbone = nn.Sequential(
|
||||
nn.Conv2d(in_channels=3, out_channels=16, kernel_size=(5, 5), stride=(1, 1)), # 0
|
||||
torch.nn.BatchNorm2d(16),
|
||||
nn.ReLU(),
|
||||
nn.MaxPool2d(kernel_size=(2, 2)),
|
||||
nn.Dropout(0),
|
||||
nn.Flatten(),
|
||||
nn.Linear(480, 64),
|
||||
nn.Dropout(0),
|
||||
nn.ReLU(),
|
||||
nn.Linear(64, class_num),
|
||||
nn.Dropout(0),
|
||||
nn.Softmax(1)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
logits = self.backbone(x)
|
||||
|
||||
return logits
|
||||
|
||||
|
||||
class myNet_ocr_color(nn.Module):
|
||||
def __init__(self,cfg=None,num_classes=78,export=False,color_num=None):
|
||||
super(myNet_ocr_color, self).__init__()
|
||||
if cfg is None:
|
||||
cfg =[32,32,64,64,'M',128,128,'M',196,196,'M',256,256]
|
||||
# cfg =[32,32,'M',64,64,'M',128,128,'M',256,256]
|
||||
self.feature = self.make_layers(cfg, True)
|
||||
self.export = export
|
||||
self.color_num=color_num
|
||||
self.conv_out_num=12 #颜色第一个卷积层输出通道12
|
||||
if self.color_num:
|
||||
self.conv1=nn.Conv2d(cfg[-1],self.conv_out_num,kernel_size=3,stride=2)
|
||||
self.bn1=nn.BatchNorm2d(self.conv_out_num)
|
||||
self.relu1=nn.ReLU(inplace=True)
|
||||
self.gap =nn.AdaptiveAvgPool2d(output_size=1)
|
||||
self.color_classifier=nn.Conv2d(self.conv_out_num,self.color_num,kernel_size=1,stride=1)
|
||||
self.color_bn = nn.BatchNorm2d(self.color_num)
|
||||
self.flatten = nn.Flatten()
|
||||
self.loc = nn.MaxPool2d((5, 2), (1, 1),(0,1),ceil_mode=False)
|
||||
self.newCnn=nn.Conv2d(cfg[-1],num_classes,1,1)
|
||||
# self.newBn=nn.BatchNorm2d(num_classes)
|
||||
def make_layers(self, cfg, batch_norm=False):
|
||||
layers = []
|
||||
in_channels = 3
|
||||
for i in range(len(cfg)):
|
||||
if i == 0:
|
||||
conv2d =nn.Conv2d(in_channels, cfg[i], kernel_size=5,stride =1)
|
||||
if batch_norm:
|
||||
layers += [conv2d, nn.BatchNorm2d(cfg[i]), nn.ReLU(inplace=True)]
|
||||
else:
|
||||
layers += [conv2d, nn.ReLU(inplace=True)]
|
||||
in_channels = cfg[i]
|
||||
else :
|
||||
if cfg[i] == 'M':
|
||||
layers += [nn.MaxPool2d(kernel_size=3, stride=2,ceil_mode=True)]
|
||||
else:
|
||||
conv2d = nn.Conv2d(in_channels, cfg[i], kernel_size=3, padding=(1,1),stride =1)
|
||||
if batch_norm:
|
||||
layers += [conv2d, nn.BatchNorm2d(cfg[i]), nn.ReLU(inplace=True)]
|
||||
else:
|
||||
layers += [conv2d, nn.ReLU(inplace=True)]
|
||||
in_channels = cfg[i]
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.feature(x)
|
||||
if self.color_num:
|
||||
x_color=self.conv1(x)
|
||||
x_color=self.bn1(x_color)
|
||||
x_color =self.relu1(x_color)
|
||||
x_color = self.color_classifier(x_color)
|
||||
x_color = self.color_bn(x_color)
|
||||
x_color =self.gap(x_color)
|
||||
x_color = self.flatten(x_color)
|
||||
x=self.loc(x)
|
||||
x=self.newCnn(x)
|
||||
|
||||
if self.export:
|
||||
conv = x.squeeze(2) # b *512 * width
|
||||
conv = conv.transpose(2,1) # [w, b, c]
|
||||
if self.color_num:
|
||||
return conv,x_color
|
||||
return conv
|
||||
else:
|
||||
b, c, h, w = x.size()
|
||||
assert h == 1, "the height of conv must be 1"
|
||||
conv = x.squeeze(2) # b *512 * width
|
||||
conv = conv.permute(2, 0, 1) # [w, b, c]
|
||||
output = F.log_softmax(conv, dim=2)
|
||||
if self.color_num:
|
||||
return output,x_color
|
||||
return output
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
x = torch.randn(1,3,48,216)
|
||||
model = myNet_ocr(num_classes=78,export=True)
|
||||
out = model(x)
|
||||
print(out.shape)
|
||||
119
plate_recognition/plate_rec.py
Normal file
119
plate_recognition/plate_rec.py
Normal file
@ -0,0 +1,119 @@
|
||||
from plate_recognition.plateNet import myNet_ocr,myNet_ocr_color
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import cv2
|
||||
import numpy as np
|
||||
import os
|
||||
import time
|
||||
import sys
|
||||
|
||||
def cv_imread(path): #可以读取中文路径的图片
|
||||
img=cv2.imdecode(np.fromfile(path,dtype=np.uint8),-1)
|
||||
return img
|
||||
|
||||
def allFilePath(rootPath,allFIleList):
|
||||
fileList = os.listdir(rootPath)
|
||||
for temp in fileList:
|
||||
if os.path.isfile(os.path.join(rootPath,temp)):
|
||||
if temp.endswith('.jpg') or temp.endswith('.png') or temp.endswith('.JPG'):
|
||||
allFIleList.append(os.path.join(rootPath,temp))
|
||||
else:
|
||||
allFilePath(os.path.join(rootPath,temp),allFIleList)
|
||||
device = torch.device('cuda') if torch.cuda.is_available() else torch.device("cpu")
|
||||
color=['黑色','蓝色','绿色','白色','黄色']
|
||||
plateName=r"#京沪津渝冀晋蒙辽吉黑苏浙皖闽赣鲁豫鄂湘粤桂琼川贵云藏陕甘青宁新学警港澳挂使领民航危0123456789ABCDEFGHJKLMNPQRSTUVWXYZ险品"
|
||||
mean_value,std_value=(0.588,0.193)
|
||||
def decodePlate(preds):
|
||||
pre=0
|
||||
newPreds=[]
|
||||
index=[]
|
||||
for i in range(len(preds)):
|
||||
if preds[i]!=0 and preds[i]!=pre:
|
||||
newPreds.append(preds[i])
|
||||
index.append(i)
|
||||
pre=preds[i]
|
||||
return newPreds,index
|
||||
|
||||
def image_processing(img,device):
|
||||
img = cv2.resize(img, (168,48))
|
||||
img = np.reshape(img, (48, 168, 3))
|
||||
|
||||
# normalize
|
||||
img = img.astype(np.float32)
|
||||
img = (img / 255. - mean_value) / std_value
|
||||
img = img.transpose([2, 0, 1])
|
||||
img = torch.from_numpy(img)
|
||||
|
||||
img = img.to(device)
|
||||
img = img.view(1, *img.size())
|
||||
return img
|
||||
|
||||
def get_plate_result(img,device,model,is_color=False):
|
||||
input = image_processing(img,device)
|
||||
if is_color: #是否识别颜色
|
||||
preds,color_preds = model(input)
|
||||
color_preds = torch.softmax(color_preds,dim=-1)
|
||||
color_conf,color_index = torch.max(color_preds,dim=-1)
|
||||
color_conf=color_conf.item()
|
||||
else:
|
||||
preds = model(input)
|
||||
preds=torch.softmax(preds,dim=-1)
|
||||
prob,index=preds.max(dim=-1)
|
||||
index = index.view(-1).detach().cpu().numpy()
|
||||
prob=prob.view(-1).detach().cpu().numpy()
|
||||
|
||||
|
||||
# preds=preds.view(-1).detach().cpu().numpy()
|
||||
newPreds,new_index=decodePlate(index)
|
||||
prob=prob[new_index]
|
||||
plate=""
|
||||
for i in newPreds:
|
||||
plate+=plateName[i]
|
||||
# if not (plate[0] in plateName[1:44] ):
|
||||
# return ""
|
||||
if is_color:
|
||||
return plate,prob,color[color_index],color_conf #返回车牌号以及每个字符的概率,以及颜色，和颜色的概率
|
||||
else:
|
||||
return plate,prob
|
||||
|
||||
def init_model(device,model_path,is_color = False):
|
||||
# print( print(sys.path))
|
||||
# model_path ="plate_recognition/model/checkpoint_61_acc_0.9715.pth"
|
||||
check_point = torch.load(model_path,map_location=device)
|
||||
model_state=check_point['state_dict']
|
||||
cfg=check_point['cfg']
|
||||
color_classes=0
|
||||
if is_color:
|
||||
color_classes=5 #颜色类别数
|
||||
model = myNet_ocr_color(num_classes=len(plateName),export=True,cfg=cfg,color_num=color_classes)
|
||||
|
||||
model.load_state_dict(model_state,strict=False)
|
||||
model.to(device)
|
||||
model.eval()
|
||||
return model
|
||||
|
||||
# model = init_model(device)
|
||||
if __name__ == '__main__':
|
||||
model_path = r"weights/plate_rec_color.pth"
|
||||
image_path ="images/tmp2424.png"
|
||||
testPath = r"/mnt/Gpan/Mydata/pytorchPorject/CRNN/crnn_plate_recognition/images"
|
||||
fileList=[]
|
||||
allFilePath(testPath,fileList)
|
||||
# result = get_plate_result(image_path,device)
|
||||
# print(result)
|
||||
is_color = False
|
||||
model = init_model(device,model_path,is_color=is_color)
|
||||
right=0
|
||||
begin = time.time()
|
||||
|
||||
for imge_path in fileList:
|
||||
img=cv2.imread(imge_path)
|
||||
if is_color:
|
||||
plate,_,plate_color,_=get_plate_result(img,device,model,is_color=is_color)
|
||||
print(plate)
|
||||
else:
|
||||
plate,_=get_plate_result(img,device,model,is_color=is_color)
|
||||
print(plate,imge_path)
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user