def show_example(path, results, output, num_models):
classes = None
tf = tfs.Compose([
tfs.UnitSpherePointCloud(),
tfs.RandomRotatePointCloud(type='upright')
])
ds = ModelNetPointCloud(basedir=path,
split='test',
categories=classes,
device='cuda',
transform=tf,
num_points=2**10,
sample_points=2**12)
# Get data
loader = DataLoader(ds, batch_size=num_models, shuffle=True)
for item in loader:
break
# Set up figure
rows = 3
fig = plt.figure(figsize=(4 * rows, 4 * num_models), dpi=200)
# Draw inputs
for i, pts in enumerate(item.clone().cpu()):
ax = fig.add_subplot(num_models, rows, i * rows + 1, projection='3d')
if i == 0:
ax.set_title('Input')
ax_points(ax, pts, s=3)
# Predict & Draw
measure, net = result
net.cuda().eval()
s1, s2 = net.hack_forward(item)
for i, (s1, s2) in enumerate(zip(s1, s2)):
pts_ax = fig.add_subplot(num_models,
rows,
i * rows + 2,
projection='3d')
pred_ax = fig.add_subplot(num_models,
rows,
i * rows + 3,
projection='3d')
if i == 0:
pts_ax.set_title('')
pred_ax.set_title('Prediction')
colors = ['b', 'g', 'r', 'c', 'm', 'y', 'k']
for j in range(s1.size(0)):
col = colors[j % len(colors)]
p = s1[j, :].unsqueeze(0)
ax_points(pts_ax, p.tolist(), s=8, c=col)
for j in range(s2.size(0)):
col = colors[j % len(colors)]
lcl_pred = s2[j, :, :]
ax_points(pred_ax, lcl_pred.tolist(), s=3, c=col)
plt.tight_layout()
plt.savefig(output)
def load_item(path, num_models):
classes = ['chair', 'bathtub', 'toilet', 'night_stand']
tf = tfs.Compose([tfs.UnitSpherePointCloud()])
"""
tf = tfs.Compose([tfs.UnitSpherePointCloud(),
tfs.RandomRotatePointCloud(type='upright')])
"""
ds = ModelNetPointCloud(basedir=path, split='test', categories=classes,
device='cuda', transform=tf,
num_points=2**10, sample_points=2**12)
# Get data
loader = DataLoader(ds, batch_size=num_models, shuffle=True)
for item in loader:
return item
def __init__(self, data_path, rotate, classes=None, jitter=None):
t = tfs.Compose([
tfs.UnitSpherePointCloud(),
tfs.RandomRotatePointCloud(type=rotate)
])
self.train_dataset = ModelNetPointCloud(basedir=data_path,
split='train',
categories=classes,
transform=t,
num_points=2**10,
sample_points=2**12)
self.valid_dataset = ModelNetPointCloud(basedir=data_path,
split='test',
categories=classes,
transform=t,
num_points=2**10,
sample_points=2**12)
self.dataset = self.train_dataset
def load_item(path):
classes = None
tf = tfs.Compose([
tfs.UnitSpherePointCloud(),
tfs.RandomRotatePointCloud(type='upright')
])
ds = ModelNetPointCloud(basedir=path,
split='test',
categories=classes,
device='cuda',
transform=tf,
num_points=2**10,
sample_points=2**12)
# Get data
loader = DataLoader(ds, batch_size=1, shuffle=True)
for item in loader:
return item