(dataset_path: str, split: str, sample_id: str, max_depth)
| 7 | |
| 8 | |
| 9 | def sample_vis(dataset_path: str, split: str, sample_id: str, max_depth): |
| 10 | dataset = ARKitScenesDataset(root=dataset_path, split=split) |
| 11 | video_id = sample_id.split('_')[0] |
| 12 | idx = None |
| 13 | for i in range(len(dataset)): |
| 14 | if dataset.samples[i][0] == video_id and dataset.samples[i][1] == sample_id: |
| 15 | idx = i |
| 16 | break |
| 17 | if idx is None: |
| 18 | raise ValueError(f'Can\'t find sample from split={split}, with video_id={video_id} and sample_id={sample_id}') |
| 19 | sample = dataset[idx] |
| 20 | |
| 21 | max_depth = np.min([max_depth, |
| 22 | sample[dataset_keys.HIGH_RES_DEPTH_IMG].max(), |
| 23 | sample[dataset_keys.LOW_RES_DEPTH_IMG].max()]) |
| 24 | fig, axes = plt.subplots(2, 2) |
| 25 | axes[0, 0].set_title('Color img') |
| 26 | axes[0, 0].axis(False) |
| 27 | axes[0, 0].imshow(image_chw_to_hwc(sample[dataset_keys.COLOR_IMG]/255)) |
| 28 | axes[0, 1].set_title('High Res img (0=no depth)') |
| 29 | axes[0, 1].axis(False) |
| 30 | img = axes[0, 1].imshow(sample[dataset_keys.HIGH_RES_DEPTH_IMG][0], vmin=0, vmax=max_depth, cmap=plt.get_cmap('turbo')) |
| 31 | fig.colorbar(img, ax=axes[0, 1]) |
| 32 | axes[1, 1].set_title('Low Res img') |
| 33 | axes[1, 1].axis(False) |
| 34 | img = axes[1, 1].imshow(sample[dataset_keys.LOW_RES_DEPTH_IMG][0], vmin=0, vmax=max_depth, cmap=plt.get_cmap('turbo')) |
| 35 | fig.colorbar(img, ax=axes[1, 1]) |
| 36 | axes[1, 0].set_title('Color and low res overlay') |
| 37 | axes[1, 0].axis(False) |
| 38 | axes[1, 0].imshow(image_chw_to_hwc(sample[dataset_keys.COLOR_IMG]/255)) |
| 39 | axes[1, 0].imshow(sample[dataset_keys.LOW_RES_DEPTH_IMG][0], vmin=0, vmax=max_depth, cmap=plt.get_cmap('turbo'), alpha=0.5) |
| 40 | plt.show() |
| 41 | plt.waitforbuttonpress() |
| 42 | |
| 43 | |
| 44 | if __name__ == '__main__': |
no test coverage detected