PyTorch实现的MTCNN/LPRNet车牌识别

这是一个在MTCNN和LPRNet中使用PYTORCH的两阶段轻量级和健壮的车牌识别。

MTCNN是一个非常著名的实时检测模型，主要用于人脸识别。修改后用于车牌检测。LPRNet是另一种实时的端到端DNN，用于模糊识别.该网络以其优越的性能和较低的计算成本而不需要初步的字符分割。在这项工作中嵌入了Spatial Transformer Layer（空间变换层 LocNet），以便有更好的识别特性。

在Nivida Quadro P4000上使用该结构在CCPD数据集上能达到~ 80 ms/image的速度，识别准确率可达99%。下面是流程框架：

MTCNN

MTCNN 基础知识

MTCNN 一开始主要是拿来做人脸识别的（不知道现在还是不是，希望大佬可以分享最新的人脸识别网络）。MTCNN人脸检测是2016年的论文提出来的，MTCNN的“MT”是指多任务学习(Multi-Task)，在同一个任务中同时学习”识别人脸“、”边框回归“、”人脸关键点识别“。

1. 首先对test图片不断进行Resize，得到图片金字塔。按照resize_factor(如0.70，这个具体根据数据集人脸大小分布来确定，基本确定在0.70-0.80之间会比较合适，设的比较大，容易延长推理时间，小了容易漏掉一些中小型人脸）对test图片进行resize，直到大等于Pnet要求的12 * 12大小。这样子你会得到原图、原图 * resize_factor、原图* resize_factor^{2…、原图*resize_factor}n（注，最后一个的图片大小会大等于12）这些不同大小的图片，堆叠起来的话像是金字塔，简单称为图片金字塔。注意，这些图像都是要一幅幅输入到Pnet中去得到候选的。

2. 图片金字塔输入Pnet，得到大量的候选（candidate）。根据上述步骤1得到的图片金字塔，将所有图片输入到Pnet，得到输出map形状是（m, n, 16(2+4+10)）。根据分类得分，筛选掉一大部分的候选，再根据得到的4个偏移量对bbox进行校准后得到bbox的左上右下的坐标，对这些候选根据IOU值再进行非极大值抑制（NMS）筛选掉一大部分候选。详细的说就是根据分类得分从大到小排，得到（num_left, 4）的张量，即num_left个bbox的左上、右下绝对坐标。每次以队列里最大分数值的bbox坐标和剩余坐标求出iou，干掉iou大于0.6（阈值是提前设置的）的框，并把这个最大分数值移到最终结果。重复这个操作，会干掉很多有大量overlap的bbox，最终得到(num_left_after_nms, 16)个候选，这些候选需要根据bbox坐标去原图截出图片后，resize为24 * 24输入到Rnet。

3. 经过Pnet筛选出来的候选图片，经过Rnet进行精调。根据Pnet输出的坐标，去原图上截取出图片（截取图片有个细节是需要截取bbox最大边长的正方形，这是为了保障resize的时候不产生形变和保留更多的人脸框周围细节），resize为24 * 24，输入到Rnet，进行精调。Rnet仍旧会输出二分类one-hot2个输出、bbox的坐标偏移量4个输出、landmark10个输出，根据二分类得分干掉大部分不是人脸的候选、对截图的bbox进行偏移量调整后（说的简单点就是对左上右下的x、y坐标进行上下左右调整），再次重复Pnet所述的IOU NMS干掉大部分的候选。最终Pnet输出的也是（num_left_after_Rnet, 16），根据bbox的坐标再去原图截出图片输入到Onet，同样也是根据最大边长的正方形截取方法，避免形变和保留更多细节。

4. 经过Rnet干掉很多候选后的图片输入到Onet，输出准确的bbox坐标和landmark坐标。大体可以重复Pnet的过程，不过有区别的是这个时候我们除了关注bbox的坐标外，也要输出landmark的坐标。（有小伙伴会问，前面不关注landmark的输出吗？嗯，作者认为关注的很有限，前面之所以也有landmark坐标的输出，主要是希望能够联合landmark坐标使得bbox更精确，换言之，推理阶段的Pnet、Rnet完全可以不用输出landmark，Onet输出即可。当然，训练阶段Pnet、Rnet还是要关注landmark的）经过分类筛选、框调整后的NMS筛选，好的，至此我们就得到准确的人脸bbox坐标和landmark点了，任务完满结束。

MTCNN车牌检测

这项工作只使用proposal net(Pnet)和output net(Onet)，因为在这种情况下跳过Rnet不会损害准确性。Onet接受24(高度)x94(宽度)bGR图像，这与LPRNet的输入一致。

修改后的MTCNN结构如下：

MTCNN车牌检测网络结构

import torch
import torch.nn as nn
import torch.nn.functional as F
from collections import OrderedDict

class Flatten(nn.Module):

    def __init__(self):
        super(Flatten, self).__init__()

    def forward(self, x):
        """
        Arguments:
            x: a float tensor with shape [batch_size, c, h, w].
        Returns:
            a float tensor with shape [batch_size, c*h*w].
        """

        # without this pretrained model isn't working
        x = x.transpose(3, 2).contiguous()

        return x.view(x.size(0), -1)


class PNet(nn.Module):

    def __init__(self, is_train=False):

        super(PNet, self).__init__()
        self.is_train = is_train

        self.features = nn.Sequential(OrderedDict([
            ('conv1', nn.Conv2d(3, 10, 3, 1)),
            ('prelu1', nn.PReLU(10)),
            ('pool1', nn.MaxPool2d((2,5), ceil_mode=True)),

            ('conv2', nn.Conv2d(10, 16, (3,5), 1)),
            ('prelu2', nn.PReLU(16)),

            ('conv3', nn.Conv2d(16, 32, (3,5), 1)),
            ('prelu3', nn.PReLU(32))
        ]))

        self.conv4_1 = nn.Conv2d(32, 2, 1, 1)
        self.conv4_2 = nn.Conv2d(32, 4, 1, 1)

    def forward(self, x):
        """
        Arguments:
            x: a float tensor with shape [batch_size, 3, h, w].
        Returns:
            b: a float tensor with shape [batch_size, 4, h', w'].
            a: a float tensor with shape [batch_size, 2, h', w'].
        """
        x = self.features(x)
        a = self