def __init__(self,
height=64,
width=64,
with_r=False,
with_boundary=False):
super(AddCoordsTh, self).__init__()
self.with_r = with_r
self.with_boundary = with_boundary
device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
with torch.no_grad():
x_coords = torch.arange(height).unsqueeze(1).expand(
height, width).float()
y_coords = torch.arange(width).unsqueeze(0).expand(
height, width).float()
x_coords = (x_coords / (height - 1)) * 2 - 1
y_coords = (y_coords / (width - 1)) * 2 - 1
coords = torch.stack([x_coords, y_coords],
dim=0) # (2, height, width)
if self.with_r:
rr = torch.sqrt(
torch.pow(x_coords, 2) +
torch.pow(y_coords, 2)) # (height, width)
rr = (rr / torch.max(rr)).unsqueeze(0)
coords = torch.cat([coords, rr], dim=0)
self.coords = coords.unsqueeze(0).to(
device) # (1, 2 or 3, height, width)
self.x_coords = x_coords.to(device)
self.y_coords = y_coords.to(device)
def normalize(x, eps=1e-6):
"""Apply min-max normalization."""
x = x.contiguous()
N, C, H, W = x.size()
x_ = x.view(N * C, -1)
max_val = porch.max(x_, dim=1, keepdim=True)[0]
min_val = porch.min(x_, dim=1, keepdim=True)[0]
x_ = (x_ - min_val) / (max_val - min_val + eps)
out = x_.view(N, C, H, W)
return out
def normalize(x, eps=1e-6):
"""Apply min-max normalization."""
# x = x.contiguous()
x = torch.varbase_to_tensor(x)
N, C, H, W = x.shape
x_ = x.view(N * C, -1)
max_val = torch.max(x_, dim=1, keepdim=True)[0]
min_val = torch.min(x_, dim=1, keepdim=True)[0]
x_ = (x_ - min_val) / (max_val - min_val + eps)
x_ = torch.varbase_to_tensor(x_)
out = x_.view(N, C, H, W)
return out
def get_preds_fromhm(hm):
max, idx = torch.max(
hm.view(hm.size(0), hm.size(1), hm.size(2) * hm.size(3)), 2)
idx += 1
preds = idx.view(idx.size(0), idx.size(1), 1).repeat(1, 1, 2).float()
preds[..., 0].apply_(lambda x: (x - 1) % hm.size(3) + 1)
preds[..., 1].add_(-1).div_(hm.size(2)).floor_().add_(1)
for i in range(preds.size(0)):
for j in range(preds.size(1)):
hm_ = hm[i, j, :]
pX, pY = int(preds[i, j, 0]) - 1, int(preds[i, j, 1]) - 1
if pX > 0 and pX < 63 and pY > 0 and pY < 63:
diff = torch.FloatTensor(
[hm_[pY, pX + 1] - hm_[pY, pX - 1],
hm_[pY + 1, pX] - hm_[pY - 1, pX]])
preds[i, j].add_(diff.sign_().mul_(.25))
preds.add_(-0.5)
return preds
def max(self, dim=None, keepdim=False):
return torch.max(self, dim=dim, keepdim=keepdim)