Kleinhe/SemanticBoost

1import math2import numpy as np3import matplotlib4import matplotlib.pyplot as plt5from mpl_toolkits.mplot3d.art3d import Poly3DCollection6import mpl_toolkits.mplot3d.axes3d as p37from textwrap import wrap8from tqdm import tqdm9 10def list_cut_average(ll, intervals):11    if intervals == 1:12        return ll13 14    bins = math.ceil(len(ll) * 1.0 / intervals)15    ll_new = []16    for i in range(bins):17        l_low = intervals * i18        l_high = l_low + intervals19        l_high = l_high if l_high < len(ll) else len(ll)20        ll_new.append(np.mean(ll[l_low:l_high]))21    return ll_new22 23 24def plot_3d_motion(kinematic_tree, joints, title, dataset="humanml", figsize=(10.24, 10.24), radius=3,25                   vis_mode='default', gt_frames=[]):26    matplotlib.use('Agg')27    title = '\n'.join(wrap(title, 40))28 29    def init():30        ax.set_xlim3d([-radius / 2, radius / 2])31        ax.set_ylim3d([0, radius])32        ax.set_zlim3d([-radius / 3., radius * 2 / 3.])33        # print(title)34        fig.suptitle(title, fontsize=20)35        ax.grid(b=False)36 37    def plot_xzPlane(minx, maxx, miny, minz, maxz):38        ## Plot a plane XZ39        verts = [40            [minx, miny, minz],41            [minx, miny, maxz],42            [maxx, miny, maxz],43            [maxx, miny, minz]44        ]45        xz_plane = Poly3DCollection([verts])46        xz_plane.set_facecolor((0.5, 0.5, 0.5, 0.5))47        ax.add_collection3d(xz_plane)48 49    #         return ax50 51    # (seq_len, joints_num, 3)52    data = joints.copy().reshape(len(joints), -1, 3)53 54    # preparation related to specific datasets55    if dataset == 'kit':56        data *= 0.003  # scale for visualization57    elif dataset == 'humanml':58        data *= 1.3  # scale for visualization59    elif dataset in ['humanact12', 'uestc']:60        data *= -1.5 # reverse axes, scale for visualization61 62    fig = plt.figure(figsize=figsize)63    plt.tight_layout()64    ax = p3.Axes3D(fig)65    init()66    MINS = data.min(axis=0).min(axis=0)67    MAXS = data.max(axis=0).max(axis=0)68    colors_blue = ["#4D84AA", "#5B9965", "#61CEB9", "#34C1E2", "#80B79A"]  # GT color69    colors_orange = ["#DD5A37", "#D69E00", "#B75A39", "#FF6D00", "#DDB50E"]  # Generation color70    colors = colors_orange71    if vis_mode == 'upper_body':  # lower body taken fixed to input motion72        colors[0] = colors_blue[0]73        colors[1] = colors_blue[1]74    elif vis_mode == 'gt':75        colors = colors_blue76 77    frame_number = data.shape[0]78    #     print(dataset.shape)79 80    height_offset = MINS[1]81    data[:, :, 1] -= height_offset82    trajec = data[:, 0, [0, 2]]83 84    data[..., 0] -= data[:, 0:1, 0]85    data[..., 2] -= data[:, 0:1, 2]86 87    #     print(trajec.shape)88 89    def update(index):90        #         print(index)91        ax.lines = []92        ax.collections = []93        ax.view_init(elev=120, azim=-90)94        ax.dist = 7.595        #         ax =96        plot_xzPlane(MINS[0] - trajec[index, 0], MAXS[0] - trajec[index, 0], 0, MINS[2] - trajec[index, 1],97                     MAXS[2] - trajec[index, 1])98 99        used_colors = colors_blue if index in gt_frames else colors100        for i, (chain, color) in enumerate(zip(kinematic_tree, used_colors)):101            if i < 5:102                linewidth = 4.0103            else:104                linewidth = 2.0105            ax.plot3D(data[index, chain, 0], data[index, chain, 1], data[index, chain, 2], linewidth=linewidth,106                      color=color)107        #         print(trajec[:index, 0].shape)108 109        plt.axis('off')110        ax.set_xticklabels([])111        ax.set_yticklabels([])112        ax.set_zticklabels([])113    114    for i in tqdm(range(frame_number)):115        update(i)116        plt.savefig("temp/%06d.png"%(i))117 118    plt.close()