原文:DBNET文字检测 - 知乎 (zhihu.com)
一、DBNET介绍
DBNET核心采用的是基于分割的做法进行文本检测,即将每个文本块都进行语义分割,然后对分割概率图进行简单二值化、最终转化得为box或者poly格式的检测结果。除去网络设计方面的差异,最大特点是引入了Differentiable Binarization(DB)模块来优化分割预测结果。常规的基于语义分割的文本检测算法都是直接输出二值语义概率图或者其他辅助信息,然后经过阈值二值化得到最终结果,要想得到比较好的文本检测效果,一般都需要复杂的后处理,例如PSENet和PANet,会导致速度很慢。DBNET将阈值二值化过程变得可微,这一小小改动不仅可以增加错误预测梯度,也可以联合优化各个分支,得到更好的语义概率图。
二、DBNET算法流程
与常规基于语义分割算法的区别是多了一条threshold map分支,该分支的主要目的是和分割图联合得到更接近二值化的二值图,属于辅助分支。
2.1、backbone
骨架网络采用的是resnet18或者resnet50,为了增加网络特征提取能力,在layer2、layer3和layer4模块内部引入了变形卷积dcnv2模块。
import torch.nn as nn
import math
import torch.utils.model_zoo as model_zoo
BatchNorm2d = nn.BatchNorm2d
__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 'deformable_resnet18', 'deformable_resnet50',
'resnet152']
model_urls = {
'resnet18': 'https://download.pytorch.org/modelss/resnet18-5c106cde.pth',
'resnet34': 'https://download.pytorch.org/modelss/resnet34-333f7ec4.pth',
'resnet50': 'https://download.pytorch.org/modelss/resnet50-19c8e357.pth',
'resnet101': 'https://download.pytorch.org/modelss/resnet101-5d3b4d8f.pth',
'resnet152': 'https://download.pytorch.org/modelss/resnet152-b121ed2d.pth',
}
def constant_init(module, constant, bias=0): # 常量初始化
nn.init.constant_(module.weight, constant)
if hasattr(module, 'bias'):
nn.init.constant_(module.bias, bias)
def conv3x3(in_planes, out_planes, stride=1):
"""3x3 convolution with padding"""
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
padding=1, bias=False)
class BasicBlock(nn.Module):
expansion = 1
def __init__(self, inplanes, planes, stride=1, downsample=None, dcn=None):
super(BasicBlock, self).__init__()
self.with_dcn = dcn is not None
self.conv1 = conv3x3(inplanes, planes, stride) # 正常卷积过程
self.bn1 = BatchNorm2d(planes)
self.relu = nn.ReLU(inplace=True)
self.with_modulated_dcn = False
if not self.with_dcn:
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, padding=1, bias=False)
else:
from torchvision.ops import DeformConv2d
deformable_groups = dcn.get('deformable_groups', 1)
offset_channels = 18
self.conv2_offset = nn.Conv2d(planes, deformable_groups * offset_channels, kernel_size=3, padding=1)
self.conv2 = DeformConv2d(planes, planes, kernel_size=3, padding=1, bias=False)
self.bn2 = BatchNorm2d(planes)
self.downsample = downsample
self.stride = stride
def forward(self, x):
residual = x
out = self.conv1(x)
out = self.bn1(out)
out = self.relu(out)
# out = self.conv2(out)
if not self.with_dcn:
out = self.conv2(out)
else:
offset = self.conv2_offset(out)
out = self.conv2(out, offset)
out = self.bn2(out)
if self.downsample is not None:
residual = self.downsample(x)
out += residual
out = self.relu(out)
return out
class Bottleneck(nn.Module):
expansion = 4
def __init__(self, inplanes, planes, stride=1, downsample=None, dcn=None):
super(Bottleneck, self).__init__()
self.with_dcn = dcn is not None
self.conv1 = nn.Con