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