Create submission for the Sintel leaderboard
(model, iters=32, warm_start=False, output_path='sintel_submission')
| 20 | |
| 21 | @torch.no_grad() |
| 22 | def create_sintel_submission(model, iters=32, warm_start=False, output_path='sintel_submission'): |
| 23 | """ Create submission for the Sintel leaderboard """ |
| 24 | model.eval() |
| 25 | for dstype in ['clean', 'final']: |
| 26 | test_dataset = datasets.MpiSintel(split='test', aug_params=None, dstype=dstype) |
| 27 | |
| 28 | flow_prev, sequence_prev = None, None |
| 29 | for test_id in range(len(test_dataset)): |
| 30 | image1, image2, (sequence, frame) = test_dataset[test_id] |
| 31 | if sequence != sequence_prev: |
| 32 | flow_prev = None |
| 33 | |
| 34 | padder = InputPadder(image1.shape) |
| 35 | image1, image2 = padder.pad(image1[None].cuda(), image2[None].cuda()) |
| 36 | |
| 37 | flow_low, flow_pr = model(image1, image2, iters=iters, flow_init=flow_prev, test_mode=True) |
| 38 | flow = padder.unpad(flow_pr[0]).permute(1, 2, 0).cpu().numpy() |
| 39 | |
| 40 | if warm_start: |
| 41 | flow_prev = forward_interpolate(flow_low[0])[None].cuda() |
| 42 | |
| 43 | output_dir = os.path.join(output_path, dstype, sequence) |
| 44 | output_file = os.path.join(output_dir, 'frame%04d.flo' % (frame+1)) |
| 45 | |
| 46 | if not os.path.exists(output_dir): |
| 47 | os.makedirs(output_dir) |
| 48 | |
| 49 | frame_utils.writeFlow(output_file, flow) |
| 50 | sequence_prev = sequence |
| 51 | |
| 52 | |
| 53 | @torch.no_grad() |
nothing calls this directly
no test coverage detected