diff --git a/Car_recognition.py b/Car_recognition.py index 0924cbd..b4acefc 100644 --- a/Car_recognition.py +++ b/Car_recognition.py @@ -16,6 +16,7 @@ from plate_recognition.plate_rec import get_plate_result,allFilePath,init_model, # from plate_recognition.plate_cls import cv_imread from plate_recognition.double_plate_split_merge import get_split_merge from plate_recognition.color_rec import plate_color_rec,init_color_model +from car_recognition.car_rec import init_car_rec_model,get_color_and_score clors = [(255,0,0),(0,255,0),(0,0,255),(255,255,0),(0,255,255)] danger=['危','险'] @@ -65,7 +66,7 @@ def scale_coords_landmarks(img1_shape, coords, img0_shape, ratio_pad=None): # coords[:, [0, 2, 4, 6]] -= pad[0] # x padding coords[:, [1, 3, 5, 7]] -= pad[1] # y padding - coords[:, :10] /= gain + coords[:, :8] /= gain #clip_coords(coords, img0_shape) coords[:, 0].clamp_(0, img0_shape[1]) # x1 coords[:, 1].clamp_(0, img0_shape[0]) # y1 @@ -79,7 +80,7 @@ def scale_coords_landmarks(img1_shape, coords, img0_shape, ratio_pad=None): # # coords[:, 9].clamp_(0, img0_shape[0]) # y5 return coords -def get_plate_rec_landmark(img, xyxy, conf, landmarks, class_num,device,plate_rec_model,plate_color_model=None): +def get_plate_rec_landmark(img, xyxy, conf, landmarks, class_num,device,plate_rec_model,car_rec_model): h,w,c = img.shape result_dict={} x1 = int(xyxy[0]) @@ -90,10 +91,15 @@ def get_plate_rec_landmark(img, xyxy, conf, landmarks, class_num,device,plate_re rect=[x1,y1,x2,y2] if int(class_num) ==2: + # + car_roi_img = img[y1:y2,x1:x2] + car_color,color_conf=get_color_and_score(car_rec_model,car_roi_img,device) result_dict['class_type']=class_type[int(class_num)] result_dict['rect']=rect #车辆roi - result_dict['score']=conf + result_dict['score']=conf #车牌区域检测得分 result_dict['object_no']=int(class_num) + result_dict['car_color']=car_color + result_dict['color_conf']=color_conf return result_dict for i in range(4): @@ -103,10 +109,9 @@ def get_plate_rec_landmark(img, xyxy, conf, landmarks, class_num,device,plate_re class_label= int(class_num) #车牌的的类型0代表单牌，1代表双层车牌 roi_img = four_point_transform(img,landmarks_np) #透视变换得到车牌小图 - color_code = plate_color_rec(roi_img,plate_color_model,device) #车牌颜色识别 if class_label: #判断是否是双层车牌，是双牌的话进行分割后然后拼接 roi_img=get_split_merge(roi_img) - plate_number = get_plate_result(roi_img,device,plate_rec_model) #对车牌小图进行识别 + plate_number ,plate_color= get_plate_result(roi_img,device,plate_rec_model) #对车牌小图进行识别,得到颜色和车牌号 for dan in danger: #只要出现‘危’或者‘险’就是危险品车牌 if dan in plate_number: plate_number='危险品' @@ -116,14 +121,14 @@ def get_plate_rec_landmark(img, xyxy, conf, landmarks, class_num,device,plate_re result_dict['landmarks']=landmarks_np.tolist() #车牌角点坐标 result_dict['plate_no']=plate_number #车牌号 result_dict['roi_height']=roi_img.shape[0] #车牌高度 - result_dict['plate_color']=color_code #车牌颜色 + result_dict['plate_color']=plate_color #车牌颜色 result_dict['object_no']=class_label #单双层 0单层 1双层 - result_dict['score']=conf + result_dict['score']=conf #车牌区域检测得分 return result_dict -def detect_Recognition_plate(model, orgimg, device,plate_rec_model,img_size,plate_color_model=None): +def detect_Recognition_plate(model, orgimg, device,plate_rec_model,img_size,car_rec_model=None): # Load model # img_size = opt_img_size conf_thres = 0.3 @@ -183,7 +188,7 @@ def detect_Recognition_plate(model, orgimg, device,plate_rec_model,img_size,plat conf = det[j, 4].cpu().numpy() landmarks = det[j, 5:13].view(-1).tolist() class_num = det[j, 13].cpu().numpy() - result_dict = get_plate_rec_landmark(orgimg, xyxy, conf, landmarks, class_num,device,plate_rec_model,plate_color_model) + result_dict = get_plate_rec_landmark(orgimg, xyxy, conf, landmarks, class_num,device,plate_rec_model,car_rec_model) dict_list.append(result_dict) return dict_list # cv2.imwrite('result.jpg', orgimg) @@ -202,7 +207,7 @@ def draw_result(orgimg,dict_list): 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)) - height_area = result['roi_height'] + height_area = int(result['roi_height']/2) landmarks=result['landmarks'] result_p = result['plate_no'] if result['object_no']==0:#单层 @@ -218,6 +223,13 @@ def draw_result(orgimg,dict_list): orgimg=cv2ImgAddText(orgimg,result_p,rect_area[0],rect_area[3],(0,255,0),height_area) else: orgimg=cv2ImgAddText(orgimg,result_p,rect_area[0]-height_area,rect_area[1]-height_area-10,(0,255,0),height_area) + else: + height_area=int((rect_area[3]-rect_area[1])/20) + car_color = result['car_color'] + car_color_str="车辆颜色:" + car_color_str+=car_color + orgimg=cv2ImgAddText(orgimg,car_color_str,rect_area[0],rect_area[1],(0,255,0),height_area) + cv2.rectangle(orgimg,(rect_area[0],rect_area[1]),(rect_area[2],rect_area[3]),object_color[object_no],2) #画框 print(result_str) return orgimg @@ -233,8 +245,8 @@ def get_second(capture): if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--detect_model', nargs='+', type=str, default='weights/detect.pt', help='model.pt path(s)') #检测模型 - parser.add_argument('--rec_model', type=str, default='weights/plate_rec.pth', help='model.pt path(s)')#识别模型 - parser.add_argument('--color_model',type=str,default='weights/color_classify.pth',help='plate color')#颜色识别模型 + parser.add_argument('--rec_model', type=str, default='weights/plate_rec_color.pth', help='model.pt path(s)')#车牌识别+车牌颜色识别模型 + parser.add_argument('--car_rec_model',type=str,default='weights/car_rec_color.pth',help='car_rec_model') #车辆识别模型 parser.add_argument('--image_path', type=str, default='imgs', help='source') parser.add_argument('--img_size', type=int, default=384, help='inference size (pixels)') parser.add_argument('--output', type=str, default='result1', help='source') @@ -250,12 +262,12 @@ if __name__ == '__main__': detect_model = load_model(opt.detect_model, device) #初始化检测模型 plate_rec_model=init_model(device,opt.rec_model) #初始化识别模型 + car_rec_model = init_car_rec_model(opt.car_rec_model,device) #初始化车辆识别模型 #算参数量 total = sum(p.numel() for p in detect_model.parameters()) total_1 = sum(p.numel() for p in plate_rec_model.parameters()) print("detect params: %.2fM,rec params: %.2fM" % (total/1e6,total_1/1e6)) - plate_color_model =init_color_model(opt.color_model,device) time_all = 0 time_begin=time.time() if not opt.video: #处理图片 @@ -273,7 +285,8 @@ if __name__ == '__main__': if img.shape[-1]==4: img=cv2.cvtColor(img,cv2.COLOR_BGRA2BGR) # detect_one(model,img_path,device) - dict_list=detect_Recognition_plate(detect_model, img, device,plate_rec_model,opt.img_size,plate_color_model) + dict_list=detect_Recognition_plate(detect_model, img, device,plate_rec_model,opt.img_size,car_rec_model) + # print(dict_list) ori_img=draw_result(img,dict_list) img_name = os.path.basename(img_path) save_img_path = os.path.join(save_path,img_name) diff --git a/car_recognition/car_rec.py b/car_recognition/car_rec.py new file mode 100644 index 0000000..b5b48eb --- /dev/null +++ b/car_recognition/car_rec.py @@ -0,0 +1,63 @@ +from car_recognition.myNet import myNet +import torch +import cv2 +import torch.nn.functional as F +import os +import numpy as np + +colors = ['黑色','蓝色','黄色','棕色','绿色','灰色','橙色','粉色','紫色','红色','白色'] +def init_car_rec_model(model_path,device): + check_point = torch.load(model_path) + cfg= check_point['cfg'] + model = myNet(num_classes=11,cfg=cfg) + model.load_state_dict(check_point['state_dict']) + model.to(device) + model.eval() + return model + +def imge_processing(img,device): + img = cv2.resize(img,(64,64)) + img = img.transpose([2,0,1]) + img = torch.from_numpy(img).float().to(device) + img = img-127.5 + img = img.unsqueeze(0) + return img + +def allFilePath(rootPath,allFIleList): + fileList = os.listdir(rootPath) + for temp in fileList: + if os.path.isfile(os.path.join(rootPath,temp)): + allFIleList.append(os.path.join(rootPath,temp)) + else: + allFilePath(os.path.join(rootPath,temp),allFIleList) + +def get_color_and_score(model,img,device): + img = imge_processing(img,device) + result = model(img) + out =F.softmax( result) + _, predicted = torch.max(out.data, 1) + out=out.data.cpu().numpy().tolist()[0] + predicted = predicted.item() + car_color= colors[predicted] + color_conf = out[predicted] + # print(pic_,colors[predicted[0]]) + return car_color,color_conf + + +if __name__ == '__main__': + # root_file =r"/mnt/Gpan/BaiduNetdiskDownload/VehicleColour/VehicleColour/class/7" + root_file =r"imgs" + file_list=[] + allFilePath(root_file,file_list) + device = torch.device("cuda" if torch.cuda.is_available else "cpu") + model_path = r"/mnt/Gpan/Mydata/pytorchPorject/Car_system/car_color/color_model/0.8682285244554049_epoth_117_model.pth" + model = init_car_rec_model(model_path,device) + for pic_ in file_list: + img = cv2.imread(pic_) + # img = imge_processing(img,device) + color,conf = get_color_and_score(model,img,device) + print(pic_,color,conf) + + + + \ No newline at end of file diff --git a/car_recognition/myNet.py b/car_recognition/myNet.py new file mode 100644 index 0000000..76947ea --- /dev/null +++ b/car_recognition/myNet.py @@ -0,0 +1,95 @@ +import math + +import torch +import torch.nn as nn +from torch.autograd import Variable +from torchvision import models + + +__all__ = ['myNet','myResNet18'] + +# defaultcfg = { +# 11 : [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512], +# 13 : [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512], +# 16 : [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512], +# 19 : [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512], +# } +# myCfg = [32,'M',64,'M',96,'M',128,'M',192,'M',256] +myCfg = [32,'M',64,'M',96,'M',128,'M',256] +# myCfg = [8,'M',16,'M',32,'M',64,'M',96] +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.gap =nn.AdaptiveAvgPool2d((1,1)) + self.classifier = nn.Linear(cfg[-1], num_classes) + # self.classifier = nn.Conv2d(cfg[-1],num_classes,kernel_size=1,stride=1) + # self.bn_c= nn.BatchNorm2d(num_classes) + # self.flatten = nn.Flatten() + 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): + y = self.feature(x) + y = nn.AvgPool2d(kernel_size=3, stride=1)(y) + y = y.view(x.size(0), -1) + y = self.classifier(y) + + # y = self.flatten(y) + return y + +class myResNet18(nn.Module): + def __init__(self,num_classes=1000): + super(myResNet18,self).__init__() + model_ft = models.resnet18(pretrained=True) + self.model =model_ft + self.model.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1,ceil_mode=True) + self.model.averagePool = nn.AvgPool2d((5,5),stride=1,ceil_mode=True) + self.cls=nn.Linear(512,num_classes) + + def forward(self,x): + x = self.model.conv1(x) + x = self.model.bn1(x) + x = self.model.relu(x) + x = self.model.maxpool(x) + + x = self.model.layer1(x) + x = self.model.layer2(x) + x = self.model.layer3(x) + x = self.model.layer4(x) + + x = self.model.averagePool(x) + x = x.view(x.size(0), -1) + x = self.cls(x) + + return x +if __name__ == '__main__': + net = myNet(num_classes=2) + # infeatures = net.cls.in_features + # net.cls=nn.Linear(infeatures,2) + x = torch.FloatTensor(16, 3, 64, 64) + y = net(x) + print(y.shape) + # print(net) \ No newline at end of file diff --git a/detect_demo.py b/detect_demo.py index 46dace6..e08834f 100644 --- a/detect_demo.py +++ b/detect_demo.py @@ -169,7 +169,7 @@ def draw_result(orgimg,dict_list): return orgimg if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument('--detect_model', nargs='+', type=str, default='weights/plate_detect.pt', help='model.pt path(s)') #检测模型 + parser.add_argument('--detect_model', nargs='+', type=str, default='weights/detect.pt', help='model.pt path(s)') #检测模型 parser.add_argument('--image_path', type=str, default='imgs', help='source') parser.add_argument('--img_size', type=int, default=640, help='inference size (pixels)') parser.add_argument('--output', type=str, default='result1', help='source') diff --git a/plate_recognition/plateNet.py b/plate_recognition/plateNet.py index aaad71a..b4e8979 100644 --- a/plate_recognition/plateNet.py +++ b/plate_recognition/plateNet.py @@ -1,5 +1,6 @@ import torch.nn as nn import torch +import torch.nn.functional as F class myNet_ocr(nn.Module): @@ -120,6 +121,87 @@ class MyNet_color(nn.Module): 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.relu = nn.ReLU(inplace=True) + # 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) + 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) diff --git a/plate_recognition/plate_rec.py b/plate_recognition/plate_rec.py index 8a6a036..db9f73f 100644 --- a/plate_recognition/plate_rec.py +++ b/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 @@ -21,6 +21,7 @@ def allFilePath(rootPath,allFIleList): allFilePath(os.path.join(rootPath,temp),allFIleList) device = torch.device('cuda') if torch.cuda.is_available() else torch.device("cpu") plateName=r"#京沪津渝冀晋蒙辽吉黑苏浙皖闽赣鲁豫鄂湘粤桂琼川贵云藏陕甘青宁新学警港澳挂使领民航危0123456789ABCDEFGHJKLMNPQRSTUVWXYZ险品" +color_list=['黑色','蓝色','绿色','白色','黄色'] mean_value,std_value=(0.588,0.193) def decodePlate(preds): pre=0 @@ -47,17 +48,19 @@ def image_processing(img,device): def get_plate_result(img,device,model): input = image_processing(img,device) - preds = model(input) - # preds =preds.argmax(dim=2) #找出概率最大的那个字符 + preds,color_preds = model(input) + preds =preds.argmax(dim=2) #找出概率最大的那个字符 + color_preds = color_preds.argmax(dim=-1) # print(preds) preds=preds.view(-1).detach().cpu().numpy() + color_preds=color_preds.item() newPreds=decodePlate(preds) plate="" for i in newPreds: plate+=plateName[i] # if not (plate[0] in plateName[1:44] ): # return "" - return plate + return plate,color_list[color_preds] def init_model(device,model_path): # print( print(sys.path)) @@ -66,7 +69,7 @@ def init_model(device,model_path): 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=len(plateName),export=True,cfg=cfg) + model = myNet_ocr_color(num_classes=len(plateName),export=True,cfg=cfg,color_num=len(color_list)) model.load_state_dict(model_state) model.to(device) diff --git a/weights/color_classify.pth b/weights/color_classify.pth deleted file mode 100644 index 649e7c7..0000000 Binary files a/weights/color_classify.pth and /dev/null differ diff --git a/weights/plate_rec_color.pth b/weights/plate_rec_color.pth new file mode 100644 index 0000000..2aa9938 Binary files /dev/null and b/weights/plate_rec_color.pth differ