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)
def forward(self, *inputs):
features = self.in_cnn(inputs[0])
# These lists will store the outputs from each stage
xy_heatmaps = []
zy_heatmaps = []
xz_heatmaps = []
inp = features
for t in range(self.n_stages):
if t > 0:
combined_hm_features = self.hm_combiners[t - 1](
xy_heatmaps[t - 1],
zy_heatmaps[t - 1],
xz_heatmaps[t - 1],
)
inp = inp + combined_hm_features
xy_heatmaps.append(flat_softmax(self.xy_hm_cnns[t](inp)))
zy_heatmaps.append(flat_softmax(self.zy_hm_cnns[t](inp)))
xz_heatmaps.append(flat_softmax(self.xz_hm_cnns[t](inp)))
return xy_heatmaps, zy_heatmaps, xz_heatmaps