-
Notifications
You must be signed in to change notification settings - Fork 1
/
show_pt_yw.py
executable file
·140 lines (94 loc) · 3.33 KB
/
show_pt_yw.py
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
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import sys
from PIL import Image
import os
import os.path
import errno
import torch
import argparse
import json
import codecs
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
from torch.autograd import Variable
from datasets import PartDataset
import torch.nn.functional as F
from pointnet import FoldingNet,ChamferLoss
if __name__=='__main__':
np.random.seed(100)
pt = np.random.rand(250,3)
# fig = plt.figure()
# ax = fig.add_subplot(111,projection='3d')
#ax.scatter(pt[:,0],pt[:,1],pt[:,2])
#plt.show()
class_choice = 'Lamp'
pt_root = 'shapenetcore_partanno_segmentation_benchmark_v0'
npoints = 2500
shapenet_dataset = PartDataset(root = pt_root,class_choice = class_choice, classification = True,train = False)
print('len(shapenet_dataset) :',len(shapenet_dataset))
dataloader = torch.utils.data.DataLoader(shapenet_dataset,batch_size=1,shuffle=False)
li = list(enumerate(dataloader))
print(len(li))
# ps,cls = shapenet_dataset[0]
# print('ps.size:',ps.size())
# print('ps.type:',ps.type())
# print('cls.size',cls.size())
# print('cls.type',cls.type())
# ps2,cls2 = shapenet_dataset[1]
# ax.scatter(ps[:,0],ps[:,1],ps[:,2])
# ax.set_xlabel('X label')
# ax.set_ylabel('Y label')
# ax.set_zlabel('Z label')
# # fig2 = plt.figure()
# # a2 = fig2.add_subplot(111,projection='3d')
# # a2.scatter(ps2[:,0],ps2[:,1],ps2[:,2])
# plt.show()
foldingnet = FoldingNet()
foldingnet.load_state_dict(torch.load('cls/foldingnet_model_50.pth'))
foldingnet.cuda()
chamferloss = ChamferLoss()
chamferloss = chamferloss.cuda()
#print(foldingnet)
foldingnet.eval()
i, data = li[4]
points, target = data
points = points.transpose(2,1)
points = points.cuda()
recon_pc, mid_pc = foldingnet(points)
points_show = points.cpu().detach().numpy()
re_show = recon_pc.cpu().detach().numpy()
fig_ori = plt.figure()
a1 = fig_ori.add_subplot(111,projection='3d')
a1.scatter(points_show[0,0,:],points_show[0,1,:],points_show[0,2,:])
#plt.savefig('points_show.png')
fig_re = plt.figure()
a2 = fig_re.add_subplot(111,projection='3d')
a2.scatter(re_show[0,0,:],re_show[0,1,:],re_show[0,2,:])
#plt.savefig('re_show.png')
plt.show()
print(points.size())
print(recon_pc.size())
loss = chamferloss(points.transpose(2,1),recon_pc.transpose(2,1))
print('loss',loss.item())
for i,data in enumerate(dataloader):
points, target = data
points = points.transpose(2,1)
points = points.cuda()
recon_pc, _ = foldingnet(points)
points_show = points.cpu().detach().numpy()
#print(points_show.shape)
points_show = points_show.transpose(0,2,1)
re_show = recon_pc.cpu().detach().numpy()
re_show = re_show.transpose(0,2,1)
#batch = points.size(0)
np.savetxt('recon_pc/ori_%s_%d.pts'%(class_choice,i),points_show[0])
np.savetxt('recon_pc/rec_%s_%d.pts'%(class_choice,i),re_show[0])