PSENet详解+代码解释+测试

论文地址:https://arxiv.org/pdf/1806.02559.pdf

开源代码pytorch版本:https://github.com/whai362/PSENet

本文主要内容

该模型主要想解决形状鲁棒的文本检测的挑战,主要包括两个方面的问题:1)已有的文本检测器难以准确定位任意形状的文本,难以将文本完美地包含在矩形框内;2)大部分基于分割的文本检测器难以将紧挨的文本实例分割开。该模型是一种基于分割的文本检测器,针对每个文本实例应用多个预测。这些预测与不同的“核(kernel)”相关,这些核是通过将原始文本实例收缩到多个尺度生成的。随后,利用先进的尺度扩展(PSE)算法得到最终的检测结果。PSE算法采用逐渐扩展核尺寸的方式,从小尺度到大尺度,以获取文本最大、完整的形状。使用这种方法,PSENet能够有效区分相邻的文本实例,对任意形状的文本都具有较好的鲁棒性。

传统的基于包围框回归的方法能够定位矩形或带角度的四边形文本,但无法有效检测具有任意形状的文本实例,如弯曲文本。像第二张图。

对于这类任意形状文本,基于语义分割的方法成为首选,但像素级别的分割方法无法处理空间上紧邻的两个文本实例,因为这会导致边界信息合并成一个文本实例。像第三张图。

但是,PSENet解决了上述问题,PSENet是基于分割的方法,能够定位任意形状的文本。提出PSE算法,能够将空间上紧邻的文本实例区分开。像第四张图

模型结构详解

