mingyuan/MotionDiffuse

1import math2import numpy as np3import matplotlib4import matplotlib.pyplot as plt5from mpl_toolkits.mplot3d import Axes3D6from matplotlib.animation import FuncAnimation, FFMpegFileWriter7from mpl_toolkits.mplot3d.art3d import Poly3DCollection8import mpl_toolkits.mplot3d.axes3d as p39# import cv210 11 12def list_cut_average(ll, intervals):13    if intervals == 1:14        return ll15 16    bins = math.ceil(len(ll) * 1.0 / intervals)17    ll_new = []18    for i in range(bins):19        l_low = intervals * i20        l_high = l_low + intervals21        l_high = l_high if l_high < len(ll) else len(ll)22        ll_new.append(np.mean(ll[l_low:l_high]))23    return ll_new24 25 26def plot_3d_motion(save_path, kinematic_tree, joints, title, figsize=(10, 10), fps=120, radius=4):27    matplotlib.use('Agg')28 29    title_sp = title.split(' ')30    if len(title_sp) > 20:31        title = '\n'.join([' '.join(title_sp[:10]), ' '.join(title_sp[10:20]), ' '.join(title_sp[20:])])32    elif len(title_sp) > 10:33        title = '\n'.join([' '.join(title_sp[:10]), ' '.join(title_sp[10:])])34 35    def init():36        ax.set_xlim3d([-radius / 4, radius / 4])37        ax.set_ylim3d([0, radius / 2])38        ax.set_zlim3d([0, radius / 2])39        # print(title)40        fig.suptitle(title, fontsize=20)41        ax.grid(b=False)42 43    def plot_xzPlane(minx, maxx, miny, minz, maxz):44        ## Plot a plane XZ45        verts = [46            [minx, miny, minz],47            [minx, miny, maxz],48            [maxx, miny, maxz],49            [maxx, miny, minz]50        ]51        xz_plane = Poly3DCollection([verts])52        xz_plane.set_facecolor((0.5, 0.5, 0.5, 0.5))53        ax.add_collection3d(xz_plane)54 55    #         return ax56 57    # (seq_len, joints_num, 3)58    data = joints.copy().reshape(len(joints), -1, 3)59    fig = plt.figure(figsize=figsize)60    ax = p3.Axes3D(fig)61    init()62    MINS = data.min(axis=0).min(axis=0)63    MAXS = data.max(axis=0).max(axis=0)64    colors = ['red', 'blue', 'black', 'red', 'blue',65              'darkblue', 'darkblue', 'darkblue', 'darkblue', 'darkblue',66              'darkred', 'darkred', 'darkred', 'darkred', 'darkred']67    frame_number = data.shape[0]68    #     print(data.shape)69 70    height_offset = MINS[1]71    data[:, :, 1] -= height_offset72    trajec = data[:, 0, [0, 2]]73 74    data[..., 0] -= data[:, 0:1, 0]75    data[..., 2] -= data[:, 0:1, 2]76 77    #     print(trajec.shape)78 79    def update(index):80        #         print(index)81        ax.lines = []82        ax.collections = []83        ax.view_init(elev=120, azim=-90)84        ax.dist = 7.585        #         ax =86        plot_xzPlane(MINS[0] - trajec[index, 0], MAXS[0] - trajec[index, 0], 0, MINS[2] - trajec[index, 1],87                     MAXS[2] - trajec[index, 1])88        #         ax.scatter(data[index, :22, 0], data[index, :22, 1], data[index, :22, 2], color='black', s=3)89 90        if index > 1:91            ax.plot3D(trajec[:index, 0] - trajec[index, 0], np.zeros_like(trajec[:index, 0]),92                      trajec[:index, 1] - trajec[index, 1], linewidth=1.0,93                      color='blue')94        #             ax = plot_xzPlane(ax, MINS[0], MAXS[0], 0, MINS[2], MAXS[2])95 96        for i, (chain, color) in enumerate(zip(kinematic_tree, colors)):97            #             print(color)98            if i < 5:99                linewidth = 4.0100            else:101                linewidth = 2.0102            ax.plot3D(data[index, chain, 0], data[index, chain, 1], data[index, chain, 2], linewidth=linewidth,103                      color=color)104        #         print(trajec[:index, 0].shape)105 106        plt.axis('off')107        ax.set_xticklabels([])108        ax.set_yticklabels([])109        ax.set_zticklabels([])110 111    ani = FuncAnimation(fig, update, frames=frame_number, interval=1000 / fps, repeat=False)112 113    # writer = FFMpegFileWriter(fps=fps)114    ani.save(save_path, fps=fps)115    plt.close()116