像素聚合网络PAN原理与代码解析

最新推荐文章于 2024-07-23 18:01:13 发布

00000cj

最新推荐文章于 2024-07-23 18:01:13 发布

阅读量3.3k

点赞数

分类专栏：文本检测

本文链接：https://blog.csdn.net/ooooocj/article/details/109199242

版权

文本检测专栏收录该内容

7 篇文章 1 订阅

订阅专栏

Pipeline

Loss

像素聚合算法（Pixel Aggregation）

论文：https://arxiv.org/abs/1908.05900

官方代码：https://github.com/whai362/pan_pp.pytorch

像素聚合网络Pixel Aggregation Network是PSENet的改进版，依旧是segmentation-based文本检测方法，可以检测任意形状的文本。主要改进了PSENet速度慢的缺点，在CTW1500数据集上，PAN-320可以达到84.2FPS，同时还可以保证79.9%的F-measure。而PSENet-1s只有3.9FPS和78.0%的F-measure。

PAN主要做了两点改进来提升模型检测速度

用 ResNet-18 作为backbone，并提出了低计算量的 head 以解决因为使用 ResNet-18 而导致的特征提取能力较弱，进而带来的特征感受野较小且表征能力不足的缺点。
提出了一个可学习的后处理方法——像素聚合法，它能够通过预测出的相似向量来引导文字像素去纠正核参数。

Pipeline

1. Backbone 取ResNet18，假设Input的shape为(16,3,736,736)，16为batch_size，backbone从左到右输出的shape依次为(16,64,184,184)、(16,128,92,92)、(16,256,46,46)、(16,512,23,23)

2. Reducing Channel，每个backbone的输出接1*1*128conv、bn、relu得到 $F_{r}$ ， $F_{r}$ 从左到右的shape依次为(16,128,184,184)、(16,128,92,92)、(16,128,46,46)、(16,128,23,23)

3. FPEM

如上图，FPEM是一个 U形模组，由两个阶段组成，up-scale 增强、down-scale 增强。up-scale 增强作用于输入的特征金字塔，它以步长 32,16,8,4 像素在特征图上迭代增强。在 down-scale 阶段，输入的是由 up-scale 增强生成的特征金字塔，增强的步长从 4 到 32，同时，down-scale 增强输出的的特征金字塔就是最终 FPEM 的输出。

代码如下，输入f1~f4就是上一步的得到的 $F_{r}$

class FPEM_v1(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(FPEM_v1, self).__init__()
        planes = out_channels  # 128
        self.dwconv3_1 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, groups=planes, bias=False)
        self.smooth_layer3_1 = Conv_BN_ReLU(planes, planes)

        self.dwconv2_1 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, groups=planes, bias=False)
        self.smooth_layer2_1 = Conv_BN_ReLU(planes, planes)

        self.dwconv1_1 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, padding=1, groups=planes, bias=False)
        self.smooth_layer1_1 = Conv_BN_ReLU(planes, planes)

        self.dwconv2_2 = nn.Conv2d(planes, planes, kernel_size=3, stride=2, padding=1, groups=planes, bias=False)
        self.smooth_layer2_2 = Conv_BN_ReLU(planes, planes)

        self.dwconv3_2 = nn.Conv2d(planes, planes, kernel_size=3, stride=2, padding=1, groups=planes, bias=False)
        self.smooth_layer3_2 = Conv_BN_ReLU(planes, planes)

        self.dwconv4_2 = nn.Conv2d(planes, planes, kernel_size=3, stride=2, padding=1, groups=planes, bias=False)
        self.smooth_layer4_2 = Conv_BN_ReLU(planes, planes)

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

    def forward(self, f1, f2, f3, f4):
        f3 = self.smooth_layer3_1(self.dwconv3_1(self._upsample_add(f4, f3)))
        f2 = self.smooth_layer2_1(self.dwconv2_1(self._upsample_add(f3, f2)))
        f1 = self.smooth_layer1_1(self.dwconv1_1(self._upsample_add(f2, f1)))

        f2 = self.smooth_layer2_2(self.dwconv2_2(self._upsample_add(f2, f1)))
        f3 = self.smooth_layer3_2(self.dwconv3_2(self._upsample_add(f3, f2)))
        f4 = self.smooth_layer4_2(self.dwconv4_2(self._upsample_add(f4, f3)))

        return f1, f2, f3, f4

4. FFM

