def plot_evaluation()

in src/eval.py [0:0]

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

def plot_evaluation(X, y, out, data, config: Train_Config, untracked_only=True, parent_indices=None, vispy=None, mask=None, **kwargs):
    p_X = config.in_features.features_to_pose().solve_batch(X, ref_positions=None, mask=None, body=X, **kwargs)
    p_out = config.out_features.features_to_pose().solve_batch(out, ref_positions=None, mask=None, body=X, **kwargs)
    p_y = config.out_features.features_to_pose().solve_batch(y, ref_positions=None, mask=None, body=X, **kwargs)

    plot_smoothed = False
    if plot_smoothed:
        p_X = pos_smooth_past_momentum(p_out, smooth_frames=5, max_speed=5)
        p_y = pos_smooth_past_momentum(p_out, smooth_frames=3, max_speed=5)
        p_out = pos_smooth_past_momentum(p_out, smooth_frames=2, max_speed=5)

    X = X.cpu().detach().numpy()
    p_X = p_X.detach().cpu().numpy()
    p_out = p_out.detach().cpu().numpy()
    p_y = p_y.detach().cpu().numpy()

    p_out = p_out.reshape((len(X), -1, 3))
    p_y = p_y.reshape((len(X), -1, 3))
    p_X = p_X.reshape((len(X), -1, 3))

    p_out[:, :, [0, 2]] -= p_out[:, None, config.skeleton.Idx.root, [0, 2]]
    p_y[:, :, [0, 2]] -= p_y[:, None, config.skeleton.Idx.root, [0, 2]]
    p_X[:, :, [0, 2]] -= p_X[:, None, config.skeleton.Idx.root, [0, 2]]

    p_out[:, :, 0] -= 1
    p_X[:, :, 0] += 1

    if vispy is None:
        vispy = Vispy3DScatter()

    X_col = np.array([.2, .2, .2, 1])

    out_colors = np.zeros((X.shape[0], len(config.skeleton.Idx.all), 4), dtype=np.float32)
    out_colors[:, :] = np.array([0, 0, 1, 1])
    Y_colors = np.zeros(out_colors.shape, dtype=np.float32)
    Y_colors[:, :] = np.array([0, 1, 0, 1])
    if config.eval_occluded_joints:
        Y_colors[X[:, config.in_features.features_wo_occlusions():] > 0.00001] = np.array([1, 0, 0, 1])

    out_colors = Y_colors

    X_colors = np.zeros(out_colors.shape, dtype=np.float32)
    X_colors[:, :] = X_col
    if config.eval_occluded_joints:
        X_colors[X[:, config.in_features.features_wo_occlusions():] > 0.00001] = np.array([1, 0, 0, 1])

    points = np.concatenate((p_out, p_y, p_X), axis=1)
    colors = np.concatenate((out_colors, Y_colors, X_colors), axis=1)

    bones = None
    if parent_indices is not None:
        if not isinstance(parent_indices, dict):
            new_par_idx = {}
            for i, par in enumerate(parent_indices):
                new_par_idx[i] = par
            parent_indices = new_par_idx

        bones = np.zeros((y.shape[0], p_y.shape[1] * 2, 3))
        keys = list(parent_indices.keys())
        for self_idx, par_idx in parent_indices.items():
            if par_idx < 0:
                continue
            bones[:, keys.index(self_idx) * 2] = p_y[:, keys.index(par_idx), [0, 2, 1]]
            bones[:, keys.index(self_idx) * 2 + 1] = p_y[:, keys.index(self_idx), [0, 2, 1]]

    if untracked_only:
        untracked_indices = np.max(X[:, config.in_features.features_wo_occlusions():], axis=1) > 0.00001
        points = points[untracked_indices]
        if bones is not None:
            bones = bones[untracked_indices]
        colors = colors[untracked_indices]

    occ_mask = get_joint_mask(X, config)

    vispy.plot_skeleton_with_bones(
        points[:, :, [0, 2, 1]],
        config.skeleton,
        p_colors=colors,
        l_colors=np.array([[0., 0., 1., 1.], [0., 1., 0., 1.], list(X_col)]),
        speed=2,
        fps=data.dataset_fps * 2 // config.sequence_distance,
        n_skeletons=3,
        red_bones_for_skeleton=[0, 1, 2] if config.eval_occluded_joints else None,
        occlusions=occ_mask,
        center=np.array([0, 0, 1])
    )

    return vispy