PSENet(
  (backbone): ResNet(
    (conv1): Conv2d(3, 64, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
    (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (relu1): ReLU(inplace=True)
    (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (relu2): ReLU(inplace=True)
    (conv3): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
    (bn3): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (relu3): ReLU(inplace=True)
    (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
    (layer1): Sequential(
      (0): Bottleneck(
        (conv1): Conv2d(128, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
        (downsample): Sequential(
          (0): Conv2d(128, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
      )
      (1): Bottleneck(
        (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
      )
      (2): Bottleneck(
        (conv1): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv3): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
      )
    )
    (layer2): Sequential(
      (0): Bottleneck(
        (conv1): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
        (downsample): Sequential(
          (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)
          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
      )
      (1): Bottleneck(
        (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
      )
      (2): Bottleneck(
        (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
      )
      (3): Bottleneck(
        (conv1): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv3): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
      )
    )
    (layer3): Sequential(
      (0): Bottleneck(
        (conv1): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
        (downsample): Sequential(
          (0): Conv2d(512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False)
          (1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
      )
      (1): Bottleneck(
        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
      )
      (2): Bottleneck(
        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
      )
      (3): Bottleneck(
        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
      )
      (4): Bottleneck(
        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
      )
      (5): Bottleneck(
        (conv1): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv3): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
      )
    )
    (layer4): Sequential(
      (0): Bottleneck(
        (conv1): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
        (downsample): Sequential(
          (0): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False)
          (1): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        )
      )
      (1): Bottleneck(
        (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
      )
      (2): Bottleneck(
        (conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
        (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (relu): ReLU(inplace=True)
      )
    )
  )
  (fpn): FPN(
    (toplayer_): Conv_BN_ReLU(
      (conv): Conv2d(2048, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
    )
    (smooth1_): Conv_BN_ReLU(
      (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
    )
    (smooth2_): Conv_BN_ReLU(
      (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
    )
    (smooth3_): Conv_BN_ReLU(
      (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
    )
    (latlayer1_): Conv_BN_ReLU(
      (conv): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
    )
    (latlayer2_): Conv_BN_ReLU(
      (conv): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
    )
    (latlayer3_): Conv_BN_ReLU(
      (conv): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)
      (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (relu): ReLU(inplace=True)
    )
  )
  (det_head): PSENet_Head(
    (conv1): Conv2d(1024, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (relu1): ReLU(inplace=True)
    (conv2): Conv2d(256, 7, kernel_size=(1, 1), stride=(1, 1))
    (text_loss): DiceLoss()
    (kernel_loss): DiceLoss()
  )
)

ResNet Backbone:

这个模型的主干网络采用了 ResNet 结构,其中包含多个残差块(Bottleneck)。

conv1: 输入通道数为3(RGB图像),输出通道数为64,使用3x3的卷积核进行卷积,步幅为2。

后续通过堆叠多个残差块,逐渐增加通道数和减小空间维度。

FPN (Feature Pyramid Network):

FPN 用于融合不同层次的特征图,提供多尺度的信息。

toplayer_: 从ResNet的最高层(最后的Bottleneck块)得到一个256通道的特征图。

smooth1_, smooth2_, smooth3_: 分别用于平滑不同层次的特征图。

latlayer1_, latlayer2_, latlayer3_: 从不同层次提取的特征图。

PSENet_Head:

这是文本检测任务的头部,负责生成文本实例的掩码。

conv1: 用于进一步处理特征图,输出通道数为256。

conv2: 输出通道数为7,用于预测文本实例的几何信息,这可能包括几何形状等。

text_loss, kernel_loss: 使用 Dice Loss 作为文本和内核(kernel)的损失函数。

受到特征金字塔网络(FPN)的启发。PSENet使用高级特征图与低级特征图进行拼接,共有4个拼接特征图。这些特征图经过融合,包含了多个感受野的视角,有助于生成多尺度核。这些融合后的特征图通过映射到n个分支,生成n个分割结果S1,S2,...,Sn。每个Si代表在一个特定尺度上的文本实例的分割掩膜。分割掩膜的尺度由超参数确定。在这些掩膜中,S1代表最小尺度的文本实例分割结果,Sn表示初始分割掩膜(最大核)。在获取分割掩膜后,使用PSE算法逐渐扩展所有文本实例的核,从S1到Sn,最终得到检测结果R。

由于基于分割的方法对相邻很近的文本实例分别检测是很困难的。为了解决这个问题,提出了PSE算法

先进尺度扩展算法(PSE, Progressive Scale Expansion Algorithm)

受广度优先搜索(BFS)启发。使用三个不同尺度的分割结果S = {S₁, S₂, S₃},其中S₁是最小核。首先,在S₁上找到4个连接组件C = {c₁, c₂, c₃, c₄}。接着,通过合并S₂和S₃中的像素进行尺度扩展。最终,提取连接组件作为文本实例的最终预测结果。

尺度扩展的具体流程如下:

1、从多个核的像素开始,使用BFS算法以迭代方式合并邻近的文本像素。

2、处理可能的冲突像素,采用先来先服务的方式合并。

3、尺度扩展过程中的边界冲突不影响最终检测性能。

小辉问:能详细说说一下广度优先搜索(BFS)算法吗?

小G答:广度优先搜索(BFS)是一种图遍历算法,用于从图的某个起始节点开始,逐层访问该节点的邻居节点,直到访问到目标节点为止。该算法同时适用于树的层次遍历。

一、以下是广度优先搜索算法的基本步骤:

1、初始化队列: 创建一个空队列(通常是先进先出的队列),用于存储待访问的节点。

2、将起始节点入队: 将起始节点放入队列中。

3、迭代: 循环执行以下步骤,直到队列为空:

a. 出队并访问: 从队列中取出一个节点,访问该节点。

b. 将未访问的邻居节点入队: 对于该节点的每个邻居节点,如果它还没有被访问过且未入队,将其入队。

4、结束: 当队列为空时,算法结束。

二、BFS的关键点是确保节点按照它们的距离从起始节点排列,即先访问距离起始节点为1的节点,然后是距离为2的节点,以此类推。这样可以保证找到的路径是最短路径。

BFS常用于解决以下问题:

1、最短路径问题: 查找两个节点之间的最短路径。

2、连通性问题: 判断两个节点是否相互连通。

3、状态空间搜索: 在状态空间中搜索,例如在迷宫中寻找最短路径。

在PSENet训练过程中,需要为不同核尺度生成相应的分割结果和对应的标签值。为了获得标签值,采用了收缩原始文本实例的方法

标签生成(Label Generation)

1、初始多边形与分割标签: 在(b)中,蓝色边界多边形表示初始文本实例,对应最大的分割标签掩膜((c)中右上特征图)。

2、Vatti算法收缩: 为了获得(c)中收缩的掩膜序列,使用Vatti算法来收缩原始多边形,记为pn。算法通过向内收缩每个像素,得到一系列收缩后的多边形pi,如图(a)所示。

3、多边形到二值掩膜图像: 每个收缩多边形pi被转换为二值掩膜图像,作为分割标签的结果。标签特征图标记为G₁, G₂, ..., Gn。

上面两个是多边形Pn和Pi之间的间距计算和缩放比例计算公式,其中Area()表示多边形的面积,Perimeter()表示多边形的周长,参数m和n是算法的超参数,分别取值为0.5和6。

损失函数(Loss Function)

PSENet损失函数

这里Lc表示完整文本实例损失, Ls 表示收缩文本实例损失。λ表示Lc和Ls间平衡参数,论文中λ = 0.7。

文本实例采用Dice coefficient损失函数

L s 表示收缩文本实例的损失

总结:论文使用PSENet进行文本检测,其中的backbone是FPN。在网络的构建中,首先从FPN中得到4个256通道特征图(P2, P3, P4, P5),然后通过层叠和上采样等操作得到一个256通道的特征图F,使用了函数C(∗)得到输出特征图F。接下来,通过一系列卷积和上采样操作生成n个分割结果(S1, S2, ..., Sn)。在标签生成中,设置n=6, m=0.5,得到尺度{0.5, 0.6, 0.7, 0.8, 0.9, 1.0}。在训练中,模糊文本区域被忽略,损失平衡参数λ=0.7。Negative-positive比率被设置为3。训练数据进行了多种增强,包括随机缩放、水平翻转、旋转、随机剪切等。对四边形文本数据集,最终预测结果通过计算最小面积矩阵提取出包围框。对弯曲文本数据集,采用Ramer-Douglas-Peucker算法生成任意形状的包围框。

小辉问:能详细说说一下Ramer-Douglas-Peucker算法吗?

小G答:Ramer-Douglas-Peucker(RDP)算法是一种用于曲线近似的算法,通常用于减少离散曲线的点数,同时保持曲线的形状。RDP 算法的基本思想是通过在曲线上选择适当的数据点来近似原始曲线。它的步骤如下:

1、选择起始点和结束点: 在给定的曲线上选择起始点和结束点,并将它们添加到近似曲线中。

2、计算垂直距离: 对于曲线上的每个点,计算该点到直线段(由起始点和结束点组成)的垂直距离。保留距离最大的点及其索引。

3、判断垂直距离是否超过阈值: 如果最大垂直距离小于给定的阈值(误差容忍度),则在起始点和结束点之间的线段可以用一条直线来近似,无需添加其他点。

4、递归: 如果最大垂直距离大于阈值,则将距离最大的点作为新的端点,将曲线分成两个子曲线,对每个子曲线递归地应用 RDP 算法。

5、合并结果: 将所有递归步骤中选定的点合并,形成近似曲线。

RDP 算法的优点是能够在一定误差容忍度内保留曲线的形状,同时减少点的数量,从而降低存储和计算成本

训练及测试

使用的数据集:ICDAR2015

1、先到pse目录下,准备编译构建 C 扩展模块,这是由于pse的实现是用C++实现的

cd ./models/post_processing/pse/

2、使用 build_ext 命令,setuptools将按照 setup.py 中的配置构建 C 扩展模块。--inplace 的工作方式是基于当前工作目录,它会在当前工作目录中查找 setup.py 文件,并将构建好的扩展模块放置在这个目录中。这使得你可以在当前目录中直接测试和使用这些模块,而不需要将其安装到系统范围。

python setup.py build_ext --inplace

3、回到项目根目录

cd ../../../

4、环境中安装mmcv有点棘手,所以在这说一下,要装人家的东西先,用那个东西安装。mmcv 中包含了一系列用于图像和视频处理的模块,例如图像增强、数据加载、模型构建、训练和推断的功能。该库的设计目标是提供一个灵活且高效的工具箱,以支持研究和实际项目中的各种需求。

pip install -U openmim 
mim install mmcv-full

5、这个过程由于我用的都是新版的环境,以及自定义的数据集,等等所以改的地方比较多,有问题可以评论区问(自己的笔记本只有一块显卡,超级慢,但是代码支持多块显卡,我用的公司服务器4块显卡同时训练,效果还行)

CUDA_VISIBLE_DEVICES=0 python train.py config/psenet/psenet_r50_ic15_736.py

模型搭建代码解析

PSENet整体结构搭建

class PSENet(nn.Module):
    def __init__(self,
                 backbone,
                 neck,
                 detection_head):
        super(PSENet, self).__init__()
        # 构建模型的三个组件:backbone、neck、detection_head
        self.backbone = build_backbone(backbone)
        self.fpn = build_neck(neck)
        self.det_head = build_head(detection_head)

    def _upsample(self, x, size, scale=1):
        _, _, H, W = size
        # 使用双线性插值进行上采样
        return F.interpolate(x, size=(H // scale, W // scale), mode='bilinear')

    def forward(self,
                imgs,
                gt_texts=None,
                gt_kernels=None,
                training_masks=None,
                img_metas=None,
                cfg=None):
        # 存储模型输出的字典
        outputs = dict()

        # 如果不是在训练模式且需要报告速度,则记录开始时间
        if not self.training and cfg.report_speed:
            torch.cuda.synchronize()
            start = time.time()

        # backbone
        f = self.backbone(imgs)
        # 如果不是在训练模式且需要报告速度,则记录 backbone 运行时间
        if not self.training and cfg.report_speed:
            torch.cuda.synchronize()
            outputs.update(dict(
                backbone_time=time.time() - start
            ))
            start = time.time()

        # FPN
        f1, f2, f3, f4, = self.fpn(f[0], f[1], f[2], f[3])

        # 将 FPN 输出的特征图进行拼接
        f = torch.cat((f1, f2, f3, f4), 1)

        # 如果不是在训练模式且需要报告速度,则记录 FPN 运行时间
        if not self.training and cfg.report_speed:
            torch.cuda.synchronize()
            outputs.update(dict(
                neck_time=time.time() - start
            ))
            start = time.time()

        # detection
        det_out = self.det_head(f)

        # 如果不是在训练模式且需要报告速度,则记录 detection head 运行时间
        if not self.training and cfg.report_speed:
            torch.cuda.synchronize()
            outputs.update(dict(
                det_head_time=time.time() - start
            ))

        # 如果处于训练模式
        if self.training:
            # 对检测输出进行上采样,与标签匹配
            det_out = self._upsample(det_out, imgs.size())
            # 计算检测头的损失
            det_loss = self.det_head.loss(det_out, gt_texts, gt_kernels, training_masks)
            outputs.update(det_loss)
        else:
            # 对检测输出进行上采样,得到最终结果
            det_out = self._upsample(det_out, imgs.size(), 1)
            # 获取检测结果
            det_res = self.det_head.get_results(det_out, img_metas, cfg)
            outputs.update(det_res)

        return outputs

backbone

def resnet50(pretrained=False, **kwargs):
    """
    构建一个 ResNet-50 模型。

    参数:
        pretrained (bool): 如果为 True,则返回在 Places 数据集上预训练的模型。
        **kwargs: 其他关键字参数,用于传递给 ResNet 类的构造函数。
    """
    # 使用 ResNet 类构建 ResNet-50 模型,Bottleneck 表示残差块的类型,[3, 4, 6, 3] 表示每个残差阶段的块数量
    model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
    
    # 如果需要预训练的模型,加载预训练参数
    if pretrained:
        model.load_state_dict(load_url(model_urls['resnet50']), strict=False)
    
    return model
    class ResNet(nn.Module):

    def __init__(self, block, layers):
        super(ResNet, self).__init__()
        self.inplanes = 128
        
        # 第一层:卷积、批归一化、ReLU激活
        self.conv1 = conv3x3(3, 64, stride=2)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu1 = nn.ReLU(inplace=True)
        
        # 第二层:卷积、批归一化、ReLU激活
        self.conv2 = conv3x3(64, 64)
        self.bn2 = nn.BatchNorm2d(64)
        self.relu2 = nn.ReLU(inplace=True)
        
        # 第三层:卷积、批归一化、ReLU激活
        self.conv3 = conv3x3(64, 128)
        self.bn3 = nn.BatchNorm2d(128)
        self.relu3 = nn.ReLU(inplace=True)
        
        # 最大池化层,用于下采样
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)

        # 创建ResNet的四个层,每层包含多个基本块(BasicBlock)
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)

        # 初始化网络参数
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            # 如果需要下采样,构建下采样层
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, planes * block.expansion,
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(planes * block.expansion),
            )

        # 构建包含多个基本块的层
        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)

    def forward(self, x):
        # 前向传播过程
        # 第一层
        x = self.relu1(self.bn1(self.conv1(x)))        
        # 第二层
        x = self.relu2(self.bn2(self.conv2(x)))        
        # 第三层
        x = self.relu3(self.bn3(self.conv3(x)))        
        # 最大池化
        x = self.maxpool(x)
        # 存储每个阶段的特征
        f = []
        x = self.layer1(x)
        f.append(x)
        x = self.layer2(x)
        f.append(x)
        x = self.layer3(x)
        f.append(x)
        x = self.layer4(x)
        f.append(x)
        # 返回每个阶段的特征作为元组
        return tuple(f)
        
class Bottleneck(nn.Module):
    expansion = 4  # 扩张因子,用于指示此瓶颈块的通道数相对于输入通道数的倍数

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(Bottleneck, self).__init__()
        # 第一层卷积,1x1的卷积核,用于降维
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes)
        
        # 第二层卷积,3x3的卷积核,用于提取特征
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes)
        # 第三层卷积,1x1的卷积核,用于升维
        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(planes * 4)
        # ReLU激活函数,引入非线性特性
        self.relu = nn.ReLU(inplace=True)
        # downsample参数用于在需要时对输入进行下采样,以匹配维度
        self.downsample = downsample
        # 步幅参数,控制卷积层的步幅
        self.stride = stride

    def forward(self, x):
        residual = x  # 保存输入的残差连接
        # 第一层卷积、批归一化、ReLU激活
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        # 第二层卷积、批归一化、ReLU激活
        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)
        # 第三层卷积、批归一化
        out = self.conv3(out)
        out = self.bn3(out)
        # 如果存在下采样,则对输入进行下采样
        if self.downsample is not None:
            residual = self.downsample(x)
        # 将残差与处理后的特征相加
        out += residual        
        # 再次应用ReLU激活函数
        out = self.relu(out)
        return out

neck

class FPN(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(FPN, self).__init__()

        # Top layer,对应最高层级的特征图
        self.toplayer_ = Conv_BN_ReLU(2048, 256, kernel_size=1, stride=1, padding=0)

        # Smooth layers,对应下一层级的特征图,用于减少上采样的混叠效应
        self.smooth1_ = Conv_BN_ReLU(256, 256, kernel_size=3, stride=1, padding=1)
        self.smooth2_ = Conv_BN_ReLU(256, 256, kernel_size=3, stride=1, padding=1)
        self.smooth3_ = Conv_BN_ReLU(256, 256, kernel_size=3, stride=1, padding=1)

        # Lateral layers,对应低层级的特征图,用于传递信息到更高的层级
        self.latlayer1_ = Conv_BN_ReLU(1024, 256, kernel_size=1, stride=1, padding=0)
        self.latlayer2_ = Conv_BN_ReLU(512, 256, kernel_size=1, stride=1, padding=0)
        self.latlayer3_ = Conv_BN_ReLU(256, 256, kernel_size=1, stride=1, padding=0)

        # 初始化权重和标准化层参数
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def _upsample(self, x, y, scale=1):
        _, _, H, W = y.size()
        return F.interpolate(x, size=(H // scale, W // scale), mode='bilinear')

    def _upsample_add(self, x, y):
        _, _, H, W = y.size()
        return F.interpolate(x, size=(H, W), mode='bilinear') + y

    def forward(self, f2, f3, f4, f5):
        # Top-down pathway,自顶向下的路径
        p5 = self.toplayer_(f5)

        # Lateral connections,横向连接
        f4 = self.latlayer1_(f4)
        p4 = self._upsample_add(p5, f4)
        p4 = self.smooth1_(p4)

        f3 = self.latlayer2_(f3)
        p3 = self._upsample_add(p4, f3)
        p3 = self.smooth2_(p3)

        f2 = self.latlayer3_(f2)
        p2 = self._upsample_add(p3, f2)
        p2 = self.smooth3_(p2)

        # Upsample,上采样
        p3 = self._upsample(p3, p2)
        p4 = self._upsample(p4, p2)
        p5 = self._upsample(p5, p2)

        return p2, p3, p4, p5

head

class PSENet_Head(nn.Module):
    def __init__(self,
                 in_channels,
                 hidden_dim,
                 num_classes,
                 loss_text,
                 loss_kernel):
        super(PSENet_Head, self).__init__()
        # 定义卷积层和标准化层
        self.conv1 = nn.Conv2d(in_channels, hidden_dim, kernel_size=3, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(hidden_dim)
        self.relu1 = nn.ReLU(inplace=True)

        # 用于输出最终结果的卷积层
        self.conv2 = nn.Conv2d(hidden_dim, num_classes, kernel_size=1, stride=1, padding=0)

        # 定义文本和内核损失函数
        self.text_loss = build_loss(loss_text)
        self.kernel_loss = build_loss(loss_kernel)

        # 初始化权重和标准化层参数
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def forward(self, f):
        # 前向传播
        out = self.conv1(f)
        out = self.relu1(self.bn1(out))
        out = self.conv2(out)

        return out

    def get_results(self, out, img_meta, cfg):
        outputs = dict()

        # 如果不是在训练模式且需要报告速度,则记录开始时间
        if not self.training and cfg.report_speed:
            torch.cuda.synchronize()
            start = time.time()

        # 对输出进行处理,得到文本分数和内核信息
        score = torch.sigmoid(out[:, 0, :, :])
        kernels = out[:, :cfg.test_cfg.kernel_num, :, :] > 0
        text_mask = kernels[:, :1, :, :]
        kernels[:, 1:, :, :] = kernels[:, 1:, :, :] * text_mask

        # 将结果转为 numpy 数组
        score = score.data.cpu().numpy()[0].astype(np.float32)
        kernels = kernels.data.cpu().numpy()[0].astype(np.uint8)

        # 对预测的内核进行分割,得到文本的分割标签
        label = pse(kernels, cfg.test_cfg.min_area)

        # 获取原始图像尺寸和处理后的尺寸
        org_img_size = img_meta['org_img_size'][0]
        img_size = img_meta['img_size'][0]

        # 调整标签和分数的尺寸
        label_num = np.max(label) + 1
        label = cv2.resize(label, (img_size[1], img_size[0]), interpolation=cv2.INTER_NEAREST)
        score = cv2.resize(score, (img_size[1], img_size[0]), interpolation=cv2.INTER_NEAREST)

        # 如果不是在训练模式且需要报告速度,则记录处理时间
        if not self.training and cfg.report_speed:
            torch.cuda.synchronize()
            outputs.update(dict(
                det_pse_time=time.time() - start
            ))

        # 计算尺度变换比例
        scale = (float(org_img_size[1]) / float(img_size[1]),
                 float(org_img_size[0]) / float(img_size[0]))

        # 初始化存储边界框和得分的列表
        bboxes = []
        scores = []

        # 遍历标签,提取文本区域信息
        for i in range(1, label_num):
            ind = label == i
            points = np.array(np.where(ind)).transpose((1, 0))

            # 过滤面积小于阈值的区域
            if points.shape[0] < cfg.test_cfg.min_area:
                label[ind] = 0
                continue

            # 计算区域内的平均分数
            score_i = np.mean(score[ind])
            
            # 过滤分数低于阈值的区域
            if score_i < cfg.test_cfg.min_score:
                label[ind] = 0
                continue

            # 根据配置选择输出类型,矩形或多边形
            if cfg.test_cfg.bbox_type == 'rect':
                rect = cv2.minAreaRect(points[:, ::-1])
                bbox = cv2.boxPoints(rect) * scale
            elif cfg.test_cfg.bbox_type == 'poly':
                binary = np.zeros(label.shape, dtype='uint8')
                binary[ind] = 1
                _, contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
                bbox = contours[0] * scale

            bbox = bbox.astype('int32')
            bboxes.append(bbox.reshape(-1))
            scores.append(score_i)

        outputs.update(dict(
            bboxes=bboxes,
            scores=scores
        ))

        return outputs

    def loss(self, out, gt_texts, gt_kernels, training_masks):
        # 获取文本和内核的输出
        texts = out[:, 0, :, :]
        kernels = out[:, 1:, :, :]

        # 计算文本损失
        selected_masks = ohem_batch(texts, gt_texts, training_masks)
        loss_text = self.text_loss(texts, gt_texts, selected_masks, reduce=False)
        iou_text = iou((texts > 0).long(), gt_texts, training_masks, reduce=False)
        losses = dict(
            loss_text=loss_text,
            iou_text=iou_text
        )

        # 计算内核损失
        loss_kernels = []
        selected_masks = gt_texts * training_masks
        for i in range(kernels.size(1)):
            kernel_i = kernels[:, i, :, :]
            gt_kernel_i = gt_kernels[:, i, :, :]
            loss_kernel_i = self.kernel_loss(kernel_i, gt_kernel_i, selected_masks, reduce=False)
            loss_kernels.append(loss_kernel_i)
        loss_kernels = torch.mean(torch.stack(loss_kernels, dim=1), dim=1)
        iou_kernel = iou(
            (kernels[:, -1, :, :] > 0).long(), gt_kernels[:, -1, :, :], training_masks * gt_texts, reduce=False)
        losses.update(dict(
            loss_kernels=loss_kernels,
            iou_kernel=iou_kernel
        ))

        return losses

测试结果

测试结果我展示的是还可以的,其实还有少量不是那么好的....

评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值