GenerativePointClouds/test.py at main · MatuTia/GenerativePointClouds · 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
import torch
from matplotlib import pyplot as plt

from generator import StyleTreeGenerator


class PlotUtils:

    def __init__(self, title: str):
        self.ax = plt.figure().add_subplot(projection='3d')
        self.ax.view_init(elev=0, azim=0, roll=0)
        self.ax.set_title(title)
        plt.axis('off')

    def scatter_plot(self, x, y, z):
        self.ax.scatter3D(x, y, z, s=6)

    def perimeter_plot(self, x, y, z):
        self.ax.plot3D(x, y, z)

    def plane_plot(self, x, y, z):
        self.ax.plot_surface(x, y, z)

    @staticmethod
    def show_plot():
        plt.show()


if __name__ == '__main__':

    # Default Setting
    torch.manual_seed(42)
    torch.set_default_dtype(torch.float32)

    # Training-Name
    name = 'Surface-Dynamic-3-StyleGAN-64'
    epoch = 500
    dir_name = "generator"

    device = 'cpu'

    model = torch.load(f'model/{name}/{dir_name}/generator-{epoch}.pt')

    ada_in_after = False
    mapping_branching = False
    truncate_style = False

    # gen = TreeGenerator(degree == 32).to(device)
    gen = StyleTreeGenerator(ada_in_after, mapping_branching, truncate_style, False, device).to(device)
    gen.load_state_dict(model)

    noise = torch.randn((10, 1, 96), device=device)
    style = torch.randn((10, 1, 96), device=device)

    clouds = gen.forward(style, [noise]).cpu().detach().numpy()

    for cloud in clouds:
        plot = PlotUtils("")
        plot.scatter_plot(cloud[:, 0], cloud[:, 2], cloud[:, 1])
        plot.show_plot()