| 37 | return torch.stack([b1, b2, b3], dim=1).permute(0, 2, 1) |
| 38 | |
| 39 | def debug_vis_input(batch_data, cfg, prd_data, iter_counts): |
| 40 | for i in range(cfg.exp.batch_size): |
| 41 | # print("i is,", i) |
| 42 | save_dir = cfg.exp.vis_dir |
| 43 | vis_dir = os.path.join(save_dir, 'vis_B_input') |
| 44 | if not os.path.exists((vis_dir)): |
| 45 | os.mkdir(vis_dir) |
| 46 | # src_pc = batch_data['src_pc'][i] # (3, 1024) |
| 47 | tgt_pc = batch_data['tgt_pc'][i] # (3, 1024) |
| 48 | |
| 49 | device = tgt_pc.device |
| 50 | num_points = tgt_pc.shape[1] |
| 51 | |
| 52 | total_pc = np.array(tgt_pc.detach().cpu()) |
| 53 | color_mask_tgt = ['g'] * num_points |
| 54 | |
| 55 | fig = plt.figure() |
| 56 | ax = fig.add_subplot(projection='3d') |
| 57 | ax.scatter3D(list(total_pc[0]), list(total_pc[1]), list(total_pc[2]), c=color_mask_tgt, s=10, alpha=0.9) |
| 58 | ax.axis('scaled') |
| 59 | ax.set_zlabel('Z', fontdict={'size': 20, 'color': 'red'}) |
| 60 | ax.set_ylabel('Y', fontdict={'size': 20, 'color': 'red'}) |
| 61 | ax.set_xlabel('X', fontdict={'size': 20, 'color': 'red'}) |
| 62 | |
| 63 | fig.savefig(os.path.join(vis_dir, 'batch_{}_{}_input'.format(iter_counts,i))) |
| 64 | plt.close(fig) |
| 65 | |
| 66 | def debug_vis_output(batch_data, cfg, pred_data, iter_counts): |
| 67 | |