FPEM是可以级联的模块，官方代码中堆叠了两个FPEM。FFM用于融合不同深度的特征金字塔，首先通过逐元素相加结合了每个FPEM输出中相应 scale 的特征图，然后对特征图进行降采样，连接成最终 4*128 通道的特征图。

        # FPEM
        f1_1, f2_1, f3_1, f4_1 = self.fpem1(f1, f2, f3, f4)
        f1_2, f2_2, f3_2, f4_2 = self.fpem2(f1_1, f2_1, f3_1, f4_1)

        # FFM
        f1 = f1_1 + f1_2
        f2 = f2_1 + f2_2
        f3 = f3_1 + f3_2
        f4 = f4_1 + f4_2
        f2 = self._upsample(f2, f1.size())
        f3 = self._upsample(f3, f1.size())
        f4 = self._upsample(f4, f1.size())
        f = torch.cat((f1, f2, f3, f4), 1)  # torch.Size([16, 512, 184, 184])

5. 最后经过3*3*128conv、bn、relu、1*1*num_class的conv，最后再stride=4 upsample成原始输入大小即得到最终输出。这里的num_class=6，其中0通道对应完整文本、1对应kernel、2~5对应similar vector

Loss

完整的Loss函数如下

其中 $L_{tex}$ 和 $L_{ker}$ 分别是完整文本区域和核的损失，计算方法和PSENet一样。

只有一点不同，计算 $L_{ker}$ 时的mask不同

# PSENet取网络输出的完整本文预测图中大于0.5的区域，同时排除标注为忽略的部分
mask0 = torch.sigmoid(texts).data.cpu().numpy()
mask1 = training_masks.data.cpu().numpy()
selected_masks = ((mask0 > 0.5) & (mask1 > 0.5)).astype('float32')

# PAN取完整文本标注的区域，同时排除标注为忽略的部分
selected_masks = gt_texts * training_masks

这里对照代码重点讲一下 $L_{agg}$ 和 $L_{dis}$

其中N是文本实例的个数， $T_{i}$ 是第i个文本实例， $K_{i}$ 是第i个文本实例对应的kernel， $D(p, K_{i})$ 定义了文本实例 $T_{i}$ 内的像素p和 $K_{i}$ 之间的距离，公式如下

官方计算 $L_{agg}$ 代码解读

def forward_single(self, emb, instance, kernel, training_mask):
    # emb就是similar vector，shape为(4,w,h)，wh为网络输入的宽和高
    # instance是文本实例的ground truth，第i个文本实例区域的值为i，背景的值为0，shape为(w,h)
    # kernel是instance中每个文本实例缩放后的图，且每个文本实例的kernel区域的都为1，背景的值为0，shape为(w,h)
    # training_mask是标注为DO NOT CARE的文本实例区域值为0，其余部分值为1的图，shape为(w,h)
    training_mask = (training_mask > 0.5).long()
    kernel = (kernel > 0.5).long()
    instance = instance * training_mask  # 去掉标注为忽略的文本实例区域
    instance_kernel = (instance * kernel).view(-1)  # 第i个kernel的值为i
    instance = instance.view(-1)
    emb = emb.view(self.feature_dim, -1)  # (4,541696)

    unique_labels, unique_ids = torch.unique(instance_kernel, sorted=True, return_inverse=True)
    # 假设图中有5个文本实例，unique_label==tensor([0, 1, 2, 3, 4，5], device='cuda:0')，0是背景
    num_instance = unique_labels.size(0)
    if num_instance <= 1:
        return 0

    emb_mean = emb.new_zeros((self.feature_dim, num_instance), dtype=torch.float32)  # shape=(4,6)
    for i, lb in enumerate(unique_labels):
        if lb == 0:  # 背景
            continue
        ind_k = instance_kernel == lb  # 第i个kernel所有像素的索引
        emb_mean[:, i] = torch.mean(emb[:, ind_k], dim=1)  # 公式中的G(Ki)

    l_agg = emb.new_zeros(num_instance, dtype=torch.float32)  # bug （自带的不是我加的）
    for i, lb in enumerate(unique_labels):  # 遍历每一个文本实例
        if lb == 0:  # 0是背景
            continue
        ind = instance == lb
        emb_ = emb[:, ind]  # 公式中的F(p)
        # 单个文本instance的所有像素的similar vector，例如这张图片的第一个文本instance共有1012个像素，每个像素的similar vector (4,)
        dist = (emb_ - emb_mean[:, i:i + 1]).norm(p=2, dim=0)  # (torch.Size([4, 1012]) - torch.Size([4, 1])) -> torch.Size([1012]) 这里是单个文本实例的每个像素的距离
        # 注意emb_mean[:, i].shape==torch.Size([4]), emb_mean[:, i:i+1].shape==torch.Size([4, 1]), 这里必须用后者，前者会报错
        dist = F.relu(dist - self.delta_v) ** 2
        l_agg[i] = torch.mean(torch.log(dist + 1.0))
    l_agg = torch.mean(l_agg[1:])  # 对应公式里求N个文本实例的平均值

$L_{dis}$ 的公式如下

其中

官方代码的 $L_{dis}$ 实现如下

