def heatmaps_to_coords(xy_hm, zy_hm, xz_hm):
xy = dsnt(xy_hm)
zy = dsnt(zy_hm)
xz = dsnt(xz_hm)
x, y = xy.split(1, -1)
z = 0.5 * (zy[:, :, 0:1] + xz[:, :, 1:2])
return torch.cat([x, y, z], -1)
def forward(self, *inputs):
t = inputs[0]
t = self.in_cnn(t)
self.xy_heatmaps = [flat_softmax(self.xy_hm_cnn(t))]
self.zy_heatmaps = [flat_softmax(self.zy_hm_cnn(t))]
self.xz_heatmaps = [flat_softmax(self.xz_hm_cnn(t))]
xy = dsnt(self.xy_heatmaps[-1])
zy = dsnt(self.zy_heatmaps[-1])
xz = dsnt(self.xz_heatmaps[-1])
x = xy.narrow(-1, 0, 1)
y = xy.narrow(-1, 1, 1)
z = 0.5 * (zy.narrow(-1, 0, 1) + xz.narrow(-1, 1, 1))
return torch.cat([x, y, z], -1)