├── data └── image │ ├── 0.jpg │ ├── 1.jpg │ ├── 2.jpg │ ├── 3.jpg │ ├── 4.jpg │ ├── 5.jpg │ └── 6.jpg ├── weight └── Description.txt ├── README.md ├── roi_pooling.py ├── load_data.py ├── train_wR2.py ├── main.py └── train_rpnet.py /data/image/0.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/872699467/CCPD_CNN/HEAD/data/image/0.jpg -------------------------------------------------------------------------------- /data/image/1.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/872699467/CCPD_CNN/HEAD/data/image/1.jpg -------------------------------------------------------------------------------- /data/image/2.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/872699467/CCPD_CNN/HEAD/data/image/2.jpg -------------------------------------------------------------------------------- /data/image/3.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/872699467/CCPD_CNN/HEAD/data/image/3.jpg -------------------------------------------------------------------------------- /data/image/4.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/872699467/CCPD_CNN/HEAD/data/image/4.jpg -------------------------------------------------------------------------------- /data/image/5.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/872699467/CCPD_CNN/HEAD/data/image/5.jpg -------------------------------------------------------------------------------- /data/image/6.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/872699467/CCPD_CNN/HEAD/data/image/6.jpg -------------------------------------------------------------------------------- /weight/Description.txt: -------------------------------------------------------------------------------- 1 | Place these files in this directory (fh02.out、fh02.pth、wR2.pth) -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # CCPD_CNN 2 | License Plate Recognition 3 | 链接：https://pan.baidu.com/s/1F_rTFN6cKCh-boqTBzV_vg 4 | 提取码：y2y0 5 | -------------------------------------------------------------------------------- /roi_pooling.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import torch.nn.functional as F 4 | 5 | def roi_pooling(input, rois, size=(7, 7), spatial_scale=1.0): 6 | assert (rois.dim() == 2) 7 | assert (rois.size(1) == 5) 8 | output = [] 9 | rois = rois.data.float() 10 | num_rois = rois.size(0) 11 | 12 | rois[:, 1:].mul_(spatial_scale) 13 | rois = rois.long() 14 | for i in range(num_rois): 15 | roi = rois[i] 16 | im_idx = roi[0] 17 | # im = input.narrow(0, im_idx, 1) 18 | im = input.narrow(0, im_idx, 1)[..., roi[2]:(roi[4] + 1), roi[1]:(roi[3] + 1)] 19 | output.append(F.adaptive_max_pool2d(im, size)) 20 | 21 | return torch.cat(output, 0) 22 | 23 | 24 | def roi_pooling_ims(input, rois, size=(7, 7), spatial_scale=1.0): 25 | # written for one roi one image 26 | # size: (w, h) 27 | assert (rois.dim() == 2) 28 | assert len(input) == len(rois) 29 | assert (rois.size(1) == 4) 30 | output = [] 31 | rois = rois.data.float() 32 | num_rois = rois.size(0) 33 | 34 | rois[:, 1:].mul_(spatial_scale) 35 | rois = rois.long() 36 | for i in range(num_rois): 37 | roi = rois[i] 38 | im = input.narrow(0, i, 1)[..., roi[1]:(roi[3] + 1), roi[0]:(roi[2] + 1)] 39 | output.append(F.adaptive_max_pool2d(im, size)) 40 | 41 | return torch.cat(output, 0) 42 | 43 | if __name__ == '__main__': 44 | ''' 45 | input = torch.rand(2, 1, 10, 10) 46 | rois = torch.LongTensor([[1, 2, 7, 8], [3, 3, 8, 8]]) 47 | 48 | out = roi_pooling_ims(input, rois, size=(8, 8)) 49 | out.backward(out.data.clone().uniform_()) 50 | ''' 51 | # input = ag.Variable(torch.rand(2, 1, 10, 10), requires_grad=True) 52 | # rois = ag.Variable(torch.LongTensor([[0, 1, 2, 7, 8], [0, 3, 3, 8, 8], [1, 3, 3, 8, 8]]), requires_grad=False) 53 | # rois = ag.Variable(torch.LongTensor([[0,3,3,8,8]]),requires_grad=False) 54 | 55 | # out = adaptive_max_pool(input, (3, 3)) 56 | # out.backward(out.data.clone().uniform_()) 57 | 58 | # out = roi_pooling(input, rois, size=(3, 3)) 59 | # out.backward(out.data.clone().uniform_()) 60 | -------------------------------------------------------------------------------- /load_data.py: -------------------------------------------------------------------------------- 1 | from torch.utils.data import * 2 | from imutils import paths 3 | import cv2 4 | import numpy as np 5 | 6 | # pytorch自定义训练集数据加载器 7 | class labelFpsDataLoader(Dataset): 8 | def __init__(self, img_dir, imgSize, is_transform=None): 9 | self.img_dir = img_dir 10 | self.img_paths = [] 11 | for i in range(len(img_dir)): 12 | self.img_paths += [el for el in paths.list_images(img_dir[i])] 13 | # self.img_paths = os.listdir(img_dir) 14 | # print self.img_paths 15 | self.img_size = imgSize 16 | self.is_transform = is_transform 17 | 18 | def __len__(self): 19 | return len(self.img_paths) 20 | 21 | def __getitem__(self, index): 22 | img_name = self.img_paths[index] 23 | img = cv2.imread(img_name) 24 | resizedImage = cv2.resize(img, self.img_size) 25 | resizedImage = np.transpose(resizedImage, (2, 0, 1)) # [H,W,C]--->[C,H,W] 26 | resizedImage = resizedImage.astype('float32') 27 | resizedImage /= 255.0 28 | # img_name例如025-95_113-154&383_386&473-386&473_177&454_154&383_363&402-0_0_22_27_27_33_16-37-15.jpg 29 | lbl = img_name.split('/')[-1].rsplit('.', 1)[0].split('-')[-3] # 对应车牌号(0_0_22_27_27_33_16) 30 | 31 | iname = img_name.rsplit('/', 1)[-1].rsplit('.', 1)[0].split('-') 32 | [leftUp, rightDown] = [[int(eel) for eel in el.split('&')] for el in iname[2].split('_')] # 对应左上角和右下角坐标 33 | ori_w, ori_h = [float(int(el)) for el in [img.shape[1], img.shape[0]]] 34 | new_labels = [(leftUp[0] + rightDown[0]) / (2 * ori_w), (leftUp[1] + rightDown[1]) / (2 * ori_h), 35 | (rightDown[0] - leftUp[0]) / ori_w, (rightDown[1] - leftUp[1]) / ori_h] # 真实值[cen_x,cen_y,w,h] 36 | 37 | return resizedImage, new_labels, lbl, img_name 38 | 39 | # pytorch自定义测试集数据加载器 40 | class labelTestDataLoader(Dataset): 41 | def __init__(self, img_dir, imgSize, is_transform=None): 42 | self.img_dir = img_dir 43 | self.img_paths = [] 44 | for i in range(len(img_dir)): 45 | self.img_paths += [el for el in paths.list_images(img_dir[i])] 46 | # self.img_paths = os.listdir(img_dir) 47 | # print self.img_paths 48 | self.img_size = imgSize 49 | self.is_transform = is_transform 50 | 51 | def __len__(self): 52 | return len(self.img_paths) 53 | 54 | def __getitem__(self, index): 55 | img_name = self.img_paths[index] 56 | img = cv2.imread(img_name) 57 | # img = img.astype('float32') 58 | resizedImage = cv2.resize(img, self.img_size) 59 | resizedImage = np.transpose(resizedImage, (2, 0, 1)) 60 | resizedImage = resizedImage.astype('float32') 61 | resizedImage /= 255.0 62 | lbl = img_name.split('/')[-1].split('.')[0].split('-')[-3] 63 | return resizedImage, lbl, img_name 64 | 65 | 66 | class ChaLocDataLoader(Dataset): 67 | def __init__(self, img_dir, imgSize, is_transform=None): 68 | self.img_dir = img_dir 69 | self.img_paths = [] 70 | for i in range(len(img_dir)): 71 | self.img_paths += [el for el in paths.list_images(img_dir[i])] 72 | # self.img_paths = os.listdir(img_dir) 73 | # print self.img_paths 74 | self.img_size = imgSize 75 | self.is_transform = is_transform 76 | 77 | def __len__(self): 78 | return len(self.img_paths) 79 | 80 | def __getitem__(self, index): 81 | img_name = self.img_paths[index] 82 | img = cv2.imread(img_name) 83 | resizedImage = cv2.resize(img, self.img_size) 84 | resizedImage = np.reshape(resizedImage, (resizedImage.shape[2], resizedImage.shape[0], resizedImage.shape[1])) 85 | 86 | iname = img_name.rsplit('/', 1)[-1].rsplit('.', 1)[0].split('-') 87 | [leftUp, rightDown] = [[int(eel) for eel in el.split('&')] for el in iname[2].split('_')] 88 | 89 | 90 | ori_w, ori_h = float(img.shape[1]), float(img.shape[0]) 91 | assert img.shape[0] == 1160 92 | new_labels = [(leftUp[0] + rightDown[0]) / (2 * ori_w), (leftUp[1] + rightDown[1]) / (2 * ori_h), 93 | (rightDown[0] - leftUp[0]) / ori_w, (rightDown[1] - leftUp[1]) / ori_h] 94 | 95 | resizedImage = resizedImage.astype('float32') 96 | resizedImage /= 255.0 97 | return resizedImage, new_labels 98 | 99 | 100 | class demoTestDataLoader(Dataset): 101 | def __init__(self, img_dir, imgSize, is_transform=None): 102 | self.img_dir = img_dir 103 | self.img_paths = [] 104 | for i in range(len(img_dir)): 105 | self.img_paths += [el for el in paths.list_images(img_dir[i])] 106 | self.img_size = imgSize 107 | self.is_transform = is_transform 108 | 109 | def __len__(self): 110 | return len(self.img_paths) 111 | 112 | def __getitem__(self, index): 113 | img_name = self.img_paths[index] 114 | img = cv2.imread(img_name) 115 | # img = img.astype('float32') 116 | resizedImage = cv2.resize(img, self.img_size) 117 | resizedImage = np.transpose(resizedImage, (2, 0, 1)) 118 | resizedImage = resizedImage.astype('float32') 119 | resizedImage /= 255.0 120 | return resizedImage, img_name 121 | -------------------------------------------------------------------------------- /train_wR2.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import torch 3 | import torch.nn as nn 4 | import torch.optim as optim 5 | from torch.autograd import Variable 6 | import numpy as np 7 | import os 8 | import argparse 9 | from time import time 10 | from load_data import * 11 | from torch.optim import lr_scheduler 12 | 13 | def get_n_params(model): 14 | pp=0 15 | for p in list(model.parameters()): 16 | nn=1 17 | for s in list(p.size()): 18 | nn = nn*s 19 | pp += nn 20 | return pp 21 | 22 | class wR2(nn.Module): 23 | def __init__(self, num_classes=1000): 24 | super(wR2, self).__init__() 25 | hidden1 = nn.Sequential( 26 | nn.Conv2d(in_channels=3, out_channels=48, kernel_size=5, padding=2, stride=2), 27 | nn.BatchNorm2d(num_features=48), 28 | nn.ReLU(), 29 | nn.MaxPool2d(kernel_size=2, stride=2, padding=1), 30 | nn.Dropout(0.2) 31 | ) 32 | hidden2 = nn.Sequential( 33 | nn.Conv2d(in_channels=48, out_channels=64, kernel_size=5, padding=2), 34 | nn.BatchNorm2d(num_features=64), 35 | nn.ReLU(), 36 | nn.MaxPool2d(kernel_size=2, stride=1, padding=1), 37 | nn.Dropout(0.2) 38 | ) 39 | hidden3 = nn.Sequential( 40 | nn.Conv2d(in_channels=64, out_channels=128, kernel_size=5, padding=2), 41 | nn.BatchNorm2d(num_features=128), 42 | nn.ReLU(), 43 | nn.MaxPool2d(kernel_size=2, stride=2, padding=1), 44 | nn.Dropout(0.2) 45 | ) 46 | hidden4 = nn.Sequential( 47 | nn.Conv2d(in_channels=128, out_channels=160, kernel_size=5, padding=2), 48 | nn.BatchNorm2d(num_features=160), 49 | nn.ReLU(), 50 | nn.MaxPool2d(kernel_size=2, stride=1, padding=1), 51 | nn.Dropout(0.2) 52 | ) 53 | hidden5 = nn.Sequential( 54 | nn.Conv2d(in_channels=160, out_channels=192, kernel_size=5, padding=2), 55 | nn.BatchNorm2d(num_features=192), 56 | nn.ReLU(), 57 | nn.MaxPool2d(kernel_size=2, stride=2, padding=1), 58 | nn.Dropout(0.2) 59 | ) 60 | hidden6 = nn.Sequential( 61 | nn.Conv2d(in_channels=192, out_channels=192, kernel_size=5, padding=2), 62 | nn.BatchNorm2d(num_features=192), 63 | nn.ReLU(), 64 | nn.MaxPool2d(kernel_size=2, stride=1, padding=1), 65 | nn.Dropout(0.2) 66 | ) 67 | hidden7 = nn.Sequential( 68 | nn.Conv2d(in_channels=192, out_channels=192, kernel_size=5, padding=2), 69 | nn.BatchNorm2d(num_features=192), 70 | nn.ReLU(), 71 | nn.MaxPool2d(kernel_size=2, stride=2, padding=1), 72 | nn.Dropout(0.2) 73 | ) 74 | hidden8 = nn.Sequential( 75 | nn.Conv2d(in_channels=192, out_channels=192, kernel_size=5, padding=2), 76 | nn.BatchNorm2d(num_features=192), 77 | nn.ReLU(), 78 | nn.MaxPool2d(kernel_size=2, stride=1, padding=1), 79 | nn.Dropout(0.2) 80 | ) 81 | hidden9 = nn.Sequential( 82 | nn.Conv2d(in_channels=192, out_channels=192, kernel_size=3, padding=1), 83 | nn.BatchNorm2d(num_features=192), 84 | nn.ReLU(), 85 | nn.MaxPool2d(kernel_size=2, stride=2, padding=1), 86 | nn.Dropout(0.2) 87 | ) 88 | hidden10 = nn.Sequential( 89 | nn.Conv2d(in_channels=192, out_channels=192, kernel_size=3, padding=1), 90 | nn.BatchNorm2d(num_features=192), 91 | nn.ReLU(), 92 | nn.MaxPool2d(kernel_size=2, stride=1, padding=1), 93 | nn.Dropout(0.2) 94 | ) 95 | self.features = nn.Sequential( 96 | hidden1, 97 | hidden2, 98 | hidden3, 99 | hidden4, 100 | hidden5, 101 | hidden6, 102 | hidden7, 103 | hidden8, 104 | hidden9, 105 | hidden10 106 | ) 107 | self.classifier = nn.Sequential( 108 | nn.Linear(23232, 100), 109 | # nn.ReLU(inplace=True), 110 | nn.Linear(100, 100), 111 | # nn.ReLU(inplace=True), 112 | nn.Linear(100, num_classes), 113 | ) 114 | 115 | def forward(self, x): 116 | x1 = self.features(x) 117 | x11 = x1.view(x1.size(0), -1) 118 | x = self.classifier(x11) 119 | return x 120 | 121 | def train_model(model, criterion, optimizer, num_epochs=25): 122 | # since = time.time() 123 | for epoch in range(epoch_start, num_epochs): 124 | lossAver = [] 125 | model.train(True) 126 | lrScheduler.step() 127 | start = time() 128 | 129 | for i, (XI, YI) in enumerate(trainloader): 130 | # print('%s/%s %s' % (i, times, time()-start)) 131 | YI = np.array([el.numpy() for el in YI]).T 132 | if use_gpu: 133 | x = Variable(XI.cuda(0)) 134 | y = Variable(torch.FloatTensor(YI).cuda(0), requires_grad=False) 135 | else: 136 | x = Variable(XI) 137 | y = Variable(torch.FloatTensor(YI), requires_grad=False) 138 | # Forward pass: Compute predicted y by passing x to the model 139 | y_pred = model(x) 140 | 141 | # Compute and print loss 142 | loss = 0.0 143 | if len(y_pred) == batchSize: 144 | loss += 0.8 * nn.L1Loss().cuda()(y_pred[:][:2], y[:][:2]) 145 | loss += 0.2 * nn.L1Loss().cuda()(y_pred[:][2:], y[:][2:]) 146 | lossAver.append(loss.data[0]) 147 | 148 | # Zero gradients, perform a backward pass, and update the weights. 149 | optimizer.zero_grad() 150 | loss.backward() 151 | optimizer.step() 152 | torch.save(model.state_dict(), storeName) 153 | if i % 50 == 1: 154 | with open(args['writeFile'], 'a') as outF: 155 | outF.write('train %s images, use %s seconds, loss %s\n' % (i*batchSize, time() - start, sum(lossAver[-50:]) / len(lossAver[-50:]))) 156 | print ('%s %s %s\n' % (epoch, sum(lossAver) / len(lossAver), time()-start)) 157 | with open(args['writeFile'], 'a') as outF: 158 | outF.write('Epoch: %s %s %s\n' % (epoch, sum(lossAver) / len(lossAver), time()-start)) 159 | torch.save(model.state_dict(), storeName + str(epoch)) 160 | return model 161 | 162 | 163 | if __name__=='__main__': 164 | ap = argparse.ArgumentParser() 165 | ap.add_argument("-i", "--images", required=True, 166 | help="path to the input file") 167 | ap.add_argument("-n", "--epochs", default=25, 168 | help="epochs for train") 169 | ap.add_argument("-b", "--batchsize", default=4, 170 | help="batch size for train") 171 | ap.add_argument("-r", "--resume", default='111', 172 | help="file for re-train") 173 | ap.add_argument("-w", "--writeFile", default='wR2.out', 174 | help="file for output") 175 | args = vars(ap.parse_args()) 176 | 177 | use_gpu = torch.cuda.is_available() 178 | print(use_gpu) 179 | 180 | numClasses = 4 181 | imgSize = (480, 480) 182 | batchSize = int(args["batchsize"]) if use_gpu else 8 183 | modelFolder = 'wR2/' 184 | storeName = modelFolder + 'wR2.pth' 185 | if not os.path.isdir(modelFolder): 186 | os.mkdir(modelFolder) 187 | 188 | epochs = int(args["epochs"]) 189 | with open(args['writeFile'], 'wb') as outF: 190 | pass 191 | epoch_start = 0 192 | resume_file = str(args["resume"]) 193 | if not resume_file == '111':#resume表示从上次训练完保存的模型继续训练 194 | # epoch_start = int(resume_file[resume_file.find('pth') + 3:]) + 1 195 | if not os.path.isfile(resume_file): 196 | print("fail to load existed model! Existing ...") 197 | exit(0) 198 | print("Load existed model! %s" % resume_file) 199 | model_conv = wR2(numClasses) 200 | model_conv = torch.nn.DataParallel(model_conv, device_ids=range(torch.cuda.device_count())) 201 | model_conv.load_state_dict(torch.load(resume_file)) 202 | model_conv = model_conv.cuda() 203 | else: #从头开始训练 204 | model_conv = wR2(numClasses) 205 | if use_gpu: 206 | model_conv = torch.nn.DataParallel(model_conv, device_ids=range(torch.cuda.device_count())) 207 | model_conv = model_conv.cuda() 208 | 209 | print(model_conv) 210 | print(get_n_params(model_conv)) 211 | #识别模块采用交叉熵损失函数 212 | criterion = nn.MSELoss() 213 | #优化函数 214 | optimizer_conv = optim.SGD(model_conv.parameters(), lr=0.001, momentum=0.9) 215 | #训练过程中调整学习率,epoch次数每增大5次，学习率乘以0.1 216 | lrScheduler = lr_scheduler.StepLR(optimizer_conv, step_size=5, gamma=0.1) 217 | # optimizer_conv = optim.Adam(model_conv.parameters(), lr=0.01) 218 | # dst = LocDataLoader([args["images"]], imgSize) 219 | dst = ChaLocDataLoader(args["images"].split(','), imgSize) 220 | trainloader = DataLoader(dst, batch_size=batchSize, shuffle=True, num_workers=4) 221 | 222 | model_conv = train_model(model_conv, criterion, optimizer_conv, num_epochs=epochs) 223 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | # encoding:utf-8 2 | import cv2 3 | import torch 4 | import torch.nn as nn 5 | import argparse 6 | from load_data import * 7 | from time import time 8 | from roi_pooling import roi_pooling_ims 9 | from PIL import Image, ImageDraw, ImageFont 10 | from sklearn.decomposition import LatentDirichletAllocation 11 | 12 | 13 | class wR2(nn.Module): 14 | def __init__(self, num_classes=1000): 15 | super(wR2, self).__init__() 16 | hidden1 = nn.Sequential( 17 | nn.Conv2d(in_channels=3, out_channels=48, kernel_size=5, padding=2, stride=2), 18 | nn.BatchNorm2d(num_features=48), 19 | nn.ReLU(), 20 | nn.MaxPool2d(kernel_size=2, stride=2, padding=1), 21 | nn.Dropout(0.2) 22 | ) 23 | hidden2 = nn.Sequential( 24 | nn.Conv2d(in_channels=48, out_channels=64, kernel_size=5, padding=2), 25 | nn.BatchNorm2d(num_features=64), 26 | nn.ReLU(), 27 | nn.MaxPool2d(kernel_size=2, stride=1, padding=1), 28 | nn.Dropout(0.2) 29 | ) 30 | hidden3 = nn.Sequential( 31 | nn.Conv2d(in_channels=64, out_channels=128, kernel_size=5, padding=2), 32 | nn.BatchNorm2d(num_features=128), 33 | nn.ReLU(), 34 | nn.MaxPool2d(kernel_size=2, stride=2, padding=1), 35 | nn.Dropout(0.2) 36 | ) 37 | hidden4 = nn.Sequential( 38 | nn.Conv2d(in_channels=128, out_channels=160, kernel_size=5, padding=2), 39 | nn.BatchNorm2d(num_features=160), 40 | nn.ReLU(), 41 | nn.MaxPool2d(kernel_size=2, stride=1, padding=1), 42 | nn.Dropout(0.2) 43 | ) 44 | hidden5 = nn.Sequential( 45 | nn.Conv2d(in_channels=160, out_channels=192, kernel_size=5, padding=2), 46 | nn.BatchNorm2d(num_features=192), 47 | nn.ReLU(), 48 | nn.MaxPool2d(kernel_size=2, stride=2, padding=1), 49 | nn.Dropout(0.2) 50 | ) 51 | hidden6 = nn.Sequential( 52 | nn.Conv2d(in_channels=192, out_channels=192, kernel_size=5, padding=2), 53 | nn.BatchNorm2d(num_features=192), 54 | nn.ReLU(), 55 | nn.MaxPool2d(kernel_size=2, stride=1, padding=1), 56 | nn.Dropout(0.2) 57 | ) 58 | hidden7 = nn.Sequential( 59 | nn.Conv2d(in_channels=192, out_channels=192, kernel_size=5, padding=2), 60 | nn.BatchNorm2d(num_features=192), 61 | nn.ReLU(), 62 | nn.MaxPool2d(kernel_size=2, stride=2, padding=1), 63 | nn.Dropout(0.2) 64 | ) 65 | hidden8 = nn.Sequential( 66 | nn.Conv2d(in_channels=192, out_channels=192, kernel_size=5, padding=2), 67 | nn.BatchNorm2d(num_features=192), 68 | nn.ReLU(), 69 | nn.MaxPool2d(kernel_size=2, stride=1, padding=1), 70 | nn.Dropout(0.2) 71 | ) 72 | hidden9 = nn.Sequential( 73 | nn.Conv2d(in_channels=192, out_channels=192, kernel_size=3, padding=1), 74 | nn.BatchNorm2d(num_features=192), 75 | nn.ReLU(), 76 | nn.MaxPool2d(kernel_size=2, stride=2, padding=1), 77 | nn.Dropout(0.2) 78 | ) 79 | hidden10 = nn.Sequential( 80 | nn.Conv2d(in_channels=192, out_channels=192, kernel_size=3, padding=1), 81 | nn.BatchNorm2d(num_features=192), 82 | nn.ReLU(), 83 | nn.MaxPool2d(kernel_size=2, stride=1, padding=1), 84 | nn.Dropout(0.2) 85 | ) 86 | self.features = nn.Sequential( 87 | hidden1, 88 | hidden2, 89 | hidden3, 90 | hidden4, 91 | hidden5, 92 | hidden6, 93 | hidden7, 94 | hidden8, 95 | hidden9, 96 | hidden10 97 | ) 98 | self.classifier = nn.Sequential( 99 | nn.Linear(23232, 100), 100 | # nn.ReLU(inplace=True), 101 | nn.Linear(100, 100), 102 | # nn.ReLU(inplace=True), 103 | nn.Linear(100, num_classes), 104 | ) 105 | 106 | def forward(self, x): 107 | x1 = self.features(x) # [1,3,480,480]----->[1,192,11,11] 108 | x11 = x1.view(x1.size(0), -1) # [1,192,11,11]----->[1,23232] 109 | x = self.classifier(x11) # [1,23232]----->[1,1000] 110 | return x 111 | 112 | 113 | class fh02(nn.Module): 114 | def __init__(self, num_points, num_classes, wrPath=None): 115 | super(fh02, self).__init__() 116 | self.load_wR2(wrPath) 117 | self.classifier1 = nn.Sequential( 118 | # nn.Dropout(), 119 | nn.Linear(53248, 128), 120 | # nn.ReLU(inplace=True), 121 | # nn.Dropout(), 122 | nn.Linear(128, provNum), 123 | ) 124 | self.classifier2 = nn.Sequential( 125 | # nn.Dropout(), 126 | nn.Linear(53248, 128), 127 | # nn.ReLU(inplace=True), 128 | # nn.Dropout(), 129 | nn.Linear(128, alphaNum), 130 | ) 131 | self.classifier3 = nn.Sequential( 132 | # nn.Dropout(), 133 | nn.Linear(53248, 128), 134 | # nn.ReLU(inplace=True), 135 | # nn.Dropout(), 136 | nn.Linear(128, adNum), 137 | ) 138 | self.classifier4 = nn.Sequential( 139 | # nn.Dropout(), 140 | nn.Linear(53248, 128), 141 | # nn.ReLU(inplace=True), 142 | # nn.Dropout(), 143 | nn.Linear(128, adNum), 144 | ) 145 | self.classifier5 = nn.Sequential( 146 | # nn.Dropout(), 147 | nn.Linear(53248, 128), 148 | # nn.ReLU(inplace=True), 149 | # nn.Dropout(), 150 | nn.Linear(128, adNum), 151 | ) 152 | self.classifier6 = nn.Sequential( 153 | # nn.Dropout(), 154 | nn.Linear(53248, 128), 155 | # nn.ReLU(inplace=True), 156 | # nn.Dropout(), 157 | nn.Linear(128, adNum), 158 | ) 159 | self.classifier7 = nn.Sequential( 160 | # nn.Dropout(), 161 | nn.Linear(53248, 128), 162 | # nn.ReLU(inplace=True), 163 | # nn.Dropout(), 164 | nn.Linear(128, adNum), 165 | ) 166 | 167 | def load_wR2(self, path): 168 | self.wR2 = wR2(numPoints) 169 | self.wR2 = torch.nn.DataParallel(self.wR2, device_ids=range(torch.cuda.device_count())) 170 | if not path is None: 171 | self.wR2.load_state_dict(torch.load(path)) 172 | # self.wR2 = self.wR2.cuda() 173 | # for param in self.wR2.parameters(): 174 | # param.requires_grad = False 175 | 176 | def forward(self, x): 177 | x0 = self.wR2.module.features[0](x) # [1,3,480,480]--->[1,48,121,121] 178 | _x1 = self.wR2.module.features[1](x0) # [1,48,121,121]--->[1,64,122,122] 179 | x2 = self.wR2.module.features[2](_x1) # [1,64,122,122]--->[1,128,62,62] 180 | _x3 = self.wR2.module.features[3](x2) # [1,128,62,62]--->[1,160,63,63] 181 | x4 = self.wR2.module.features[4](_x3) # [1,160,63,63]--->[1,192,32,32] 182 | _x5 = self.wR2.module.features[5](x4) # [1,192,32,32]--->[1,192,33,33] 183 | 184 | x6 = self.wR2.module.features[6](_x5) # [1,192,33,33]--->[1,192,17,17] 185 | x7 = self.wR2.module.features[7](x6) # [1,192,17,17]--->[1,192,18,18] 186 | x8 = self.wR2.module.features[8](x7) # [1,192,18,18]--->[1,192,10,10] 187 | x9 = self.wR2.module.features[9](x8) # [1,192,10,10]--->[1,192,11,11] 188 | x9 = x9.view(x9.size(0), -1) # [1,192,11,11]--->[1,23232] 189 | boxLoc = self.wR2.module.classifier(x9) # [1,23232]--->[1,4] #预测车牌位置 190 | 191 | # 多尺度卷积：在不同层数的feature map在通道维数上进行拼接，然后预测 192 | with torch.no_grad(): 193 | h1, w1 = _x1.data.size()[2], _x1.data.size()[3] 194 | p1 = torch.FloatTensor([[w1, 0, 0, 0], [0, h1, 0, 0], [0, 0, w1, 0], [0, 0, 0, h1]]).cuda() 195 | h2, w2 = _x3.data.size()[2], _x3.data.size()[3] 196 | p2 = torch.FloatTensor([[w2, 0, 0, 0], [0, h2, 0, 0], [0, 0, w2, 0], [0, 0, 0, h2]]).cuda() 197 | h3, w3 = _x5.data.size()[2], _x5.data.size()[3] 198 | p3 = torch.FloatTensor([[w3, 0, 0, 0], [0, h3, 0, 0], [0, 0, w3, 0], [0, 0, 0, h3]]).cuda() 199 | 200 | # x, y, w, h --> x1, y1, x2, y2 201 | assert boxLoc.data.size()[1] == 4 202 | postfix = torch.FloatTensor([[1, 0, 1, 0], [0, 1, 0, 1], [-0.5, 0, 0.5, 0], [0, -0.5, 0, 0.5]]).cuda() 203 | boxNew = boxLoc.mm(postfix).clamp(min=0, max=1) 204 | 205 | # input = torch.rand(2, 1, 10, 10) 206 | # rois = torch.LongTensor([[0, 1, 2, 7, 8], [0, 3, 3, 8, 8], [1, 3, 3, 8, 8]]) 207 | roi1 = roi_pooling_ims(_x1, boxNew.mm(p1), size=(16, 8)) # [1,64,16,8] 208 | roi2 = roi_pooling_ims(_x3, boxNew.mm(p2), size=(16, 8)) # [1,160,16,8] 209 | roi3 = roi_pooling_ims(_x5, boxNew.mm(p3), size=(16, 8)) # [1,192,16,8] 210 | rois = torch.cat((roi1, roi2, roi3), 1) # [1,416,16,8] 211 | 212 | _rois = rois.view(rois.size(0), -1) # [1,53248] 213 | 214 | y0 = self.classifier1(_rois) # [1,38] 预测车牌的第一个字符：省份 215 | y1 = self.classifier2(_rois) # [1,25] 预测车牌的第二个字符：市 216 | y2 = self.classifier3(_rois) # [1,35] 预测车牌的第三个字符 217 | y3 = self.classifier4(_rois) # [1,35] 预测车牌的第四个字符 218 | y4 = self.classifier5(_rois) # [1,35] 预测车牌的第五个字符 219 | y5 = self.classifier6(_rois) # [1,35] 预测车牌的第六个字符 220 | y6 = self.classifier7(_rois) # [1,35] 预测车牌的第七个字符 221 | return boxLoc, [y0, y1, y2, y3, y4, y5, y6] 222 | 223 | 224 | def isEqual(labelGT, labelP): 225 | print(labelGT) 226 | print(labelP) 227 | compare = [1 if int(labelGT[i]) == int(labelP[i]) else 0 for i in range(7)] 228 | # print(sum(compare)) 229 | return sum(compare) 230 | 231 | 232 | if __name__ == '__main__': 233 | ap = argparse.ArgumentParser() 234 | ap.add_argument("-i", "--input", default='data/image/', 235 | help="path to the input folder") 236 | ap.add_argument("-m", "--model", default='weight/fh02.pth', 237 | help="path to the model file") 238 | args = vars(ap.parse_args()) 239 | # use_gpu = torch.cuda.is_available() 240 | use_gpu = False 241 | print(use_gpu) 242 | 243 | numClasses = 4 244 | numPoints = 4 245 | imgSize = (480, 480) 246 | batchSize = 8 if use_gpu else 1 247 | resume_file = str(args["model"]) 248 | 249 | provNum, alphaNum, adNum = 38, 25, 35 250 | provinces = ["皖", "沪", "津", "渝", "冀", "晋", "蒙", "辽", "吉", "黑", "苏", "浙", "京", "闽", "赣", "鲁", "豫", "鄂", "湘", "粤", 251 | "桂", 252 | "琼", "川", "贵", "云", "藏", "陕", "甘", "青", "宁", "新", "警", "学", "O"] 253 | alphabets = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 254 | 'W', 255 | 'X', 'Y', 'Z', 'O'] 256 | ads = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'J', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 257 | 'Y', 'Z', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'O'] 258 | 259 | model_conv = fh02(numPoints, numClasses) 260 | model_conv = torch.nn.DataParallel(model_conv, device_ids=range(torch.cuda.device_count())) 261 | model_conv.load_state_dict(torch.load(resume_file)) 262 | model_conv = model_conv.cuda() 263 | model_conv.eval() 264 | 265 | dst = demoTestDataLoader(args["input"].split(','), imgSize) 266 | trainloader = DataLoader(dst, batch_size=1, shuffle=False, num_workers=0) 267 | 268 | start = time() 269 | for i, (XI, ims) in enumerate(trainloader): 270 | 271 | if use_gpu: 272 | x = XI.cuda(0) 273 | else: 274 | x = XI 275 | # Forward pass: Compute predicted y by passing x to the model 276 | fps_pred, y_pred = model_conv(x) 277 | 278 | outputY = [el.data.cpu().numpy().tolist() for el in y_pred] 279 | labelPred = [t[0].index(max(t[0])) for t in outputY] 280 | 281 | [cx, cy, w, h] = fps_pred.data.cpu().numpy()[0].tolist() 282 | 283 | cv2Img = cv2.imread(ims[0]) 284 | left_up = [(cx - w / 2) * cv2Img.shape[1], (cy - h / 2) * cv2Img.shape[0]] 285 | right_down = [(cx + w / 2) * cv2Img.shape[1], (cy + h / 2) * cv2Img.shape[0]] 286 | cv2.rectangle(cv2Img, (int(left_up[0]), int(left_up[1])), (int(right_down[0]), int(right_down[1])), (0, 0, 255), 287 | 2) 288 | # The first character is Chinese character, can not be printed normally, thus is omitted. 289 | lpn = provinces[labelPred[0]] + alphabets[labelPred[1]] + ads[labelPred[2]] + ads[labelPred[3]] + ads[ 290 | labelPred[4]] + ads[labelPred[5]] + ads[labelPred[6]] 291 | print('图片名称:{0}识别结果{1}'.format(ims[0], lpn)) 292 | # PIL图片上打印汉字 293 | pilImg = Image.fromarray(cv2.cvtColor(cv2Img, cv2.COLOR_BGR2RGB)) 294 | draw = ImageDraw.Draw(pilImg) 295 | font = ImageFont.truetype("simhei.ttf", 40, encoding="utf-8") # 参数1：字体文件路径，参数2：字体大小 296 | draw.text((int(left_up[0]), int(left_up[1]) - 40), lpn, (255, 0, 0), 297 | font=font) # 参数1：打印坐标，参数2：文本，参数3：字体颜色，参数4：字体 298 | # PIL图片转cv2 图片 299 | cv2charimg = cv2.cvtColor(np.array(pilImg), cv2.COLOR_RGB2BGR) 300 | # cv2.putText(cv2Img, lpn, (int(left_up[0]), int(left_up[1])-20), cv2.FONT_ITALIC, 2, (0, 0, 255)) 301 | dstFileName = 'result/' + ims[0][-5:-4] + '.jpg' 302 | cv2.imshow('charimg',cv2charimg) 303 | cv2.waitKey(0) 304 | # cv2.imwrite(dstFileName, cv2charimg) 305 | # print('图片保存地址', dstFileName) 306 | -------------------------------------------------------------------------------- /train_rpnet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.optim as optim 4 | import numpy as np 5 | import os 6 | from visdom import Visdom 7 | import argparse 8 | from time import time 9 | from load_data import * 10 | from roi_pooling import roi_pooling_ims 11 | from torch.optim import lr_scheduler 12 | 13 | # 计算模型参数数量 14 | def get_n_params(model): 15 | pp = 0 16 | for p in list(model.parameters()): 17 | nn = 1 18 | for s in list(p.size()): 19 | nn = nn * s 20 | pp += nn 21 | return pp 22 | 23 | class wR2(nn.Module): 24 | def __init__(self, num_classes=1000): 25 | super(wR2, self).__init__() 26 | hidden1 = nn.Sequential( 27 | nn.Conv2d(in_channels=3, out_channels=48, kernel_size=5, padding=2, stride=2), 28 | nn.BatchNorm2d(num_features=48), 29 | nn.ReLU(), 30 | nn.MaxPool2d(kernel_size=2, stride=2, padding=1), 31 | nn.Dropout(0.2) 32 | ) 33 | hidden2 = nn.Sequential( 34 | nn.Conv2d(in_channels=48, out_channels=64, kernel_size=5, padding=2), 35 | nn.BatchNorm2d(num_features=64), 36 | nn.ReLU(), 37 | nn.MaxPool2d(kernel_size=2, stride=1, padding=1), 38 | nn.Dropout(0.2) 39 | ) 40 | hidden3 = nn.Sequential( 41 | nn.Conv2d(in_channels=64, out_channels=128, kernel_size=5, padding=2), 42 | nn.BatchNorm2d(num_features=128), 43 | nn.ReLU(), 44 | nn.MaxPool2d(kernel_size=2, stride=2, padding=1), 45 | nn.Dropout(0.2) 46 | ) 47 | hidden4 = nn.Sequential( 48 | nn.Conv2d(in_channels=128, out_channels=160, kernel_size=5, padding=2), 49 | nn.BatchNorm2d(num_features=160), 50 | nn.ReLU(), 51 | nn.MaxPool2d(kernel_size=2, stride=1, padding=1), 52 | nn.Dropout(0.2) 53 | ) 54 | hidden5 = nn.Sequential( 55 | nn.Conv2d(in_channels=160, out_channels=192, kernel_size=5, padding=2), 56 | nn.BatchNorm2d(num_features=192), 57 | nn.ReLU(), 58 | nn.MaxPool2d(kernel_size=2, stride=2, padding=1), 59 | nn.Dropout(0.2) 60 | ) 61 | hidden6 = nn.Sequential( 62 | nn.Conv2d(in_channels=192, out_channels=192, kernel_size=5, padding=2), 63 | nn.BatchNorm2d(num_features=192), 64 | nn.ReLU(), 65 | nn.MaxPool2d(kernel_size=2, stride=1, padding=1), 66 | nn.Dropout(0.2) 67 | ) 68 | hidden7 = nn.Sequential( 69 | nn.Conv2d(in_channels=192, out_channels=192, kernel_size=5, padding=2), 70 | nn.BatchNorm2d(num_features=192), 71 | nn.ReLU(), 72 | nn.MaxPool2d(kernel_size=2, stride=2, padding=1), 73 | nn.Dropout(0.2) 74 | ) 75 | hidden8 = nn.Sequential( 76 | nn.Conv2d(in_channels=192, out_channels=192, kernel_size=5, padding=2), 77 | nn.BatchNorm2d(num_features=192), 78 | nn.ReLU(), 79 | nn.MaxPool2d(kernel_size=2, stride=1, padding=1), 80 | nn.Dropout(0.2) 81 | ) 82 | hidden9 = nn.Sequential( 83 | nn.Conv2d(in_channels=192, out_channels=192, kernel_size=3, padding=1), 84 | nn.BatchNorm2d(num_features=192), 85 | nn.ReLU(), 86 | nn.MaxPool2d(kernel_size=2, stride=2, padding=1), 87 | nn.Dropout(0.2) 88 | ) 89 | hidden10 = nn.Sequential( 90 | nn.Conv2d(in_channels=192, out_channels=192, kernel_size=3, padding=1), 91 | nn.BatchNorm2d(num_features=192), 92 | nn.ReLU(), 93 | nn.MaxPool2d(kernel_size=2, stride=1, padding=1), 94 | nn.Dropout(0.2) 95 | ) 96 | self.features = nn.Sequential( 97 | hidden1, 98 | hidden2, 99 | hidden3, 100 | hidden4, 101 | hidden5, 102 | hidden6, 103 | hidden7, 104 | hidden8, 105 | hidden9, 106 | hidden10 107 | ) 108 | self.classifier = nn.Sequential( 109 | nn.Linear(23232, 100), 110 | # nn.ReLU(inplace=True), 111 | nn.Linear(100, 100), 112 | # nn.ReLU(inplace=True), 113 | nn.Linear(100, num_classes), 114 | ) 115 | 116 | def forward(self, x): 117 | x1 = self.features(x) 118 | x11 = x1.view(x1.size(0), -1) 119 | x = self.classifier(x11) 120 | return x 121 | 122 | class fh02(nn.Module): 123 | def __init__(self, num_points, num_classes, wrPath=None): 124 | super(fh02, self).__init__() 125 | self.load_wR2(wrPath) 126 | self.classifier1 = nn.Sequential( 127 | # nn.Dropout(), 128 | nn.Linear(53248, 128), 129 | # nn.ReLU(inplace=True), 130 | # nn.Dropout(), 131 | nn.Linear(128, provNum), 132 | ) 133 | self.classifier2 = nn.Sequential( 134 | # nn.Dropout(), 135 | nn.Linear(53248, 128), 136 | # nn.ReLU(inplace=True), 137 | # nn.Dropout(), 138 | nn.Linear(128, alphaNum), 139 | ) 140 | self.classifier3 = nn.Sequential( 141 | # nn.Dropout(), 142 | nn.Linear(53248, 128), 143 | # nn.ReLU(inplace=True), 144 | # nn.Dropout(), 145 | nn.Linear(128, adNum), 146 | ) 147 | self.classifier4 = nn.Sequential( 148 | # nn.Dropout(), 149 | nn.Linear(53248, 128), 150 | # nn.ReLU(inplace=True), 151 | # nn.Dropout(), 152 | nn.Linear(128, adNum), 153 | ) 154 | self.classifier5 = nn.Sequential( 155 | # nn.Dropout(), 156 | nn.Linear(53248, 128), 157 | # nn.ReLU(inplace=True), 158 | # nn.Dropout(), 159 | nn.Linear(128, adNum), 160 | ) 161 | self.classifier6 = nn.Sequential( 162 | # nn.Dropout(), 163 | nn.Linear(53248, 128), 164 | # nn.ReLU(inplace=True), 165 | # nn.Dropout(), 166 | nn.Linear(128, adNum), 167 | ) 168 | self.classifier7 = nn.Sequential( 169 | # nn.Dropout(), 170 | nn.Linear(53248, 128), 171 | # nn.ReLU(inplace=True), 172 | # nn.Dropout(), 173 | nn.Linear(128, adNum), 174 | ) 175 | 176 | def load_wR2(self, path): 177 | self.wR2 = wR2(numPoints) 178 | self.wR2 = torch.nn.DataParallel(self.wR2, device_ids=range(torch.cuda.device_count())) 179 | if not path is None: 180 | self.wR2.load_state_dict(torch.load(path)) 181 | # self.wR2 = self.wR2.cuda() 182 | # for param in self.wR2.parameters(): 183 | # param.requires_grad = False 184 | 185 | def forward(self, x): 186 | x0 = self.wR2.module.features[0](x) 187 | _x1 = self.wR2.module.features[1](x0) 188 | x2 = self.wR2.module.features[2](_x1) 189 | _x3 = self.wR2.module.features[3](x2) 190 | x4 = self.wR2.module.features[4](_x3) 191 | _x5 = self.wR2.module.features[5](x4) 192 | 193 | x6 = self.wR2.module.features[6](_x5) 194 | x7 = self.wR2.module.features[7](x6) 195 | x8 = self.wR2.module.features[8](x7) 196 | x9 = self.wR2.module.features[9](x8) 197 | x9 = x9.view(x9.size(0), -1) 198 | boxLoc = self.wR2.module.classifier(x9) 199 | 200 | h1, w1 = _x1.data.size()[2], _x1.data.size()[3] 201 | p1 = torch.FloatTensor([[w1, 0, 0, 0], [0, h1, 0, 0], [0, 0, w1, 0], [0, 0, 0, h1]]).cuda() 202 | h2, w2 = _x3.data.size()[2], _x3.data.size()[3] 203 | p2 = torch.FloatTensor([[w2, 0, 0, 0], [0, h2, 0, 0], [0, 0, w2, 0], [0, 0, 0, h2]]).cuda() 204 | h3, w3 = _x5.data.size()[2], _x5.data.size()[3] 205 | p3 = torch.FloatTensor([[w3, 0, 0, 0], [0, h3, 0, 0], [0, 0, w3, 0], [0, 0, 0, h3]]).cuda() 206 | 207 | # x, y, w, h --> x1, y1, x2, y2 208 | assert boxLoc.data.size()[1] == 4 209 | postfix = torch.FloatTensor([[1, 0, 1, 0], [0, 1, 0, 1], [-0.5, 0, 0.5, 0], [0, -0.5, 0, 0.5]]).cuda() 210 | boxNew = boxLoc.mm(postfix).clamp(min=0, max=1) 211 | 212 | # input = torch.rand(2, 1, 10, 10) 213 | # rois = torch.LongTensor([[0, 1, 2, 7, 8], [0, 3, 3, 8, 8], [1, 3, 3, 8, 8]]) 214 | roi1 = roi_pooling_ims(_x1, boxNew.mm(p1), size=(16, 8)) 215 | roi2 = roi_pooling_ims(_x3, boxNew.mm(p2), size=(16, 8)) 216 | roi3 = roi_pooling_ims(_x5, boxNew.mm(p3), size=(16, 8)) 217 | rois = torch.cat((roi1, roi2, roi3), 1) 218 | 219 | _rois = rois.view(rois.size(0), -1) 220 | 221 | y0 = self.classifier1(_rois) 222 | y1 = self.classifier2(_rois) 223 | y2 = self.classifier3(_rois) 224 | y3 = self.classifier4(_rois) 225 | y4 = self.classifier5(_rois) 226 | y5 = self.classifier6(_rois) 227 | y6 = self.classifier7(_rois) 228 | return boxLoc, [y0, y1, y2, y3, y4, y5, y6] 229 | 230 | 231 | def isEqual(labelGT, labelP): 232 | compare = [1 if int(labelGT[i]) == int(labelP[i]) else 0 for i in range(7)] 233 | # print(sum(compare)) 234 | return sum(compare) 235 | 236 | def eval(model, test_dirs): 237 | count, error, correct = 0, 0, 0 238 | dst = labelTestDataLoader(test_dirs, imgSize) 239 | testloader = DataLoader(dst, batch_size=1, shuffle=True, num_workers=8) 240 | start = time() 241 | for i, (XI, labels, ims) in enumerate(testloader): 242 | count += 1 243 | YI = [[int(ee) for ee in el.split('_')[:7]] for el in labels] 244 | if use_gpu: 245 | x = XI.cuda(0) 246 | else: 247 | x = XI 248 | # Forward pass: Compute predicted y by passing x to the model 249 | 250 | fps_pred, y_pred = model(x) 251 | 252 | outputY = [el.data.cpu().numpy().tolist() for el in y_pred] 253 | labelPred = [t[0].index(max(t[0])) for t in outputY] 254 | 255 | # compare YI, outputY 256 | try: 257 | if isEqual(labelPred, YI[0]) == 7: 258 | correct += 1 259 | else: 260 | pass 261 | except: 262 | error += 1 263 | return count, correct, error, float(correct) / count, (time() - start) / count 264 | 265 | #训练模型 266 | def train_model(model, criterion, optimizer, num_epochs=25): 267 | for epoch in range(epoch_start, num_epochs): 268 | lossAver = [] 269 | model.train(True) 270 | start = time() 271 | 272 | for i, (XI, Y, labels, ims) in enumerate(trainloader): 273 | if not len(XI) == batchSize: 274 | 275 | continue 276 | # labels为对应车牌号(0_0_22_27_27_33_16) 277 | YI = [[int(ee) for ee in el.split('_')[:7]] for el in labels] 278 | Y = np.array([el.numpy() for el in Y]).T # 真实值[cen_x,cen_y,w,h] 279 | if use_gpu: 280 | x = XI.cuda(0) 281 | y = torch.FloatTensor(Y).cuda(0) 282 | else: 283 | x = XI 284 | y = torch.FloatTensor(Y) 285 | # Forward pass: Compute predicted y by passing x to the model 286 | 287 | try: 288 | fps_pred, y_pred = model(x) # fps_pred为预测的[px,py,ph,pw]，y_pred为车牌7位号码的预测值 289 | except: 290 | continue 291 | 292 | # Compute and print loss 293 | loss = 0.0 294 | loss += 0.8 * nn.L1Loss().cuda()(fps_pred[:][:2], y[:][:2]) # 定位cen_x和cen_y损失 295 | loss += 0.2 * nn.L1Loss().cuda()(fps_pred[:][2:], y[:][2:]) # 定位w和h损失 296 | for j in range(7): #每个号码牌的交叉熵损失 297 | l = torch.LongTensor([el[j] for el in YI]).cuda(0) 298 | loss += criterion(y_pred[j], l) # 分类损失 299 | 300 | # Zero gradients, perform a backward pass, and update the weights. 301 | optimizer.zero_grad() 302 | loss.backward() 303 | optimizer.step() 304 | lrScheduler.step() 305 | 306 | try: 307 | lossAver.append(loss.item()) 308 | except: 309 | pass 310 | 311 | if i % 50 == 1: 312 | with open(args['writeFile'], 'a') as outF: 313 | outF.write('train %s images, use %s seconds, loss %s\n' % ( 314 | i * batchSize, time() - start, 315 | sum(lossAver) / len(lossAver) if len(lossAver) > 0 else 'NoLoss')) 316 | torch.save(model.state_dict(), storeName) 317 | print('%s %s %s\n' % (epoch, sum(lossAver) / len(lossAver), time() - start)) 318 | model.eval() 319 | count, correct, error, precision, avgTime = eval(model, testDirs) 320 | with open(args['writeFile'], 'a') as outF: 321 | outF.write('%s %s %s\n' % (epoch, sum(lossAver) / len(lossAver), time() - start)) 322 | outF.write('*** total %s error %s precision %s avgTime %s\n' % (count, error, precision, avgTime)) 323 | torch.save(model.state_dict(), storeName + str(epoch)) 324 | return model 325 | 326 | 327 | if __name__ == '__main__': 328 | ap = argparse.ArgumentParser() 329 | ap.add_argument("-i", "--images", required=True, 330 | help="path to the input file") 331 | ap.add_argument("-n", "--epochs", default=10000, 332 | help="epochs for train") 333 | ap.add_argument("-b", "--batchsize", default=2, 334 | help="batch size for train") 335 | ap.add_argument("-se", "--start_epoch", default=0, 336 | help="start epoch for train") 337 | ap.add_argument("-t", "--test", required=True, 338 | help="dirs for test") 339 | ap.add_argument("-r", "--resume", default='111', 340 | help="file for re-train") 341 | ap.add_argument("-f", "--folder", default='weight/CCPD_checkpoints/', 342 | help="folder to store model") 343 | ap.add_argument("-w", "--writeFile", default='weight/fh02.out', 344 | help="file for output") 345 | ap.add_argument("--visdom",type=str,default="rpnet",help="name of visdom") 346 | args = vars(ap.parse_args()) 347 | vis=Visdom(args["visdom"]) 348 | 349 | wR2Path = 'weight/wR2.pth' 350 | use_gpu = torch.cuda.is_available() 351 | print(use_gpu) 352 | 353 | numClasses = 7 # 车牌位数为7位 354 | numPoints = 4 # 定位点数为4个 355 | imgSize = (480, 480) # 图片大小为480*480 356 | provNum, alphaNum, adNum = 38, 25, 35 # 省份类别数量，市区类别数量，字符数量 357 | batchSize = int(args["batchsize"]) if use_gpu else 2 358 | trainDirs = args["images"].split(',') 359 | testDirs = args["test"].split(',') 360 | modelFolder = str(args["folder"]) if str(args["folder"])[-1] == '/' else str(args["folder"]) + '/' 361 | storeName = modelFolder + 'fh02.pth' 362 | if not os.path.isdir(modelFolder): 363 | os.mkdir(modelFolder) 364 | 365 | epochs = int(args["epochs"]) 366 | # initialize the output file 367 | if not os.path.isfile(args['writeFile']): 368 | with open(args['writeFile'], 'wb') as outF: 369 | pass 370 | 371 | epoch_start = int(args["start_epoch"]) 372 | resume_file = str(args["resume"]) 373 | if not resume_file == '111': # 再训练 374 | # epoch_start = int(resume_file[resume_file.find('pth') + 3:]) + 1 375 | if not os.path.isfile(resume_file): 376 | print("fail to load existed model! Existing ...") 377 | exit(0) 378 | print("Load existed model! %s" % resume_file) 379 | model_conv = fh02(numPoints, numClasses) 380 | model_conv = torch.nn.DataParallel(model_conv, device_ids=range(torch.cuda.device_count())) 381 | model_conv.load_state_dict(torch.load(resume_file)) 382 | model_conv = model_conv.cuda() 383 | else: 384 | model_conv = fh02(numPoints, numClasses, wR2Path) 385 | if use_gpu: 386 | model_conv = torch.nn.DataParallel(model_conv, device_ids=range(torch.cuda.device_count())) 387 | model_conv = model_conv.cuda() 388 | 389 | print(model_conv) 390 | print('模型参数量' + str(get_n_params(model_conv))) 391 | 392 | criterion = nn.CrossEntropyLoss() 393 | # optimizer_conv = optim.RMSprop(model_conv.parameters(), lr=0.01, momentum=0.9) 394 | optimizer_conv = optim.SGD(model_conv.parameters(), lr=0.001, momentum=0.9) 395 | 396 | dst = labelFpsDataLoader(trainDirs, imgSize) 397 | trainloader = DataLoader(dst, batch_size=batchSize, shuffle=True, num_workers=8) 398 | lrScheduler = lr_scheduler.StepLR(optimizer_conv, step_size=5, gamma=0.1) # 每隔5个epoch，学习率乘以0.1 399 | 400 | model_conv = train_model(model_conv, criterion, optimizer_conv, num_epochs=epochs) 401 | --------------------------------------------------------------------------------