def get_node_feats()

in anticipation/anticipation/datasets/epic_future_labels.py [0:0]

• 33 lines of code
• 13 McCabe index (conditional complexity)

    def get_node_feats(self, graph, trunk_features):
        node_feats = []
        node_length = []

        for node in sorted(graph.nodes()):
            visits = graph.node[node]['members']

            feats_idx = []
            for visit in visits:
                v_feat_idx = self.get_visit_feature(trunk_features, visit['start'][1], visit['stop'][1], self.fb_dim)
                feats_idx.extend(v_feat_idx)
            #         feats.append(v_feat)
            # feats = feats or [torch.zeros(self.fb_dim)]

            # # Pick last self.node_num_member visits
            # feats = feats[-self.node_num_member:]
            # length = len(feats)
            # feats = feats + [torch.zeros(self.fb_dim)]*(self.node_num_member-length)

            # uniformly pick self.node_num_member visits
            if self.node_num_member > 0:
                if self.rand_visit and len(feats_idx) > self.node_num_member:
                    idxes = sorted(np.random.choice(feats_idx, self.node_num_member, replace=False))
                    feats = [to_tensor(trunk_features[idx]) for idx in idxes]
                else:
                    if len(feats_idx) > 0:
                        idxes_ = np.round(np.linspace(0, len(feats_idx) - 1, self.node_num_member)).astype(int)
                        # print(idxes_, len(feats_idx))
                        idxes = [feats_idx[i] for i in idxes_]
                        feats = [to_tensor(trunk_features[idx]) for idx in idxes]
                    else:
                        feats = [torch.zeros(self.fb_dim)] * self.node_num_member
                
            else:
                if len(feast_idx) > 0:
                    feats = [to_tensor(trunk_features[idx]) for idx in feats_idx]
                    feats = [torch.stack(feats, 0).mean(0)]
                else:
                    feats = [torch.zeros(self.fb_dim)]
            
            length = len(feats)

            feats = torch.stack(feats, 0) # (K, 2048)
            node_feats.append(feats)
            node_length.append(length)
            
        node_feats = torch.stack(node_feats, 0)  # (N, K, 2048)
        node_length = torch.LongTensor(node_length)

        return node_feats, node_length