Create submission for the Sintel leaderboard
(model, output_path='kitti_submission', sigma=0.05)
| 131 | |
| 132 | @torch.no_grad() |
| 133 | def create_kitti_submission(model, output_path='kitti_submission', sigma=0.05): |
| 134 | """ Create submission for the Sintel leaderboard """ |
| 135 | |
| 136 | IMAGE_SIZE = [432, 1242] |
| 137 | |
| 138 | print(f"output path: {output_path}") |
| 139 | print(f"image size: {IMAGE_SIZE}") |
| 140 | print(f"training size: {TRAIN_SIZE}") |
| 141 | |
| 142 | hws = compute_grid_indices(IMAGE_SIZE) |
| 143 | weights = compute_weight(hws, (432, 1242), TRAIN_SIZE, sigma) |
| 144 | model.eval() |
| 145 | test_dataset = datasets.KITTI(split='testing', aug_params=None) |
| 146 | |
| 147 | if not os.path.exists(output_path): |
| 148 | os.makedirs(output_path) |
| 149 | |
| 150 | for test_id in range(len(test_dataset)): |
| 151 | image1, image2, (frame_id, ) = test_dataset[test_id] |
| 152 | new_shape = image1.shape[1:] |
| 153 | if new_shape[1] != IMAGE_SIZE[1]: # fix the height=432, adaptive ajust the width |
| 154 | print(f"replace {IMAGE_SIZE} with {new_shape}") |
| 155 | IMAGE_SIZE[0] = 432 |
| 156 | IMAGE_SIZE[1] = new_shape[1] |
| 157 | hws = compute_grid_indices(IMAGE_SIZE) |
| 158 | weights = compute_weight(hws, IMAGE_SIZE, TRAIN_SIZE, sigma) |
| 159 | |
| 160 | padder = InputPadder(image1.shape, mode='kitti432') # padding the image to height of 432 |
| 161 | image1, image2 = padder.pad(image1[None].cuda(), image2[None].cuda()) |
| 162 | |
| 163 | flows = 0 |
| 164 | flow_count = 0 |
| 165 | |
| 166 | for idx, (h, w) in enumerate(hws): |
| 167 | image1_tile = image1[:, :, h:h+TRAIN_SIZE[0], w:w+TRAIN_SIZE[1]] |
| 168 | image2_tile = image2[:, :, h:h+TRAIN_SIZE[0], w:w+TRAIN_SIZE[1]] |
| 169 | flow_pre, _ = model(image1_tile, image2_tile) |
| 170 | |
| 171 | padding = (w, IMAGE_SIZE[1]-w-TRAIN_SIZE[1], h, IMAGE_SIZE[0]-h-TRAIN_SIZE[0], 0, 0) |
| 172 | flows += F.pad(flow_pre * weights[idx], padding) |
| 173 | flow_count += F.pad(weights[idx], padding) |
| 174 | |
| 175 | flow_pre = flows / flow_count |
| 176 | flow = padder.unpad(flow_pre[0]).permute(1, 2, 0).cpu().numpy() |
| 177 | |
| 178 | output_filename = os.path.join(output_path, frame_id) |
| 179 | frame_utils.writeFlowKITTI(output_filename, flow) |
| 180 | |
| 181 | flow_img = flow_viz.flow_to_image(flow) |
| 182 | image = Image.fromarray(flow_img) |
| 183 | if not os.path.exists(f'vis_kitti_3patch'): |
| 184 | os.makedirs(f'vis_kitti_3patch/flow') |
| 185 | os.makedirs(f'vis_kitti_3patch/image') |
| 186 | |
| 187 | image.save(f'vis_kitti_3patch/flow/{test_id}.png') |
| 188 | imageio.imwrite(f'vis_kitti_3patch/image/{test_id}_0.png', image1[0].cpu().permute(1, 2, 0).numpy()) |
| 189 | imageio.imwrite(f'vis_kitti_3patch/image/{test_id}_1.png', image2[0].cpu().permute(1, 2, 0).numpy()) |
| 190 |
nothing calls this directly
no test coverage detected