| 40 | |
| 41 | |
| 42 | def calculate_valid_mask(depth, dataset_name): |
| 43 | if dataset_name == "hypersim": |
| 44 | q_min, q_max = [1e-4, min(MAX_FAR_PLANE, 1e9)] |
| 45 | elif dataset_name == "vkitti2": |
| 46 | q_min, q_max = [0, min(MAX_FAR_PLANE, 655)] |
| 47 | elif dataset_name == "diode": |
| 48 | q_min, q_max = [0.1, 120] |
| 49 | else: |
| 50 | raise ValueError(f"Unknown dataset {dataset_name}") |
| 51 | |
| 52 | if isinstance(depth, np.ndarray): |
| 53 | mask = np.logical_and(depth >= q_min, depth <= q_max) |
| 54 | elif isinstance(depth, torch.Tensor): |
| 55 | mask = torch.logical_and(depth >= q_min, depth <= q_max) |
| 56 | else: |
| 57 | raise ValueError(f"Invalid depth type {type(depth)}") |
| 58 | |
| 59 | return mask[None] |
| 60 | |
| 61 | |
| 62 | def resize(sample, size): |