1 Main Idea

2 Method

2.1 整体架构

• Backbone: SLANet采用轻量级backbonePP-LCNet作为体征提取器，提取多尺度（不同stage）的特征。输入图片尺寸为$$\mathbb{R}^{B \times 3 \times H \times W}$$，经过PP-LCNet 后，输出的多尺度特征：

• Neck : SLANet 采用CSP-PAN 将backbone提取的多尺度特征进行融合。输入多尺度特征$$[\mathbb{R}^{B\times 64 \times \frac{H}{4} \times \frac{W}{4}} , \mathbb{R}^{B\times 128 \times \frac{H}{8}\times \frac{W}{8}} , \mathbb{R}^{B\times 256 \times \frac{H}{16}\times \frac{W}{16}} , \mathbb{R}^{B\times 512 \times \frac{H}{32}\times \frac{W}{32}}]$$输出融合后的特征尺度为$$\mathbb{R}^{1\times 96 \times \frac{H}{32}\times \frac{W}{32}}$$

• Head : SLANet 采用GRU+Attention的形式来递归生成表格的html序列和单元格位置。

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np


class AttentionGRUCell(nn.Module):
    def __init__(self, input_size, hidden_size, num_embeddings, use_gru=False):
        super(AttentionGRUCell, self).__init__()
        self.i2h = nn.Linear(input_size, hidden_size, bias=False)
        self.h2h = nn.Linear(hidden_size, hidden_size)
        self.score = nn.Linear(hidden_size, 1, bias=False)
        
        self.rnn = nn.GRUCell(
            input_size=input_size + num_embeddings, hidden_size=hidden_size)

        self.hidden_size = hidden_size

    def forward(self, prev_hidden, batch_H, char_onehots):

        batch_H_proj = self.i2h(batch_H)
        prev_hidden_proj = torch.unsqueeze(self.h2h(prev_hidden), dim=1)

        res = torch.add(batch_H_proj, prev_hidden_proj)
        res = torch.tanh(res)
        e = self.score(res)

        alpha = F.softmax(e, dim=1)
        alpha = alpha.permute(0, 2, 1)
        # alpha = torch.transpose(alpha, [0, 2, 1])
        context = torch.squeeze(torch.matmul(alpha, batch_H), dim=1)
        concat_context = torch.concat([context, char_onehots], 1)

        cur_hidden = self.rnn(concat_context, prev_hidden)

        return (cur_hidden, cur_hidden), alpha



class SLAHead(nn.Module):
    def __init__(self,
                 in_channels,
                 hidden_size,
                 out_channels=30,
                 max_text_length=500,
                 loc_reg_num=4,
                 fc_decay=0.0,
                 **kwargs):
        """
        @param in_channels: input shape
        @param hidden_size: hidden_size for RNN and Embedding
        @param out_channels: num_classes to rec
        @param max_text_length: max text pred
        """
        super().__init__()
        in_channels = in_channels[-1]
        self.hidden_size = hidden_size
        self.max_text_length = max_text_length
        self.emb = self._char_to_onehot
        self.num_embeddings = out_channels
        self.loc_reg_num = loc_reg_num

        # structure
        self.structure_attention_cell = AttentionGRUCell(
            in_channels, hidden_size, self.num_embeddings)
        self.structure_generator = nn.Sequential(
            nn.Linear(self.hidden_size, self.hidden_size),
            nn.Linear(hidden_size, out_channels)
            )
        # loc
        self.loc_generator = nn.Sequential(
            nn.Linear(self.hidden_size, self.hidden_size),
            nn.Linear(self.hidden_size, loc_reg_num),
            nn.Sigmoid())

    def forward(self, inputs, targets=None):
        fea = inputs[-1]
        batch_size = fea.shape[0]
        # reshape
        fea = torch.reshape(fea, [fea.shape[0], fea.shape[1], -1])
        # fea = fea.transpose([0, 2, 1])  # (NTC)(batch, width, channels)
        fea = fea.permute(0, 2, 1)

        hidden = torch.zeros((batch_size, self.hidden_size))
        structure_preds = torch.zeros(
            (batch_size, self.max_text_length + 1, self.num_embeddings))
        loc_preds = torch.zeros(
            (batch_size, self.max_text_length + 1, self.loc_reg_num))
        structure_preds.requires_grad = False
        loc_preds.requires_grad = False
        if self.training and targets is not None:
            structure = targets[0]
            for i in range(self.max_text_length + 1):
                hidden, structure_step, loc_step = self._decode(structure[:, i],
                                                                fea, hidden)
                structure_preds[:, i, :] = structure_step
                loc_preds[:, i, :] = loc_step
        else:
            pre_chars = torch.zeros(batch_size, dtype=torch.int64)
            max_text_length = torch.tensor(self.max_text_length)
            # for export
            loc_step, structure_step = None, None
            for i in range(max_text_length + 1):
                hidden, structure_step, loc_step = self._decode(pre_chars, fea, hidden)
                pre_chars = structure_step.argmax(dim=1)
                structure_preds[:, i, :] = structure_step
                loc_preds[:, i, :] = loc_step
        if not self.training:
            structure_preds = F.softmax(structure_preds, dim=-1)
        return {'structure_probs': structure_preds, 'loc_preds': loc_preds}

    def _decode(self, pre_chars, features, hidden):
        """
        Predict table label and coordinates for each step
        @param pre_chars: Table label in previous step
        @param features:
        @param hidden: hidden status in previous step
        @return:
        """
        emb_feature = self.emb(pre_chars)
        # output shape is b * self.hidden_size
        (output, hidden), alpha = self.structure_attention_cell(
            hidden, features, emb_feature)

        # structure
        structure_step = self.structure_generator(output)
        # loc
        loc_step = self.loc_generator(output)
        return hidden, structure_step, loc_step

    def _char_to_onehot(self, input_char):
        input_ont_hot = F.one_hot(input_char, self.num_embeddings)
        return input_ont_hot

2.2 SLANet损失函数

注意：仅计算HTML序列为单元格处的单元格回归损失。代码中的实现是，仅当序列预测为 '<td>', '<td', '<td></td>' 是计算回归损失。

2.3 SLANet的标签体系介绍

• 引入<td></td> 来模拟空白单元格

• 引入<td 用于和rowspan，colpan来拼接，便于处理合并单元格的情况。

2.4 如何合并OCR的结果做表格结构还原

3 小结