(model, image1, image2, weights=None)
| 67 | return patch_weights |
| 68 | |
| 69 | def compute_flow(model, image1, image2, weights=None): |
| 70 | print(f"computing flow...") |
| 71 | |
| 72 | image_size = image1.shape[1:] |
| 73 | |
| 74 | image1, image2 = image1[None].cuda(), image2[None].cuda() |
| 75 | |
| 76 | hws = compute_grid_indices(image_size) |
| 77 | if weights is None: # no tile |
| 78 | padder = InputPadder(image1.shape) |
| 79 | image1, image2 = padder.pad(image1, image2) |
| 80 | |
| 81 | flow_pre, _ = model(image1, image2) |
| 82 | |
| 83 | flow_pre = padder.unpad(flow_pre) |
| 84 | flow = flow_pre[0].permute(1, 2, 0).cpu().numpy() |
| 85 | else: # tile |
| 86 | flows = 0 |
| 87 | flow_count = 0 |
| 88 | |
| 89 | for idx, (h, w) in enumerate(hws): |
| 90 | image1_tile = image1[:, :, h:h+TRAIN_SIZE[0], w:w+TRAIN_SIZE[1]] |
| 91 | image2_tile = image2[:, :, h:h+TRAIN_SIZE[0], w:w+TRAIN_SIZE[1]] |
| 92 | flow_pre, _ = model(image1_tile, image2_tile) |
| 93 | padding = (w, image_size[1]-w-TRAIN_SIZE[1], h, image_size[0]-h-TRAIN_SIZE[0], 0, 0) |
| 94 | flows += F.pad(flow_pre * weights[idx], padding) |
| 95 | flow_count += F.pad(weights[idx], padding) |
| 96 | |
| 97 | flow_pre = flows / flow_count |
| 98 | flow = flow_pre[0].permute(1, 2, 0).cpu().numpy() |
| 99 | |
| 100 | return flow |
| 101 | |
| 102 | def compute_adaptive_image_size(image_size): |
| 103 | target_size = TRAIN_SIZE |
no test coverage detected