(opt: Options)
| 139 | return predicted_depth # [B, C, H, W] |
| 140 | |
| 141 | def run_inference(opt: Options): |
| 142 | train_dataset = DAVISDataset(opt, training=True) |
| 143 | test_dataset = DAVISDataset(opt, training=False) |
| 144 | dataloader_set = [DataLoader(train_dataset, batch_size=1, shuffle=False, num_workers=8), |
| 145 | DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=8) |
| 146 | ] |
| 147 | |
| 148 | model = Predictor(opt) |
| 149 | model.eval() |
| 150 | device = "cuda" |
| 151 | model.to(device) |
| 152 | |
| 153 | print("Running inference...") |
| 154 | |
| 155 | with torch.no_grad(): |
| 156 | for dataloader in dataloader_set: |
| 157 | for batch in dataloader: |
| 158 | frames = batch["frames"].to(device) |
| 159 | depths = batch["depths"].to(device) |
| 160 | predicted_depth = model(frames, depths) # shape: [B, C, H, W] |
| 161 | |
| 162 | |
| 163 | for i in range(predicted_depth.shape[0]): |
| 164 | depth_img = predicted_depth[i].cpu().numpy() |
| 165 | depth = depth_img |
| 166 | |
| 167 | if depth_img.ndim == 3 and depth_img.shape[0] > 1: |
| 168 | depth_img = depth_img[0] |
| 169 | elif depth_img.ndim == 3: |
| 170 | depth_img = depth_img[0] |
| 171 | |
| 172 | depth_norm = cv2.normalize(depth_img, None, 0, 255, cv2.NORM_MINMAX) |
| 173 | depth_norm = depth_norm.astype(np.uint8) |
| 174 | |
| 175 | save_path = batch["predicted_depth_paths"][i] |
| 176 | os.makedirs(osp.dirname(save_path), exist_ok=True) |
| 177 | cv2.imwrite(save_path, depth_norm) |
| 178 | print(f"Saved predicted depth image to {save_path}") |
| 179 | |
| 180 | if __name__ == "__main__": |
| 181 | import tyro |
no test coverage detected