| 60 | |
| 61 | |
| 62 | def resize(sample, size): |
| 63 | if len(sample.shape) == 4: |
| 64 | return torch.nn.functional.interpolate(sample, size=size) |
| 65 | elif len(sample.shape) == 3: |
| 66 | if sample.shape[0] > 3: |
| 67 | # H x W x C -> C x H x W |
| 68 | sample = sample.permute(2, 0, 1) |
| 69 | return torch.nn.functional.interpolate(sample[None], size=size).squeeze(0) |
| 70 | elif len(sample.shape) == 2: |
| 71 | return torch.nn.functional.interpolate(sample[None, None], size=size).squeeze(0).squeeze(0) |
| 72 | else: |
| 73 | raise ValueError("Invalid shape: {}".format(sample.shape)) |
| 74 | |
| 75 | |
| 76 | def preprocess_depth(depth, dataset_name, out_channels=3, keep_raw_depth=False): |