123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152 |
- import torch
- import torch.nn as nn
- import torch.nn.functional as F
- from torchvision.models import resnet101
- import DBHead
- import einops
- class ImageMultiheadSelfAttention(nn.Module):
- def __init__(self, planes):
- super(ImageMultiheadSelfAttention, self).__init__()
- self.attn = nn.MultiheadAttention(planes, 8)
- def forward(self, x):
- res = x
- n, c, h, w = x.shape
- x = einops.rearrange(x, 'n c h w -> (h w) n c')
- x = self.attn(x, x, x)[0]
- x = einops.rearrange(x, '(h w) n c -> n c h w', n = n, c = c, h = h, w = w)
- return res + x
- class double_conv(nn.Module):
- def __init__(self, in_ch, mid_ch, out_ch, stride = 1, planes = 256):
- super(double_conv, self).__init__()
- self.planes = planes
- # down = None
- # if stride > 1:
- # down = nn.Sequential(
- # nn.AvgPool2d(2, 2),
- # nn.Conv2d(in_ch + mid_ch, self.planes * Bottleneck.expansion, kernel_size=1, stride=1, bias=False),nn.BatchNorm2d(self.planes * Bottleneck.expansion)
- # )
- self.down = None
- if stride > 1:
- self.down = nn.AvgPool2d(2,stride=2)
- self.conv = nn.Sequential(
- nn.Conv2d(in_ch + mid_ch, mid_ch, kernel_size=3, padding=1, stride = 1, bias=False),
- nn.BatchNorm2d(mid_ch),
- nn.ReLU(inplace=True),
- #Bottleneck(mid_ch, self.planes, stride, down, 2, 1, avd = True, norm_layer = nn.BatchNorm2d),
- nn.Conv2d(mid_ch, out_ch, kernel_size=3, stride = 1, padding=1, bias=False),
- nn.BatchNorm2d(out_ch),
- nn.ReLU(inplace=True),
- )
- def forward(self, x):
- if self.down is not None:
- x = self.down(x)
- x = self.conv(x)
- return x
- class double_conv_up(nn.Module):
- def __init__(self, in_ch, mid_ch, out_ch, stride = 1, planes = 256):
- super(double_conv_up, self).__init__()
- self.planes = planes
- # down = None
- # if stride > 1:
- # down = nn.Sequential(
- # nn.AvgPool2d(2, 2),
- # nn.Conv2d(in_ch + mid_ch, self.planes * Bottleneck.expansion, kernel_size=1, stride=1, bias=False),nn.BatchNorm2d(self.planes * Bottleneck.expansion)
- # )
- self.down = None
- if stride > 1:
- self.down = nn.AvgPool2d(2,stride=2)
- self.conv = nn.Sequential(
- nn.Conv2d(in_ch + mid_ch, mid_ch, kernel_size=3, padding=1, stride = 1, bias=False),
- nn.BatchNorm2d(mid_ch),
- nn.ReLU(inplace=True),
- #Bottleneck(mid_ch, self.planes, stride, down, 2, 1, avd = True, norm_layer = nn.BatchNorm2d),
- nn.Conv2d(mid_ch, mid_ch, kernel_size=3, stride = 1, padding=1, bias=False),
- nn.BatchNorm2d(mid_ch),
- nn.ReLU(inplace=True),
- nn.ConvTranspose2d(mid_ch, out_ch, kernel_size=4, stride = 2, padding=1, bias=False),
- nn.BatchNorm2d(out_ch),
- nn.ReLU(inplace=True),
- )
- def forward(self, x):
- if self.down is not None:
- x = self.down(x)
- x = self.conv(x)
- return x
- class TextDetection(nn.Module):
- def __init__(self, pretrained=None):
- super(TextDetection, self).__init__()
- self.backbone = resnet101(pretrained=True if pretrained else False)
- self.conv_db = DBHead.DBHead(64, 0)
- self.conv_mask = nn.Sequential(
- nn.Conv2d(64, 64, kernel_size=3, padding=1), nn.ReLU(inplace=True),
- nn.Conv2d(64, 64, kernel_size=3, padding=1), nn.ReLU(inplace=True),
- nn.Conv2d(64, 32, kernel_size=3, padding=1), nn.ReLU(inplace=True),
- nn.Conv2d(32, 1, kernel_size=1),
- nn.Sigmoid()
- )
- self.down_conv1 = double_conv(0, 512, 512, 2)
- self.down_conv2 = double_conv(0, 512, 512, 2)
- self.down_conv3 = double_conv(0, 512, 512, 2)
- self.upconv1 = double_conv_up(0, 512, 256)
- self.upconv2 = double_conv_up(256, 512, 256)
- self.upconv3 = double_conv_up(256, 512, 256)
- self.upconv4 = double_conv_up(256, 512, 256, planes = 128)
- self.upconv5 = double_conv_up(256, 256, 128, planes = 64)
- self.upconv6 = double_conv_up(128, 128, 64, planes = 32)
- self.upconv7 = double_conv_up(64, 64, 64, planes = 16)
- self.proj_h4 = nn.Conv2d(64 * 4, 64, 1)
- self.proj_h8 = nn.Conv2d(128 * 4, 128, 1)
- self.proj_h16 = nn.Conv2d(256 * 4, 256, 1)
- self.proj_h32 = nn.Conv2d(512 * 4, 512, 1)
- def forward(self, x):
- x = self.backbone.conv1(x)
- x = self.backbone.bn1(x)
- x = self.backbone.relu(x)
- x = self.backbone.maxpool(x) # 64@384
- h4 = self.backbone.layer1(x) # 64@384
- h8 = self.backbone.layer2(h4) # 128@192
- h16 = self.backbone.layer3(h8) # 256@96
- h32 = self.backbone.layer4(h16) # 512@48
- h4 = self.proj_h4(h4)
- h8 = self.proj_h8(h8)
- h16 = self.proj_h16(h16)
- h32 = self.proj_h32(h32)
- h64 = self.down_conv1(h32) # 512@24
- h128 = self.down_conv2(h64) # 512@12
- h256 = self.down_conv3(h128) # 512@6
- up256 = self.upconv1(h256) # 128@12
- up128 = self.upconv2(torch.cat([up256, h128], dim = 1)) # 64@24
- up64 = self.upconv3(torch.cat([up128, h64], dim = 1)) # 128@48
- up32 = self.upconv4(torch.cat([up64, h32], dim = 1)) # 64@96
- up16 = self.upconv5(torch.cat([up32, h16], dim = 1)) # 128@192
- up8 = self.upconv6(torch.cat([up16, h8], dim = 1)) # 64@384
- up4 = self.upconv7(torch.cat([up8, h4], dim = 1)) # 64@768
- return self.conv_db(up8), self.conv_mask(up4)
- if __name__ == '__main__':
- device = torch.device("cuda:0")
- net = TextDetection().to(device)
- img = torch.randn(2, 3, 1024, 1024).to(device)
- db, seg = net(img)
- print(db.shape)
- print(seg.shape)
|