if num_instance > 2:
    emb_interleave = emb_mean.permute(1, 0).repeat(num_instance, 1)
    emb_band = emb_mean.permute(1, 0).repeat(1, num_instance).view(-1, self.feature_dim)
    mask = (1 - torch.eye(num_instance, dtype=torch.int8)).view(-1, 1).repeat(1, self.feature_dim)
    mask = mask.view(num_instance, num_instance, -1)
    mask[0, :, :] = 0
    mask[:, 0, :] = 0
    mask = mask.view(num_instance * num_instance, -1)

    dist = emb_interleave - emb_band
    dist = dist[mask > 0].view(-1, self.feature_dim).norm(p=2, dim=1)
    dist = F.relu(2 * self.delta_d - dist) ** 2
    l_dis = torch.mean(torch.log(dist + 1.0))

注意代码中任意两个kernel间的距离计算了两遍

官方代码中还多算了一个论文中没有出现的loss

l_reg = torch.mean(torch.log(torch.norm(emb_mean, 2, 0) + 1.0)) * 0.001

作者的回答是 “l_reg是用来限制emb的模长不能太大，有没有这一项估计差别不是很大。”

像素聚合算法（Pixel Aggregation）

PA和PSE算法比较像，都是从最小的kernel往外expand获得完整的文本区域。不同点在于pse输出多个kernel，从小到大依次扩充，每一轮扩充的结束条件是当前kernel的所有像素都已扩充完。而pa只有一个kernel和一个完整text预测结果，扩充的条件是当前扩充像素点既在完整text预测区域内又满足和所属kernel的similar vector的欧式距离小于6（代码中为3）。官方代码pa使用pyx实现的，仿照其改成了python，代码如下并加了相应注释

def _pa(kernel, emb, label, cc, label_num, min_area=0):
    pred = np.zeros((label.shape[0], label.shape[1]), dtype=np.int32)
    mean_emb = np.zeros((label_num, 4), dtype=np.float32)
    area = np.full((label_num,), -1, dtype=np.float32)
    flag = np.zeros((label_num,), dtype=np.int32)
    inds = np.zeros((label_num, label.shape[0], label.shape[1]), dtype=np.uint8)
    p = np.zeros((label_num, 2), dtype=np.int32)

    max_rate = 1024
    for i in range(1, label_num):
        ind = label == i
        inds[i] = ind

        area[i] = np.sum(ind)  # 614.0

        if area[i] < min_area:  # 0
            label[ind] = 0
            continue

        px, py = np.where(ind)  # (614,),(614,)
        p[i] = (px[0], py[0])  # px[0]==min(px), py[0]==min(py)

        for j in range(1, i):
            if area[j] < min_area:
                continue
            if cc[p[i, 0], p[i, 1]] != cc[p[j, 0], p[j, 1]]:  # 完整的text预测图中没有把两个kernel合并成一个
                continue
            rate = area[i] / area[j]
            if rate < 1 / max_rate or rate > max_rate:
                flag[i] = 1
                mean_emb[i] = np.mean(emb[:, ind], axis=1)

                if flag[j] == 0:
                    flag[j] = 1
                    mean_emb[j] = np.mean(emb[:, inds[j].astype(np.bool)], axis=1)

    que = queue.Queue(maxsize=0)
    dx = [-1, 1, 0, 0]
    dy = [0, 0, -1, 1]

    points = np.array(np.where(label > 0)).transpose((1, 0))
    for point_idx in range(points.shape[0]):
        x, y = points[point_idx, 0], points[point_idx, 1]
        l = label[x, y]
        que.put((x, y, l))
        pred[x, y] = l

    while not que.empty():
        (x, y, l) = que.get()

        for j in range(4):
            tmpx = x + dx[j]
            tmpy = y + dy[j]
            if tmpx < 0 or tmpx >= label.shape[0] or tmpy < 0 or tmpy >= label.shape[1]:
                continue
            if kernel[0, tmpx, tmpy] == 0 or pred[tmpx, tmpy] > 0:  # 完整text预测图中这个点值为0或者已经扩充过了
                continue
            if flag[l] == 1 and np.linalg.norm(emb[:, tmpx, tmpy] - mean_emb[l]) > 3:  # 论文里是6
                continue

            que.put((tmpx, tmpy))
            pred[tmpx, tmpy] = l

    return pred


def pa(kernels, emb, min_area=0):  # (2, 184, 328)，(4, 184, 328)，0
    # kernels[0]是预测的text完整图，kernels[1]是预测的以0.5比例shrink的kernel图
    _, cc = cv2.connectedComponents(kernels[0], connectivity=4)
    label_num, label = cv2.connectedComponents(kernels[1], connectivity=4)  # label_num包含了背景，实际要-1

    return _pa(kernels[:-1], emb, label, cc, label_num, min_area)
    # (1, 184, 328)，(4, 184, 328)，(184, 328)，(184, 328)，2,3,0
    # kernels[0].shape=(184, 328), kernels[:-1].shape=(1, 184, 328)