mirror of https://github.com/open-mmlab/mmocr.git
57 lines
2.0 KiB
Python
57 lines
2.0 KiB
Python
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
|