Spaces:
Sleeping
Sleeping
File size: 2,004 Bytes
14c9181 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
import torch
import torch.nn as nn
from mmcv.cnn import ConvModule
from mmengine.model import BaseModule
from mmocr.registry import MODELS
@MODELS.register_module()
class CoordinateHead(BaseModule):
def __init__(self,
in_channel=256,
conv_num=4,
norm_cfg=dict(type='BN'),
act_cfg=dict(type='ReLU'),
init_cfg=None):
super().__init__(init_cfg=init_cfg)
mask_convs = list()
for i in range(conv_num):
if i == 0:
mask_conv = ConvModule(
in_channels=in_channel + 2, # 2 for coord
out_channels=in_channel,
kernel_size=3,
padding=1,
norm_cfg=norm_cfg,
act_cfg=act_cfg)
else:
mask_conv = ConvModule(
in_channels=in_channel,
out_channels=in_channel,
kernel_size=3,
padding=1,
norm_cfg=norm_cfg,
act_cfg=act_cfg)
mask_convs.append(mask_conv)
self.mask_convs = nn.Sequential(*mask_convs)
def forward(self, features):
coord_features = list()
for feature in features:
x_range = torch.linspace(
-1, 1, feature.shape[-1], device=feature.device)
y_range = torch.linspace(
-1, 1, feature.shape[-2], device=feature.device)
y, x = torch.meshgrid(y_range, x_range)
y = y.expand([feature.shape[0], 1, -1, -1])
x = x.expand([feature.shape[0], 1, -1, -1])
coord = torch.cat([x, y], 1)
feature_with_coord = torch.cat([feature, coord], dim=1)
feature_with_coord = self.mask_convs(feature_with_coord)
feature_with_coord = feature_with_coord + feature
coord_features.append(feature_with_coord)
return coord_features
|