| 179 | |
| 180 | |
| 181 | def visualize_pointcloud_batch(path, pointclouds, pred_labels, labels, categories, vis_label=False, target=None, elev=30, azim=225): |
| 182 | batch_size = len(pointclouds) |
| 183 | fig = plt.figure(figsize=(20,20)) |
| 184 | |
| 185 | ncols = int(np.sqrt(batch_size)) |
| 186 | nrows = max(1, (batch_size-1) // ncols+1) |
| 187 | for idx, pc in enumerate(pointclouds): |
| 188 | if vis_label: |
| 189 | label = categories[labels[idx].item()] |
| 190 | pred = categories[pred_labels[idx]] |
| 191 | colour = 'g' if label == pred else 'r' |
| 192 | elif target is None: |
| 193 | |
| 194 | colour = 'g' |
| 195 | else: |
| 196 | colour = target[idx] |
| 197 | pc = pc.cpu().numpy() |
| 198 | ax = fig.add_subplot(nrows, ncols, idx + 1, projection='3d') |
| 199 | ax.scatter(pc[:, 0], pc[:, 2], pc[:, 1], c=colour, s=5) |
| 200 | ax.view_init(elev=elev, azim=azim) |
| 201 | ax.axis('off') |
| 202 | if vis_label: |
| 203 | ax.set_title('GT: {0}\nPred: {1}'.format(label, pred)) |
| 204 | |
| 205 | plt.savefig(path) |
| 206 | plt.close(fig) |
| 207 | |
| 208 | |
| 209 | ''' |