(self, img1, img2)
| 50 | return flow |
| 51 | |
| 52 | def image_pair_process(self, img1, img2): |
| 53 | if len(img1.shape) == 2: |
| 54 | img1 = np.tile(img1[..., None], (1, 1, 3)) |
| 55 | img2 = np.tile(img2[..., None], (1, 1, 3)) |
| 56 | |
| 57 | img1 = img1.astype(np.uint8)[..., :3] |
| 58 | img2 = img2.astype(np.uint8)[..., :3] |
| 59 | |
| 60 | img1 = torch.from_numpy(img1).permute(2, 0, 1).float() |
| 61 | img2 = torch.from_numpy(img2).permute(2, 0, 1).float() |
| 62 | return img1, img2 |
| 63 | |
| 64 | @torch.no_grad() |
| 65 | def inference_model(self, im1, im2): |