This commit is contained in:
22
benchmark/PaddleOCR_DBNet/models/__init__.py
Normal file
22
benchmark/PaddleOCR_DBNet/models/__init__.py
Normal file
@@ -0,0 +1,22 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# @Time : 2019/8/23 21:55
|
||||
# @Author : zhoujun
|
||||
import copy
|
||||
from .model import Model
|
||||
from .losses import build_loss
|
||||
|
||||
__all__ = ["build_loss", "build_model"]
|
||||
support_model = ["Model"]
|
||||
|
||||
|
||||
def build_model(config):
|
||||
"""
|
||||
get architecture model class
|
||||
"""
|
||||
copy_config = copy.deepcopy(config)
|
||||
arch_type = copy_config.pop("type")
|
||||
assert (
|
||||
arch_type in support_model
|
||||
), f"{arch_type} is not developed yet!, only {support_model} are support now"
|
||||
arch_model = eval(arch_type)(copy_config)
|
||||
return arch_model
|
||||
25
benchmark/PaddleOCR_DBNet/models/backbone/__init__.py
Normal file
25
benchmark/PaddleOCR_DBNet/models/backbone/__init__.py
Normal file
@@ -0,0 +1,25 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# @Time : 2019/8/23 21:54
|
||||
# @Author : zhoujun
|
||||
|
||||
from .resnet import *
|
||||
|
||||
__all__ = ["build_backbone"]
|
||||
|
||||
support_backbone = [
|
||||
"resnet18",
|
||||
"deformable_resnet18",
|
||||
"deformable_resnet50",
|
||||
"resnet50",
|
||||
"resnet34",
|
||||
"resnet101",
|
||||
"resnet152",
|
||||
]
|
||||
|
||||
|
||||
def build_backbone(backbone_name, **kwargs):
|
||||
assert (
|
||||
backbone_name in support_backbone
|
||||
), f"all support backbone is {support_backbone}"
|
||||
backbone = eval(backbone_name)(**kwargs)
|
||||
return backbone
|
||||
366
benchmark/PaddleOCR_DBNet/models/backbone/resnet.py
Normal file
366
benchmark/PaddleOCR_DBNet/models/backbone/resnet.py
Normal file
@@ -0,0 +1,366 @@
|
||||
import math
|
||||
import paddle
|
||||
from paddle import nn
|
||||
|
||||
BatchNorm2d = nn.BatchNorm2D
|
||||
|
||||
__all__ = [
|
||||
"ResNet",
|
||||
"resnet18",
|
||||
"resnet34",
|
||||
"resnet50",
|
||||
"resnet101",
|
||||
"deformable_resnet18",
|
||||
"deformable_resnet50",
|
||||
"resnet152",
|
||||
]
|
||||
|
||||
model_urls = {
|
||||
"resnet18": "https://download.pytorch.org/models/resnet18-5c106cde.pth",
|
||||
"resnet34": "https://download.pytorch.org/models/resnet34-333f7ec4.pth",
|
||||
"resnet50": "https://download.pytorch.org/models/resnet50-19c8e357.pth",
|
||||
"resnet101": "https://download.pytorch.org/models/resnet101-5d3b4d8f.pth",
|
||||
"resnet152": "https://download.pytorch.org/models/resnet152-b121ed2d.pth",
|
||||
}
|
||||
|
||||
|
||||
def constant_init(module, constant, bias=0):
|
||||
module.weight = paddle.create_parameter(
|
||||
shape=module.weight.shape,
|
||||
dtype="float32",
|
||||
default_initializer=paddle.nn.initializer.Constant(constant),
|
||||
)
|
||||
if hasattr(module, "bias"):
|
||||
module.bias = paddle.create_parameter(
|
||||
shape=module.bias.shape,
|
||||
dtype="float32",
|
||||
default_initializer=paddle.nn.initializer.Constant(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_attr=False
|
||||
)
|
||||
|
||||
|
||||
class BasicBlock(nn.Layer):
|
||||
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, momentum=0.1)
|
||||
self.relu = nn.ReLU()
|
||||
self.with_modulated_dcn = False
|
||||
if not self.with_dcn:
|
||||
self.conv2 = nn.Conv2D(
|
||||
planes, planes, kernel_size=3, padding=1, bias_attr=False
|
||||
)
|
||||
else:
|
||||
from paddle.vision.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_attr=False
|
||||
)
|
||||
self.bn2 = BatchNorm2d(planes, momentum=0.1)
|
||||
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.Layer):
|
||||
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.Conv2D(inplanes, planes, kernel_size=1, bias_attr=False)
|
||||
self.bn1 = BatchNorm2d(planes, momentum=0.1)
|
||||
self.with_modulated_dcn = False
|
||||
if not self.with_dcn:
|
||||
self.conv2 = nn.Conv2D(
|
||||
planes, planes, kernel_size=3, stride=stride, padding=1, bias_attr=False
|
||||
)
|
||||
else:
|
||||
deformable_groups = dcn.get("deformable_groups", 1)
|
||||
from paddle.vision.ops import DeformConv2D
|
||||
|
||||
offset_channels = 18
|
||||
self.conv2_offset = nn.Conv2D(
|
||||
planes,
|
||||
deformable_groups * offset_channels,
|
||||
stride=stride,
|
||||
kernel_size=3,
|
||||
padding=1,
|
||||
)
|
||||
self.conv2 = DeformConv2D(
|
||||
planes, planes, kernel_size=3, padding=1, stride=stride, bias_attr=False
|
||||
)
|
||||
self.bn2 = BatchNorm2d(planes, momentum=0.1)
|
||||
self.conv3 = nn.Conv2D(planes, planes * 4, kernel_size=1, bias_attr=False)
|
||||
self.bn3 = BatchNorm2d(planes * 4, momentum=0.1)
|
||||
self.relu = nn.ReLU()
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
self.dcn = dcn
|
||||
self.with_dcn = dcn is not None
|
||||
|
||||
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)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
residual = self.downsample(x)
|
||||
|
||||
out += residual
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class ResNet(nn.Layer):
|
||||
def __init__(self, block, layers, in_channels=3, dcn=None):
|
||||
self.dcn = dcn
|
||||
self.inplanes = 64
|
||||
super(ResNet, self).__init__()
|
||||
self.out_channels = []
|
||||
self.conv1 = nn.Conv2D(
|
||||
in_channels, 64, kernel_size=7, stride=2, padding=3, bias_attr=False
|
||||
)
|
||||
self.bn1 = BatchNorm2d(64, momentum=0.1)
|
||||
self.relu = nn.ReLU()
|
||||
self.maxpool = nn.MaxPool2D(kernel_size=3, stride=2, padding=1)
|
||||
self.layer1 = self._make_layer(block, 64, layers[0])
|
||||
self.layer2 = self._make_layer(block, 128, layers[1], stride=2, dcn=dcn)
|
||||
self.layer3 = self._make_layer(block, 256, layers[2], stride=2, dcn=dcn)
|
||||
self.layer4 = self._make_layer(block, 512, layers[3], stride=2, dcn=dcn)
|
||||
|
||||
if self.dcn is not None:
|
||||
for m in self.modules():
|
||||
if isinstance(m, Bottleneck) or isinstance(m, BasicBlock):
|
||||
if hasattr(m, "conv2_offset"):
|
||||
constant_init(m.conv2_offset, 0)
|
||||
|
||||
def _make_layer(self, block, planes, blocks, stride=1, dcn=None):
|
||||
downsample = None
|
||||
if stride != 1 or self.inplanes != planes * block.expansion:
|
||||
downsample = nn.Sequential(
|
||||
nn.Conv2D(
|
||||
self.inplanes,
|
||||
planes * block.expansion,
|
||||
kernel_size=1,
|
||||
stride=stride,
|
||||
bias_attr=False,
|
||||
),
|
||||
BatchNorm2d(planes * block.expansion, momentum=0.1),
|
||||
)
|
||||
|
||||
layers = []
|
||||
layers.append(block(self.inplanes, planes, stride, downsample, dcn=dcn))
|
||||
self.inplanes = planes * block.expansion
|
||||
for i in range(1, blocks):
|
||||
layers.append(block(self.inplanes, planes, dcn=dcn))
|
||||
self.out_channels.append(planes * block.expansion)
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.relu(x)
|
||||
x = self.maxpool(x)
|
||||
|
||||
x2 = self.layer1(x)
|
||||
x3 = self.layer2(x2)
|
||||
x4 = self.layer3(x3)
|
||||
x5 = self.layer4(x4)
|
||||
|
||||
return x2, x3, x4, x5
|
||||
|
||||
|
||||
def load_torch_params(paddle_model, torch_patams):
|
||||
paddle_params = paddle_model.state_dict()
|
||||
|
||||
fc_names = ["classifier"]
|
||||
for key, torch_value in torch_patams.items():
|
||||
if "num_batches_tracked" in key:
|
||||
continue
|
||||
key = (
|
||||
key.replace("running_var", "_variance")
|
||||
.replace("running_mean", "_mean")
|
||||
.replace("module.", "")
|
||||
)
|
||||
torch_value = torch_value.detach().cpu().numpy()
|
||||
if key in paddle_params:
|
||||
flag = [i in key for i in fc_names]
|
||||
if any(flag) and "weight" in key: # ignore bias
|
||||
new_shape = [1, 0] + list(range(2, torch_value.ndim))
|
||||
print(
|
||||
f"name: {key}, ori shape: {torch_value.shape}, new shape: {torch_value.transpose(new_shape).shape}"
|
||||
)
|
||||
torch_value = torch_value.transpose(new_shape)
|
||||
paddle_params[key] = torch_value
|
||||
else:
|
||||
print(f"{key} not in paddle")
|
||||
paddle_model.set_state_dict(paddle_params)
|
||||
|
||||
|
||||
def load_models(model, model_name):
|
||||
import torch.utils.model_zoo as model_zoo
|
||||
|
||||
torch_patams = model_zoo.load_url(model_urls[model_name])
|
||||
load_torch_params(model, torch_patams)
|
||||
|
||||
|
||||
def resnet18(pretrained=True, **kwargs):
|
||||
"""Constructs a ResNet-18 model.
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
"""
|
||||
model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
|
||||
if pretrained:
|
||||
assert (
|
||||
kwargs.get("in_channels", 3) == 3
|
||||
), "in_channels must be 3 when pretrained is True"
|
||||
print("load from imagenet")
|
||||
load_models(model, "resnet18")
|
||||
return model
|
||||
|
||||
|
||||
def deformable_resnet18(pretrained=True, **kwargs):
|
||||
"""Constructs a ResNet-18 model.
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
"""
|
||||
model = ResNet(BasicBlock, [2, 2, 2, 2], dcn=dict(deformable_groups=1), **kwargs)
|
||||
if pretrained:
|
||||
assert (
|
||||
kwargs.get("in_channels", 3) == 3
|
||||
), "in_channels must be 3 when pretrained is True"
|
||||
print("load from imagenet")
|
||||
model.load_state_dict(model_zoo.load_url(model_urls["resnet18"]), strict=False)
|
||||
return model
|
||||
|
||||
|
||||
def resnet34(pretrained=True, **kwargs):
|
||||
"""Constructs a ResNet-34 model.
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
"""
|
||||
model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
|
||||
if pretrained:
|
||||
assert (
|
||||
kwargs.get("in_channels", 3) == 3
|
||||
), "in_channels must be 3 when pretrained is True"
|
||||
model.load_state_dict(model_zoo.load_url(model_urls["resnet34"]), strict=False)
|
||||
return model
|
||||
|
||||
|
||||
def resnet50(pretrained=True, **kwargs):
|
||||
"""Constructs a ResNet-50 model.
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
"""
|
||||
model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
|
||||
if pretrained:
|
||||
assert (
|
||||
kwargs.get("in_channels", 3) == 3
|
||||
), "in_channels must be 3 when pretrained is True"
|
||||
load_models(model, "resnet50")
|
||||
return model
|
||||
|
||||
|
||||
def deformable_resnet50(pretrained=True, **kwargs):
|
||||
"""Constructs a ResNet-50 model with deformable conv.
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
"""
|
||||
model = ResNet(Bottleneck, [3, 4, 6, 3], dcn=dict(deformable_groups=1), **kwargs)
|
||||
if pretrained:
|
||||
assert (
|
||||
kwargs.get("in_channels", 3) == 3
|
||||
), "in_channels must be 3 when pretrained is True"
|
||||
model.load_state_dict(model_zoo.load_url(model_urls["resnet50"]), strict=False)
|
||||
return model
|
||||
|
||||
|
||||
def resnet101(pretrained=True, **kwargs):
|
||||
"""Constructs a ResNet-101 model.
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
"""
|
||||
model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
|
||||
if pretrained:
|
||||
assert (
|
||||
kwargs.get("in_channels", 3) == 3
|
||||
), "in_channels must be 3 when pretrained is True"
|
||||
model.load_state_dict(model_zoo.load_url(model_urls["resnet101"]), strict=False)
|
||||
return model
|
||||
|
||||
|
||||
def resnet152(pretrained=True, **kwargs):
|
||||
"""Constructs a ResNet-152 model.
|
||||
Args:
|
||||
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
||||
"""
|
||||
model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
|
||||
if pretrained:
|
||||
assert (
|
||||
kwargs.get("in_channels", 3) == 3
|
||||
), "in_channels must be 3 when pretrained is True"
|
||||
model.load_state_dict(model_zoo.load_url(model_urls["resnet152"]), strict=False)
|
||||
return model
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
x = paddle.zeros([2, 3, 640, 640])
|
||||
net = resnet50(pretrained=True)
|
||||
y = net(x)
|
||||
for u in y:
|
||||
print(u.shape)
|
||||
|
||||
print(net.out_channels)
|
||||
40
benchmark/PaddleOCR_DBNet/models/basic.py
Normal file
40
benchmark/PaddleOCR_DBNet/models/basic.py
Normal file
@@ -0,0 +1,40 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# @Time : 2019/12/6 11:19
|
||||
# @Author : zhoujun
|
||||
from paddle import nn
|
||||
|
||||
|
||||
class ConvBnRelu(nn.Layer):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride=1,
|
||||
padding=0,
|
||||
dilation=1,
|
||||
groups=1,
|
||||
bias=True,
|
||||
padding_mode="zeros",
|
||||
inplace=True,
|
||||
):
|
||||
super().__init__()
|
||||
self.conv = nn.Conv2D(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
kernel_size=kernel_size,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
bias_attr=bias,
|
||||
padding_mode=padding_mode,
|
||||
)
|
||||
self.bn = nn.BatchNorm2D(out_channels)
|
||||
self.relu = nn.ReLU()
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv(x)
|
||||
x = self.bn(x)
|
||||
x = self.relu(x)
|
||||
return x
|
||||
132
benchmark/PaddleOCR_DBNet/models/head/DBHead.py
Normal file
132
benchmark/PaddleOCR_DBNet/models/head/DBHead.py
Normal file
@@ -0,0 +1,132 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# @Time : 2019/12/4 14:54
|
||||
# @Author : zhoujun
|
||||
import paddle
|
||||
from paddle import nn, ParamAttr
|
||||
|
||||
|
||||
class DBHead(nn.Layer):
|
||||
def __init__(self, in_channels, out_channels, k=50):
|
||||
super().__init__()
|
||||
self.k = k
|
||||
self.binarize = nn.Sequential(
|
||||
nn.Conv2D(
|
||||
in_channels,
|
||||
in_channels // 4,
|
||||
3,
|
||||
padding=1,
|
||||
weight_attr=ParamAttr(initializer=nn.initializer.KaimingNormal()),
|
||||
),
|
||||
nn.BatchNorm2D(
|
||||
in_channels // 4,
|
||||
weight_attr=ParamAttr(initializer=nn.initializer.Constant(1)),
|
||||
bias_attr=ParamAttr(initializer=nn.initializer.Constant(1e-4)),
|
||||
),
|
||||
nn.ReLU(),
|
||||
nn.Conv2DTranspose(
|
||||
in_channels // 4,
|
||||
in_channels // 4,
|
||||
2,
|
||||
2,
|
||||
weight_attr=ParamAttr(initializer=nn.initializer.KaimingNormal()),
|
||||
),
|
||||
nn.BatchNorm2D(
|
||||
in_channels // 4,
|
||||
weight_attr=ParamAttr(initializer=nn.initializer.Constant(1)),
|
||||
bias_attr=ParamAttr(initializer=nn.initializer.Constant(1e-4)),
|
||||
),
|
||||
nn.ReLU(),
|
||||
nn.Conv2DTranspose(
|
||||
in_channels // 4, 1, 2, 2, weight_attr=nn.initializer.KaimingNormal()
|
||||
),
|
||||
nn.Sigmoid(),
|
||||
)
|
||||
|
||||
self.thresh = self._init_thresh(in_channels)
|
||||
|
||||
def forward(self, x):
|
||||
shrink_maps = self.binarize(x)
|
||||
threshold_maps = self.thresh(x)
|
||||
if self.training:
|
||||
binary_maps = self.step_function(shrink_maps, threshold_maps)
|
||||
y = paddle.concat((shrink_maps, threshold_maps, binary_maps), axis=1)
|
||||
else:
|
||||
y = paddle.concat((shrink_maps, threshold_maps), axis=1)
|
||||
return y
|
||||
|
||||
def _init_thresh(self, inner_channels, serial=False, smooth=False, bias=False):
|
||||
in_channels = inner_channels
|
||||
if serial:
|
||||
in_channels += 1
|
||||
self.thresh = nn.Sequential(
|
||||
nn.Conv2D(
|
||||
in_channels,
|
||||
inner_channels // 4,
|
||||
3,
|
||||
padding=1,
|
||||
bias_attr=bias,
|
||||
weight_attr=ParamAttr(initializer=nn.initializer.KaimingNormal()),
|
||||
),
|
||||
nn.BatchNorm2D(
|
||||
inner_channels // 4,
|
||||
weight_attr=ParamAttr(initializer=nn.initializer.Constant(1)),
|
||||
bias_attr=ParamAttr(initializer=nn.initializer.Constant(1e-4)),
|
||||
),
|
||||
nn.ReLU(),
|
||||
self._init_upsample(
|
||||
inner_channels // 4, inner_channels // 4, smooth=smooth, bias=bias
|
||||
),
|
||||
nn.BatchNorm2D(
|
||||
inner_channels // 4,
|
||||
weight_attr=ParamAttr(initializer=nn.initializer.Constant(1)),
|
||||
bias_attr=ParamAttr(initializer=nn.initializer.Constant(1e-4)),
|
||||
),
|
||||
nn.ReLU(),
|
||||
self._init_upsample(inner_channels // 4, 1, smooth=smooth, bias=bias),
|
||||
nn.Sigmoid(),
|
||||
)
|
||||
return self.thresh
|
||||
|
||||
def _init_upsample(self, in_channels, out_channels, smooth=False, bias=False):
|
||||
if smooth:
|
||||
inter_out_channels = out_channels
|
||||
if out_channels == 1:
|
||||
inter_out_channels = in_channels
|
||||
module_list = [
|
||||
nn.Upsample(scale_factor=2, mode="nearest"),
|
||||
nn.Conv2D(
|
||||
in_channels,
|
||||
inter_out_channels,
|
||||
3,
|
||||
1,
|
||||
1,
|
||||
bias_attr=bias,
|
||||
weight_attr=ParamAttr(initializer=nn.initializer.KaimingNormal()),
|
||||
),
|
||||
]
|
||||
if out_channels == 1:
|
||||
module_list.append(
|
||||
nn.Conv2D(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size=1,
|
||||
stride=1,
|
||||
padding=1,
|
||||
bias_attr=True,
|
||||
weight_attr=ParamAttr(
|
||||
initializer=nn.initializer.KaimingNormal()
|
||||
),
|
||||
)
|
||||
)
|
||||
return nn.Sequential(module_list)
|
||||
else:
|
||||
return nn.Conv2DTranspose(
|
||||
in_channels,
|
||||
out_channels,
|
||||
2,
|
||||
2,
|
||||
weight_attr=ParamAttr(initializer=nn.initializer.KaimingNormal()),
|
||||
)
|
||||
|
||||
def step_function(self, x, y):
|
||||
return paddle.reciprocal(1 + paddle.exp(-self.k * (x - y)))
|
||||
13
benchmark/PaddleOCR_DBNet/models/head/__init__.py
Normal file
13
benchmark/PaddleOCR_DBNet/models/head/__init__.py
Normal file
@@ -0,0 +1,13 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# @Time : 2020/6/5 11:35
|
||||
# @Author : zhoujun
|
||||
from .DBHead import DBHead
|
||||
|
||||
__all__ = ["build_head"]
|
||||
support_head = ["DBHead"]
|
||||
|
||||
|
||||
def build_head(head_name, **kwargs):
|
||||
assert head_name in support_head, f"all support head is {support_head}"
|
||||
head = eval(head_name)(**kwargs)
|
||||
return head
|
||||
50
benchmark/PaddleOCR_DBNet/models/losses/DB_loss.py
Normal file
50
benchmark/PaddleOCR_DBNet/models/losses/DB_loss.py
Normal file
@@ -0,0 +1,50 @@
|
||||
import paddle
|
||||
from models.losses.basic_loss import BalanceCrossEntropyLoss, MaskL1Loss, DiceLoss
|
||||
|
||||
|
||||
class DBLoss(paddle.nn.Layer):
|
||||
def __init__(self, alpha=1.0, beta=10, ohem_ratio=3, reduction="mean", eps=1e-06):
|
||||
"""
|
||||
Implement PSE Loss.
|
||||
:param alpha: binary_map loss 前面的系数
|
||||
:param beta: threshold_map loss 前面的系数
|
||||
:param ohem_ratio: OHEM的比例
|
||||
:param reduction: 'mean' or 'sum'对 batch里的loss 算均值或求和
|
||||
"""
|
||||
super().__init__()
|
||||
assert reduction in ["mean", "sum"], " reduction must in ['mean','sum']"
|
||||
self.alpha = alpha
|
||||
self.beta = beta
|
||||
self.bce_loss = BalanceCrossEntropyLoss(negative_ratio=ohem_ratio)
|
||||
self.dice_loss = DiceLoss(eps=eps)
|
||||
self.l1_loss = MaskL1Loss(eps=eps)
|
||||
self.ohem_ratio = ohem_ratio
|
||||
self.reduction = reduction
|
||||
|
||||
def forward(self, pred, batch):
|
||||
shrink_maps = pred[:, 0, :, :]
|
||||
threshold_maps = pred[:, 1, :, :]
|
||||
binary_maps = pred[:, 2, :, :]
|
||||
loss_shrink_maps = self.bce_loss(
|
||||
shrink_maps, batch["shrink_map"], batch["shrink_mask"]
|
||||
)
|
||||
loss_threshold_maps = self.l1_loss(
|
||||
threshold_maps, batch["threshold_map"], batch["threshold_mask"]
|
||||
)
|
||||
metrics = dict(
|
||||
loss_shrink_maps=loss_shrink_maps, loss_threshold_maps=loss_threshold_maps
|
||||
)
|
||||
if pred.shape[1] > 2:
|
||||
loss_binary_maps = self.dice_loss(
|
||||
binary_maps, batch["shrink_map"], batch["shrink_mask"]
|
||||
)
|
||||
metrics["loss_binary_maps"] = loss_binary_maps
|
||||
loss_all = (
|
||||
self.alpha * loss_shrink_maps
|
||||
+ self.beta * loss_threshold_maps
|
||||
+ loss_binary_maps
|
||||
)
|
||||
metrics["loss"] = loss_all
|
||||
else:
|
||||
metrics["loss"] = loss_shrink_maps
|
||||
return metrics
|
||||
16
benchmark/PaddleOCR_DBNet/models/losses/__init__.py
Normal file
16
benchmark/PaddleOCR_DBNet/models/losses/__init__.py
Normal file
@@ -0,0 +1,16 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# @Time : 2020/6/5 11:36
|
||||
# @Author : zhoujun
|
||||
import copy
|
||||
from .DB_loss import DBLoss
|
||||
|
||||
__all__ = ["build_loss"]
|
||||
support_loss = ["DBLoss"]
|
||||
|
||||
|
||||
def build_loss(config):
|
||||
copy_config = copy.deepcopy(config)
|
||||
loss_type = copy_config.pop("type")
|
||||
assert loss_type in support_loss, f"all support loss is {support_loss}"
|
||||
criterion = eval(loss_type)(**copy_config)
|
||||
return criterion
|
||||
101
benchmark/PaddleOCR_DBNet/models/losses/basic_loss.py
Normal file
101
benchmark/PaddleOCR_DBNet/models/losses/basic_loss.py
Normal file
@@ -0,0 +1,101 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# @Time : 2019/12/4 14:39
|
||||
# @Author : zhoujun
|
||||
import paddle
|
||||
import paddle.nn as nn
|
||||
|
||||
|
||||
class BalanceCrossEntropyLoss(nn.Layer):
|
||||
"""
|
||||
Balanced cross entropy loss.
|
||||
Shape:
|
||||
- Input: :math:`(N, 1, H, W)`
|
||||
- GT: :math:`(N, 1, H, W)`, same shape as the input
|
||||
- Mask: :math:`(N, H, W)`, same spatial shape as the input
|
||||
- Output: scalar.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, negative_ratio=3.0, eps=1e-6):
|
||||
super(BalanceCrossEntropyLoss, self).__init__()
|
||||
self.negative_ratio = negative_ratio
|
||||
self.eps = eps
|
||||
|
||||
def forward(
|
||||
self,
|
||||
pred: paddle.Tensor,
|
||||
gt: paddle.Tensor,
|
||||
mask: paddle.Tensor,
|
||||
return_origin=False,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
pred: shape :math:`(N, 1, H, W)`, the prediction of network
|
||||
gt: shape :math:`(N, 1, H, W)`, the target
|
||||
mask: shape :math:`(N, H, W)`, the mask indicates positive regions
|
||||
"""
|
||||
positive = gt * mask
|
||||
negative = (1 - gt) * mask
|
||||
positive_count = int(positive.sum())
|
||||
negative_count = min(
|
||||
int(negative.sum()), int(positive_count * self.negative_ratio)
|
||||
)
|
||||
loss = nn.functional.binary_cross_entropy(pred, gt, reduction="none")
|
||||
positive_loss = loss * positive
|
||||
negative_loss = loss * negative
|
||||
negative_loss, _ = negative_loss.reshape([-1]).topk(negative_count)
|
||||
|
||||
balance_loss = (positive_loss.sum() + negative_loss.sum()) / (
|
||||
positive_count + negative_count + self.eps
|
||||
)
|
||||
|
||||
if return_origin:
|
||||
return balance_loss, loss
|
||||
return balance_loss
|
||||
|
||||
|
||||
class DiceLoss(nn.Layer):
|
||||
"""
|
||||
Loss function from https://arxiv.org/abs/1707.03237,
|
||||
where iou computation is introduced heatmap manner to measure the
|
||||
diversity between tow heatmaps.
|
||||
"""
|
||||
|
||||
def __init__(self, eps=1e-6):
|
||||
super(DiceLoss, self).__init__()
|
||||
self.eps = eps
|
||||
|
||||
def forward(self, pred: paddle.Tensor, gt, mask, weights=None):
|
||||
"""
|
||||
pred: one or two heatmaps of shape (N, 1, H, W),
|
||||
the losses of tow heatmaps are added together.
|
||||
gt: (N, 1, H, W)
|
||||
mask: (N, H, W)
|
||||
"""
|
||||
return self._compute(pred, gt, mask, weights)
|
||||
|
||||
def _compute(self, pred, gt, mask, weights):
|
||||
if len(pred.shape) == 4:
|
||||
pred = pred[:, 0, :, :]
|
||||
gt = gt[:, 0, :, :]
|
||||
assert pred.shape == gt.shape
|
||||
assert pred.shape == mask.shape
|
||||
if weights is not None:
|
||||
assert weights.shape == mask.shape
|
||||
mask = weights * mask
|
||||
intersection = (pred * gt * mask).sum()
|
||||
|
||||
union = (pred * mask).sum() + (gt * mask).sum() + self.eps
|
||||
loss = 1 - 2.0 * intersection / union
|
||||
assert loss <= 1
|
||||
return loss
|
||||
|
||||
|
||||
class MaskL1Loss(nn.Layer):
|
||||
def __init__(self, eps=1e-6):
|
||||
super(MaskL1Loss, self).__init__()
|
||||
self.eps = eps
|
||||
|
||||
def forward(self, pred: paddle.Tensor, gt, mask):
|
||||
loss = (paddle.abs(pred - gt) * mask).sum() / (mask.sum() + self.eps)
|
||||
return loss
|
||||
39
benchmark/PaddleOCR_DBNet/models/model.py
Normal file
39
benchmark/PaddleOCR_DBNet/models/model.py
Normal file
@@ -0,0 +1,39 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# @Time : 2019/8/23 21:57
|
||||
# @Author : zhoujun
|
||||
from addict import Dict
|
||||
from paddle import nn
|
||||
import paddle.nn.functional as F
|
||||
|
||||
from models.backbone import build_backbone
|
||||
from models.neck import build_neck
|
||||
from models.head import build_head
|
||||
|
||||
|
||||
class Model(nn.Layer):
|
||||
def __init__(self, model_config: dict):
|
||||
"""
|
||||
PANnet
|
||||
:param model_config: 模型配置
|
||||
"""
|
||||
super().__init__()
|
||||
model_config = Dict(model_config)
|
||||
backbone_type = model_config.backbone.pop("type")
|
||||
neck_type = model_config.neck.pop("type")
|
||||
head_type = model_config.head.pop("type")
|
||||
self.backbone = build_backbone(backbone_type, **model_config.backbone)
|
||||
self.neck = build_neck(
|
||||
neck_type, in_channels=self.backbone.out_channels, **model_config.neck
|
||||
)
|
||||
self.head = build_head(
|
||||
head_type, in_channels=self.neck.out_channels, **model_config.head
|
||||
)
|
||||
self.name = f"{backbone_type}_{neck_type}_{head_type}"
|
||||
|
||||
def forward(self, x):
|
||||
_, _, H, W = x.shape
|
||||
backbone_out = self.backbone(x)
|
||||
neck_out = self.neck(backbone_out)
|
||||
y = self.head(neck_out)
|
||||
y = F.interpolate(y, size=(H, W), mode="bilinear", align_corners=True)
|
||||
return y
|
||||
75
benchmark/PaddleOCR_DBNet/models/neck/FPN.py
Normal file
75
benchmark/PaddleOCR_DBNet/models/neck/FPN.py
Normal file
@@ -0,0 +1,75 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# @Time : 2019/9/13 10:29
|
||||
# @Author : zhoujun
|
||||
import paddle
|
||||
import paddle.nn.functional as F
|
||||
from paddle import nn
|
||||
|
||||
from models.basic import ConvBnRelu
|
||||
|
||||
|
||||
class FPN(nn.Layer):
|
||||
def __init__(self, in_channels, inner_channels=256, **kwargs):
|
||||
"""
|
||||
:param in_channels: 基础网络输出的维度
|
||||
:param kwargs:
|
||||
"""
|
||||
super().__init__()
|
||||
inplace = True
|
||||
self.conv_out = inner_channels
|
||||
inner_channels = inner_channels // 4
|
||||
# reduce layers
|
||||
self.reduce_conv_c2 = ConvBnRelu(
|
||||
in_channels[0], inner_channels, kernel_size=1, inplace=inplace
|
||||
)
|
||||
self.reduce_conv_c3 = ConvBnRelu(
|
||||
in_channels[1], inner_channels, kernel_size=1, inplace=inplace
|
||||
)
|
||||
self.reduce_conv_c4 = ConvBnRelu(
|
||||
in_channels[2], inner_channels, kernel_size=1, inplace=inplace
|
||||
)
|
||||
self.reduce_conv_c5 = ConvBnRelu(
|
||||
in_channels[3], inner_channels, kernel_size=1, inplace=inplace
|
||||
)
|
||||
# Smooth layers
|
||||
self.smooth_p4 = ConvBnRelu(
|
||||
inner_channels, inner_channels, kernel_size=3, padding=1, inplace=inplace
|
||||
)
|
||||
self.smooth_p3 = ConvBnRelu(
|
||||
inner_channels, inner_channels, kernel_size=3, padding=1, inplace=inplace
|
||||
)
|
||||
self.smooth_p2 = ConvBnRelu(
|
||||
inner_channels, inner_channels, kernel_size=3, padding=1, inplace=inplace
|
||||
)
|
||||
|
||||
self.conv = nn.Sequential(
|
||||
nn.Conv2D(self.conv_out, self.conv_out, kernel_size=3, padding=1, stride=1),
|
||||
nn.BatchNorm2D(self.conv_out),
|
||||
nn.ReLU(),
|
||||
)
|
||||
self.out_channels = self.conv_out
|
||||
|
||||
def forward(self, x):
|
||||
c2, c3, c4, c5 = x
|
||||
# Top-down
|
||||
p5 = self.reduce_conv_c5(c5)
|
||||
p4 = self._upsample_add(p5, self.reduce_conv_c4(c4))
|
||||
p4 = self.smooth_p4(p4)
|
||||
p3 = self._upsample_add(p4, self.reduce_conv_c3(c3))
|
||||
p3 = self.smooth_p3(p3)
|
||||
p2 = self._upsample_add(p3, self.reduce_conv_c2(c2))
|
||||
p2 = self.smooth_p2(p2)
|
||||
|
||||
x = self._upsample_cat(p2, p3, p4, p5)
|
||||
x = self.conv(x)
|
||||
return x
|
||||
|
||||
def _upsample_add(self, x, y):
|
||||
return F.interpolate(x, size=y.shape[2:]) + y
|
||||
|
||||
def _upsample_cat(self, p2, p3, p4, p5):
|
||||
h, w = p2.shape[2:]
|
||||
p3 = F.interpolate(p3, size=(h, w))
|
||||
p4 = F.interpolate(p4, size=(h, w))
|
||||
p5 = F.interpolate(p5, size=(h, w))
|
||||
return paddle.concat([p2, p3, p4, p5], axis=1)
|
||||
13
benchmark/PaddleOCR_DBNet/models/neck/__init__.py
Normal file
13
benchmark/PaddleOCR_DBNet/models/neck/__init__.py
Normal file
@@ -0,0 +1,13 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
# @Time : 2020/6/5 11:34
|
||||
# @Author : zhoujun
|
||||
from .FPN import FPN
|
||||
|
||||
__all__ = ["build_neck"]
|
||||
support_neck = ["FPN"]
|
||||
|
||||
|
||||
def build_neck(neck_name, **kwargs):
|
||||
assert neck_name in support_neck, f"all support neck is {support_neck}"
|
||||
neck = eval(neck_name)(**kwargs)
|
||||
return neck
|
||||
Reference in New Issue
Block a user