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))
当然也可以直接在文本分割支路最后直接加上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>0if 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 - 1pre_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
这里我们选取文本像素进行分类loss计算。
效果
本文链接：https://my.lmcjl.com/post/14200.html
展开阅读全文
4 评论