plate color support

.gitignore

@@ -20,6 +20,7 @@ test_imgs/
 weixian/
 moto/
 pic/
+mytest_img/
 # 不忽略下面指定的文件类型
 !*.cpp
 !*.h

detect_rec_plate.py

@@ -76,11 +76,14 @@ def get_plate_rec_landmark(img, xyxy, conf, landmarks, class_num,device,plate_re
     # h_w_r = roi_img_w/roi_img_h
     if class_label :        #判断是否是双层车牌，是双牌的话进行分割后然后拼接
         roi_img=get_split_merge(roi_img)
-    plate_number = get_plate_result(roi_img,device,plate_rec_model)                 #对车牌小图进行识别
+    plate_number,rec_prob,plate_color,color_conf = get_plate_result(roi_img,device,plate_rec_model)                 #对车牌小图进行识别
     
     result_dict['rect']=rect
     result_dict['landmarks']=landmarks_np.tolist()
     result_dict['plate_no']=plate_number
+    result_dict['rec_conf']=rec_prob   #每个字符的概率
+    result_dict['plate_color']=plate_color 
+    result_dict['color_conf']=color_conf 
     result_dict['roi_height']=roi_img.shape[0]
     result_dict['score']=conf
     result_dict['label']=class_label
@@ -122,39 +125,46 @@ def draw_result(orgimg,dict_list):
     for result in dict_list:
         rect_area = result['rect']
         
-        # x,y,w,h = rect_area[0],rect_area[1],rect_area[2]-rect_area[0],rect_area[3]-rect_area[1]
-        # padding_w = 0
-        # padding_h = 0
-        # rect_area[0]=max(0,int(x-padding_w))
-        # rect_area[1]=max(0,int(y-padding_h))
-        # rect_area[2]=min(orgimg.shape[0],int(rect_area[2]+padding_w))
-        # rect_area[3]=min(orgimg.shape[1],int(rect_area[3]+padding_h))
+        x,y,w,h = rect_area[0],rect_area[1],rect_area[2]-rect_area[0],rect_area[3]-rect_area[1]
+        padding_w = 0.05*w
+        padding_h = 0.11*h
+        rect_area[0]=max(0,int(x-padding_w))
+        rect_area[1]=max(0,int(y-padding_h))
+        rect_area[2]=min(orgimg.shape[1],int(rect_area[2]+padding_w))
+        rect_area[3]=min(orgimg.shape[0],int(rect_area[3]+padding_h))
+        rect_area = [int(x) for x in rect_area]
 
         height_area = result['roi_height']
         landmarks=result['landmarks']
-        result = result['plate_no']
-        result_str+=result+" "
+        result_p = result['plate_no']+" "+result['plate_color']
+        result_str+=result_p+" "
         cv2.rectangle(orgimg,(rect_area[0],rect_area[1]),(rect_area[2],rect_area[3]),(0,0,255),2) #画框
+        labelSize = cv2.getTextSize(result_p,cv2.FONT_HERSHEY_SIMPLEX,0.5,1)
+        if rect_area[0]+labelSize[0][0]>orgimg.shape[1]:                 #防止显示的文字越界
+            rect_area[0]=int(orgimg.shape[1]-labelSize[0][0])
+            
+        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)
+        
         if len(result)>1:
             for i in range(4):  #关键点
                 cv2.circle(orgimg, (int(landmarks[i][0]), int(landmarks[i][1])), 5, clors[i], -1)
-            
-            orgimg=cv2ImgAddText(orgimg,result,rect_area[0]-height_area,rect_area[1]-height_area-10,(0,255,0),height_area)
+            orgimg=cv2ImgAddText(orgimg,result_p,rect_area[0],int(rect_area[1]-round(1.6*labelSize[0][1])),(0,0,0),21)
+            # orgimg=cv2ImgAddText(orgimg,result,rect_area[0]-height_area,rect_area[1]-height_area-10,(0,255,0),height_area)
     print(result_str)
     return orgimg
 
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument('--detect_model', nargs='+', type=str, default='weights/yolov7-lite-s.pt', help='model.pt path(s)')
-    parser.add_argument('--rec_model', type=str, default='weights/plate_rec_small.pth', help='model.pt path(s)')
+    parser.add_argument('--detect_model', nargs='+', type=str, default='weights/yolov7-lite-s.pt', help='model.pt path(s)') #检测模型
+    parser.add_argument('--rec_model', type=str, default='weights/plate_rec_color.pth', help='model.pt path(s)')  #车牌识别 +颜色识别
     parser.add_argument('--source', type=str, default='imgs', help='source')  # file/folder, 0 for webcam
     # 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='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")
+    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)

