def main()

in fairmotion/viz/body_visualizer.py [0:0]

• 58 lines of code
• 5 McCabe index (conditional complexity)

def main(args):
    comp_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    bm = BodyModel(
        model_type="smplh", bm_fname=args.body_model_file, num_betas=10
    ).to(comp_device)
    faces = c2c(bm.f)

    img_shape = (1600, 1600)
    motion = bvh.load(
        file=args.input_file,
        scale=0.5,
        v_up_skel=np.array([0.0, 1.0, 0.0]),
        v_face_skel=np.array([0.0, 0.0, 1.0]),
        v_up_env=np.array([0.0, 0.0, 1.0]),
    )
    motion = motion_ops.rotate(
        motion, conversions.Ax2R(conversions.deg2rad(-90)),
    )
    mv = prepare_mesh_viewer(img_shape)

    out = cv2.VideoWriter(
        args.video_output_file, cv2.VideoWriter_fourcc(*"XVID"), 30, img_shape
    )

    parents = bm.kintree_table[0].long()[: 21 + 1]
    parents = parents.cpu().numpy()
    dfs_order = get_dfs_order(parents)
    for frame in tqdm.tqdm(range(motion.num_frames())):
        pose = motion.get_pose_by_frame(frame)
        R, p = conversions.T2Rp(pose.data[0])
        root_orient = conversions.R2A(R)
        trans = p

        num_joints = len(pose.data) - 1
        body_model_pose_data = np.zeros(num_joints * 3)
        for motion_joint, amass_joint in enumerate(dfs_order):
            # motion_joint is idx of joint in Motion class order
            # amass_joint is idx of joint in AMASS skeleton
            if amass_joint == 0:
                continue
            pose_idx = amass_joint - 1
            # Convert rotation matrix to axis angle
            axis_angles = conversions.R2A(
                conversions.T2R(pose.data[motion_joint])
            )
            body_model_pose_data[pose_idx * 3 : pose_idx * 3 + 3] = axis_angles

        pose_data_t = (
            torch.Tensor(body_model_pose_data).to(comp_device).unsqueeze(0)
        )
        root_orient_t = torch.Tensor(root_orient).to(comp_device).unsqueeze(0)
        trans_t = torch.Tensor(trans).to(comp_device).unsqueeze(0)
        body = bm(
            pose_body=pose_data_t, root_orient=root_orient_t, trans=trans_t
        )

        body_mesh = trimesh.Trimesh(
            vertices=c2c(body.v[0]),
            faces=faces,
            vertex_colors=np.tile(colors["grey"], (6890, 1)),
        )
        # TODO: Add floor trimesh to the scene to display the ground plane
        mv.set_static_meshes([body_mesh])
        body_image = mv.render()
        img = body_image.astype(np.uint8)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        out.write(img)
    out.release()