前言:
本文将使用pytorch从0实现DBnet文本检测网络,关于DBnet网络的基础这里不在详细介绍,只注重代码编写。如果需要训练自己的数据集,可以将自己的数据集制作成本文所使用的格式即可。
网络搭建:
网络结构图:
特征提取采用resnet18,dbnet网络的输入大小为 batch x 3 x 640 x 640,输出大小为batch x 3 x 640 x 640,其中3个通道分别代表probability map,threshold map,approximate binary map。
代码如下所示:
from torchvision import models
from torchsummary import summary
from torchvision.models.feature_extraction import create_feature_extractor
import torch.nn as nn
import torch
resnet18=models.resnet18(pretrained=False)
resnet18 = create_feature_extractor(
resnet18, {
'relu':'feature2','layer1': 'feature4', 'layer2': 'feature8',
'layer3': 'feature16','layer4': 'feature32'})
class Upsample_add(nn.Module):
def __init__(self,ins,outs,x1):
super(Upsample_add, self).__init__()
self.x1=x1
self.up=nn.Sequential(
nn.ConvTranspose2d(ins, outs, kernel_size=2, stride=2