59 lines
2.0 KiB
Python
59 lines
2.0 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class AddCoordsTh(nn.Module):
|
|
def __init__(self, x_dim=64, y_dim=64, with_r=False):
|
|
super(AddCoordsTh, self).__init__()
|
|
self.x_dim = x_dim
|
|
self.y_dim = y_dim
|
|
self.with_r = with_r
|
|
|
|
xx_channel, yy_channel = self._prepare_coords()
|
|
self.xx_channel = nn.parameter.Parameter(xx_channel, requires_grad=False)
|
|
self.yy_channel = nn.parameter.Parameter(yy_channel, requires_grad=False)
|
|
|
|
def _prepare_coords(self):
|
|
xx_ones = torch.ones([1, self.y_dim], dtype=torch.int32)
|
|
xx_ones = xx_ones.unsqueeze(-1)
|
|
|
|
xx_range = torch.arange(self.x_dim, dtype=torch.int32).unsqueeze(0)
|
|
xx_range = xx_range.unsqueeze(1)
|
|
|
|
xx_channel = torch.matmul(xx_ones, xx_range)
|
|
xx_channel = xx_channel.unsqueeze(-1)
|
|
|
|
yy_ones = torch.ones([1, self.x_dim], dtype=torch.int32)
|
|
yy_ones = yy_ones.unsqueeze(1)
|
|
|
|
yy_range = torch.arange(self.y_dim, dtype=torch.int32).unsqueeze(0)
|
|
yy_range = yy_range.unsqueeze(-1)
|
|
|
|
yy_channel = torch.matmul(yy_range, yy_ones)
|
|
yy_channel = yy_channel.unsqueeze(-1)
|
|
|
|
xx_channel = xx_channel.permute(0, 3, 2, 1)
|
|
yy_channel = yy_channel.permute(0, 3, 2, 1)
|
|
|
|
xx_channel = xx_channel.float() / (self.x_dim - 1)
|
|
yy_channel = yy_channel.float() / (self.y_dim - 1)
|
|
|
|
xx_channel = xx_channel * 2 - 1
|
|
yy_channel = yy_channel * 2 - 1
|
|
return xx_channel, yy_channel
|
|
|
|
def forward(self, input_tensor):
|
|
"""
|
|
input_tensor: (batch, c, x_dim, y_dim)
|
|
"""
|
|
batch_size_tensor = input_tensor.shape[0]
|
|
xx_channel = self.xx_channel.repeat(batch_size_tensor, 1, 1, 1)
|
|
yy_channel = self.yy_channel.repeat(batch_size_tensor, 1, 1, 1)
|
|
ret = torch.cat([input_tensor, xx_channel, yy_channel], dim=1)
|
|
|
|
if self.with_r:
|
|
rr = torch.sqrt(torch.pow(xx_channel - 0.5, 2) + torch.pow(yy_channel - 0.5, 2))
|
|
ret = torch.cat([ret, rr], dim=1)
|
|
|
|
return ret
|