class AddCoords(nn.Module):
def __init__(self, with_r=False):
super().__init__()
self.with_r = with_r
def forward(self, input_tensor):
"""
Args:
input_tensor: shape(batch, channel, x_dim, y_dim)
"""
batch_size, _, x_dim, y_dim = input_tensor.size()
xx_channel = torch.arange(x_dim).repeat(1, y_dim, 1).transpose(1, 2)
yy_channel = torch.arange(y_dim).repeat(1, x_dim, 1)
print(xx_channel)
xx_channel = xx_channel.repeat(batch_size, 1, 1, 1)
yy_channel = yy_channel.repeat(batch_size, 1, 1, 1)
xx_channel_01 = xx_channel.float() / (x_dim - 1)
yy_channel_01 = yy_channel.float() / (y_dim - 1)
xx_channel = xx_channel_01 * 2 - 1
yy_channel = yy_channel_01 * 2 - 1
ret = torch.cat([
input_tensor,
xx_channel.type_as(input_tensor),
yy_channel.type_as(input_tensor)], dim=1)
if self.with_r:
rr = torch.sqrt(torch.pow(xx_channel_01.type_as(input_tensor) - 0.5, 2) + torch.pow(yy_channel_01.type_as(input_tensor) - 0.5, 2))
ret = torch.cat([ret, rr], dim=1)
return ret