pytorchOCR之DBnet（多类别文本检测1）_pytorch文本检测

pytorchOCR之DBnet（多类别文本检测1）

代码

完整代码

动机

主要是为了对不同文本类型，既要检测文本位置，又要分类文本类型，比如中英文检测分类，手写印刷检测分类。

实现方式

1 . 模型结构
这里对DBnet结构多增加了一个分类的支路，如下：

self.classhead = nn.Sequential(
            nn.Conv2d(inner_channels, inner_channels // 4, 3, padding=1, bias=bias),
            nn.BatchNorm2d(inner_channels // 4),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(inner_channels // 4, inner_channels // 4, 2, 2),
            nn.BatchNorm2d(inner_channels // 4),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(inner_channels // 4, n_classes, 2, 2)
            )
1
2
3
4
5
6
7
8
9

当然也可以直接在文本分割支路最后直接加上n_classes，这样可以共用更多卷积，但是实验下来没有分开效果好。
2. loss
这里多增加一个分类的loss如下：

class MulClassLoss(nn.Module):
    def __init__(self, ):
        super(MulClassLoss, self).__init__()
        
    def forward(self,pre_score,gt_score,n_class):
        gt_score = gt_score.reshape(-1)
        index = gt_score>0
        if index.sum()>0:
            pre_score = pre_score.permute(0,2,3,1).reshape(-1,n_class)
            gt_score = gt_score[index]
            gt_score = gt_score - 1
            pre_score = pre_score[index]
            class_loss = F.cross_entropy(pre_score,gt_score.long(), ignore_index=-1) 
        else:
            class_loss = torch.tensor(0.0).cuda()
        return class_loss
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

这里我们选取文本像素进行分类loss计算。

效果