models/export.py

@@ -26,7 +26,7 @@ from utils.torch_utils import select_device
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument('--weights', type=str, default='./yolov5s.pt', help='weights path')
-    parser.add_argument('--img-size', nargs='+', type=int, default=[320, 320], help='image size')  # height, width
+    parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size')  # height, width
     parser.add_argument('--batch-size', type=int, default=1, help='batch size')
     parser.add_argument('--grid', action='store_true', help='export Detect() layer grid')
     parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')

onnx/yolov7_plate_onnx_infer.py

@@ -232,10 +232,10 @@ def draw_result(orgimg,dict_list):
 if __name__ == "__main__":
    
     parser = argparse.ArgumentParser()
-    parser.add_argument('--detect_model',type=str, default=r'runs/train/yolov714/weights/best.onnx', help='model.pt path(s)')  #检测模型
+    parser.add_argument('--detect_model',type=str, default=r'weights/yolov7-lite-s.onnx', help='model.pt path(s)')  #检测模型
     parser.add_argument('--rec_model', type=str, default='weights/plate_rec.onnx', help='model.pt path(s)')#识别模型
     parser.add_argument('--image_path', type=str, default=r'pic', help='source') 
-    parser.add_argument('--img_size', type=int, default=320, help='inference size (pixels)')
+    parser.add_argument('--img_size', type=int, default=640, help='inference size (pixels)')
     parser.add_argument('--output', type=str, default='result', help='source') 
     opt = parser.parse_args()
     file_list = []

plate_recognition/plateNet.py

@@ -47,7 +47,7 @@ class myNet_ocr(nn.Module):
         if self.export:
             conv = x.squeeze(2) # b *512 * width
             conv = conv.transpose(2,1)  # [w, b, c]
-            conv =conv.argmax(dim=2)
+            # conv =conv.argmax(dim=2)
             return conv
         else:
             b, c, h, w = x.size()
@@ -95,9 +95,109 @@ class myNet(nn.Module):
         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,168)
+    x = torch.randn(1,3,48,216)
     model = myNet_ocr(num_classes=78,export=True)
     out = model(x)
     print(out.shape)

plate_recognition/plate_rec.py

@@ -1,4 +1,4 @@
-from plate_recognition.plateNet import myNet_ocr
+from plate_recognition.plateNet import myNet_ocr,myNet_ocr_color
 import torch
 import torch.nn as nn
 import cv2
@@ -20,16 +20,19 @@ def allFilePath(rootPath,allFIleList):
         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
+    return newPreds,index
 
 def image_processing(img,device):
     img = cv2.resize(img, (168,48))
@@ -45,54 +48,72 @@ def image_processing(img,device):
     img = img.view(1, *img.size())
     return img
 
-def get_plate_result(img,device,model):
+def get_plate_result(img,device,model,is_color=True):
     input = image_processing(img,device)
-    preds = model(input)
-    # print(preds)
-    preds=preds.view(-1).detach().cpu().numpy()
-    newPreds=decodePlate(preds)
+    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 ""
-    return plate
+    if is_color:
+        return plate,prob,color[color_index],color_conf    #返回车牌号以及每个字符的概率,以及颜色，和颜色的概率
+    else:
+        return plate,prob
 
-def init_model(device,model_path):
+def init_model(device,model_path,is_color = True):
     # 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']
-    # model_path = os.sep.join([sys.path[0],model_path])
-    model = myNet_ocr(num_classes=78,export=True,cfg=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)
+    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"double_plate"
+   testPath = r"/mnt/Gpan/Mydata/pytorchPorject/CRNN/crnn_plate_recognition/images"
    fileList=[]
    allFilePath(testPath,fileList)
 #    result = get_plate_result(image_path,device)
 #    print(result)
-   model = init_model(device)
+   is_color = False
+   model = init_model(device,model_path,is_color=is_color)
    right=0
    begin = time.time()
+   
    for imge_path in fileList:
-        plate=get_plate_result(imge_path)
-        plate_ori = imge_path.split('/')[-1].split('_')[0]
-        # print(plate,"---",plate_ori)
-        if(plate==plate_ori):
-
-            right+=1
+        img=cv2.imread(imge_path)
+        if is_color:
+            plate,_,plate_color,_=get_plate_result(img,device,model,is_color=is_color)
+            print(plate)
         else:
-            print(plate_ori,"--->",plate,imge_path)
-   end=time.time()
-   print("sum:%d ,right:%d , accuracy: %f, time: %f"%(len(fileList),right,right/len(fileList),end-begin))
+            plate,_=get_plate_result(img,device,model,is_color=is_color)
+            print(plate,imge_path)
+        
+