def plot_skeleton_with_bones()

in src/plots.py [0:0]

• 35 lines of code
• 14 McCabe index (conditional complexity)

    def plot_skeleton_with_bones(self, p, skeleton, scale=1, p_colors="blue", l_colors="blue", center=np.zeros(3),
                                 speed=1, debug_markers=[], camera_rotation_speed=0, n_skeletons=1,
                                 red_bones_for_skeleton=None, occlusions=None, fps=None):
        v = np.empty((p.shape[0], p.shape[1] * 2, p.shape[2]))

        n_joints = len(skeleton.Idx.all)
        for skeleton_idx in range(n_skeletons):
            off = skeleton_idx * n_joints
            for joint_idx, parent_idx in enumerate(skeleton.parent_idx_vector()):
                i0 = off + (parent_idx if parent_idx >= 0 else joint_idx)
                i1 = off + joint_idx
                v[:, i1 * 2] = p[:, i0]
                v[:, i1 * 2 + 1] = p[:, i1]

        if type(l_colors) is not type("") and len(l_colors) > 1 and red_bones_for_skeleton is None:
            color_arr = []
            for color in l_colors:
                for i in range(n_joints * 2):
                    color_arr.append(color)
            l_colors = np.array(color_arr)

        if fps is None:
            fps = speed * 60

        if red_bones_for_skeleton is not None and type(l_colors) == np.ndarray:
            color_arr = np.empty((v.shape[0], v.shape[1], 4))
            red_color = np.array([1., 0., 0., 1.])[None, None, :]
            red_color = np.repeat(np.repeat(red_color, v.shape[0], axis=0), v.shape[1] // n_skeletons, axis=1)
            for skeleton_idx in range(n_skeletons):
                color_arr[:, (n_joints * 2) * skeleton_idx: (n_joints * 2) * (skeleton_idx + 1)] = l_colors[
                    skeleton_idx]
                if skeleton_idx in red_bones_for_skeleton:
                    c = color_arr[:, (n_joints * 2) * skeleton_idx: (n_joints * 2) * (skeleton_idx + 1)]
                    occ_mask = occlusions > 0.00001
                    c[:, ::2][occ_mask] = red_color[:, ::2][occ_mask]
                    c[:, 1::2][occ_mask] = red_color[:, 1::2][occ_mask]
                    color_arr[:, (n_joints * 2) * skeleton_idx: (n_joints * 2) * (skeleton_idx + 1)] = c

            l_colors = color_arr

        self.plot(p, v, scale, p_colors, l_colors, center, speed, debug_markers, camera_rotation_speed, fps=fps)