mgn-tf2/plot_cloth.py at main · DimensionLab/mgn-tf2 · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
import os
import pickle
import argparse
from pathlib import Path


import numpy as np
from matplotlib import animation
import matplotlib.pyplot as plt


def plot_cloth(data, filename=None):
    figure = plt.figure()
    axis = figure.add_subplot(121, projection="3d")
    axis2 = figure.add_subplot(122, projection="3d")
    figure.set_figheight(6)
    figure.set_figwidth(12)

    skip = 10
    num_steps = data["pred_world_pos"].shape[0]
    num_frames = num_steps // skip
    upper_bound = np.amax(data["pred_world_pos"], axis=(0,1))
    lower_bound= np.amin(data["pred_world_pos"], axis=(0,1))
    def animate(frame):
        step = (frame*skip) % num_steps
        axis.cla()
        axis.set_xlim([lower_bound[0], upper_bound[0]])
        axis.set_ylim([lower_bound[1], upper_bound[1]])
        axis.set_zlim([lower_bound[2], upper_bound[2]])
        axis.autoscale(False)
        positions = data['pred_world_pos'][step]
        faces = data["cells"][step]
        axis.plot_trisurf(positions[:,0], positions[:, 1], faces, positions[:, 2])
        axis.set_title('Predicted')

        axis2.cla()
        axis2.set_xlim([lower_bound[0], upper_bound[0]])
        axis2.set_ylim([lower_bound[1], upper_bound[1]])
        axis2.set_zlim([lower_bound[2], upper_bound[2]])
        axis2.autoscale(False)
        positions = data['true_world_pos'][step]
        faces = data["cells"][step]
        axis2.plot_trisurf(positions[:,0], positions[:, 1], faces, positions[:, 2])
        axis2.set_title('Ground Truth')

        figure.suptitle(f"Time step: {step}")

        return figure,

    _ = animation.FuncAnimation(figure, animate, frames=num_frames, interval=100)

    # ani.save(filename)
    plt.show(block=True)


def generate_all():
    results_path = os.path.join(os.path.dirname(__file__), 'results', 'flag-simple_weights-step2300000-loss0.04908.hdf5')
    output_path = os.path.join(os.path.dirname(__file__), 'animations', 'flag-simple_weights-step2300000-loss0.04908.hdf5')

    Path(output_path).mkdir(parents=True, exist_ok=True)
    for i in range(100):
        with open(os.path.join(results_path, f'{i:03d}.eval'), 'rb') as f:
            data = pickle.load(f)
        plot_cloth(data, os.path.join(output_path, f'{i:03d}.mp4'))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--datafile")
    parser.add_argument("-o", "--output", help="output file")
    args = parser.parse_args()
    if (args.output == None):
        args.output = "out.mp4"

    with open(args.datafile, "rb") as f:
        data = pickle.load(f)

    plot_cloth(data, args.output)


if __name__ == "__main__":
    # generate_all()
    main()