build_feature_files.py 文件源码

def action_label_counts(directory, data_loader, n_actions=18, n=None):
    episode_paths = frame.episode_paths(directory)
    label_counts = [0, 0]
    action_label_counts = [[0, 0] for i in range(n_actions)]
    if n is not None:
        np.random.shuffle(episode_paths)
        episode_paths = episode_paths[:n]
    for episode_path in tqdm.tqdm(episode_paths):
        try:
            features, labels = data_loader.load_features_and_labels([episode_path])
        except:
            traceback.print_exc()
        else:
            for label in range(len(label_counts)):
                label_counts[label] += np.count_nonzero(labels == label)
                for action in range(n_actions):
                    actions = np.reshape(np.array(features["action"]), [-1])
                    action_label_counts[action][label] += np.count_nonzero(
                        np.logical_and(labels == label, actions == action))
    return label_counts, action_label_counts