| 25 | #---------------------------------------------------------------------------- |
| 26 | |
| 27 | def layout_grid(img, grid_w=None, grid_h=1, float_to_uint8=True, chw_to_hwc=True, to_numpy=True): |
| 28 | batch_size, channels, img_h, img_w = img.shape |
| 29 | if grid_w is None: |
| 30 | grid_w = batch_size // grid_h |
| 31 | assert batch_size == grid_w * grid_h |
| 32 | if float_to_uint8: |
| 33 | img = (img * 127.5 + 128).clamp(0, 255).to(torch.uint8) |
| 34 | img = img.reshape(grid_h, grid_w, channels, img_h, img_w) |
| 35 | img = img.permute(2, 0, 3, 1, 4) |
| 36 | img = img.reshape(channels, grid_h * img_h, grid_w * img_w) |
| 37 | if chw_to_hwc: |
| 38 | img = img.permute(1, 2, 0) |
| 39 | if to_numpy: |
| 40 | img = img.cpu().numpy() |
| 41 | return img |
| 42 | |
| 43 | def create_samples(N=256, voxel_origin=[0, 0, 0], cube_length=2.0): |
| 44 | # NOTE: the voxel_origin is actually the (bottom, left, down) corner, not the middle |