def decode_node(self, node_latent, max_depth, full_label, is_leaf=False):
if node_latent.shape[0] != 1:
raise ValueError('Node decoding does not support batch_size > 1.')
is_leaf_logit = self.leaf_classifier(node_latent)
node_is_leaf = is_leaf_logit.item() > 0
# use maximum depth to avoid potential infinite recursion
if max_depth < 1:
is_leaf = True
# decode the current part box
box = self.box_decoder(node_latent)
if node_is_leaf or is_leaf:
ret = Tree.Node(is_leaf=True, full_label=full_label, label=full_label.split('/')[-1])
ret.set_from_box_quat(box.view(-1))
return ret
else:
child_feats, child_sem_logits, child_exists_logit, edge_exists_logits = \
self.child_decoder(node_latent)
child_sem_logits = child_sem_logits.cpu().numpy().squeeze()
# children
child_nodes = []
child_idx = {}
for ci in range(child_feats.shape[1]):
if torch.sigmoid(child_exists_logit[:, ci, :]).item() > 0.5:
idx = np.argmax(child_sem_logits[ci, Tree.part_name2cids[full_label]])
idx = Tree.part_name2cids[full_label][idx]
child_full_label = Tree.part_id2name[idx]
child_nodes.append(self.decode_node(\
child_feats[:, ci, :], max_depth-1, child_full_label, \
is_leaf=(child_full_label not in Tree.part_non_leaf_sem_names)))
child_idx[ci] = len(child_nodes) - 1
# edges
child_edges = []
nz_inds = torch.nonzero(torch.sigmoid(edge_exists_logits) > 0.5)
edge_from = nz_inds[:, 1]
edge_to = nz_inds[:, 2]
edge_type = nz_inds[:, 3]
for i in range(edge_from.numel()):
cur_edge_from = edge_from[i].item()
cur_edge_to = edge_to[i].item()
cur_edge_type = edge_type[i].item()
if cur_edge_from in child_idx and cur_edge_to in child_idx:
child_edges.append({
'part_a': child_idx[cur_edge_from],
'part_b': child_idx[cur_edge_to],
'type': self.conf.edge_types[cur_edge_type]})
node = Tree.Node(is_leaf=False, children=child_nodes, edges=child_edges, \
full_label=full_label, label=full_label.split('/')[-1])
node.set_from_box_quat(box.view(-1))
return node
def decode_node(self, node_latent, max_depth, full_label, is_leaf=False):
if node_latent.shape[0] != 1:
raise ValueError('Node decoding does not support batch_size > 1.')
is_leaf_logit = self.leaf_classifier(node_latent)
node_is_leaf = is_leaf_logit.item() > 0
# use maximum depth to avoid potential infinite recursion
if max_depth < 1:
is_leaf = True
# decode the current part box
box = self.box_decoder(node_latent)
if node_is_leaf or is_leaf:
ret = Tree.Node(is_leaf=True, \
full_label=full_label, label=full_label.split('/')[-1])
ret.set_from_box_quat(box.view(-1))
return ret
else:
child_feats, child_sem_logits, child_exists_logit = \
self.child_decoder(node_latent)
child_sem_logits = child_sem_logits.cpu().numpy().squeeze()
# children
child_nodes = []
for ci in range(child_feats.shape[1]):
if child_exists_logit[:, ci, :].item() > 0:
if full_label not in Tree.part_non_leaf_sem_names:
print(
'WARNING: predicting nonzero children for a node with leaf semantics, ignoring the children'
)
continue
idx = np.argmax(
child_sem_logits[ci, Tree.part_name2cids[full_label]])
idx = Tree.part_name2cids[full_label][idx]
child_full_label = Tree.part_id2name[idx]
child_nodes.append(self.decode_node(\
child_feats[:, ci, :], max_depth-1, child_full_label, \
is_leaf=(child_full_label not in Tree.part_non_leaf_sem_names)))
ret = Tree.Node(is_leaf=len(child_nodes) == 0, children=child_nodes, \
full_label=full_label, label=full_label.split('/')[-1])
ret.set_from_box_quat(box.view(-1))
return ret
