Input: im1, im2: HxWx3 numpy (it is not necessary for im1 and im2 to be same shape, because im2 will be resized to the shape of im1) Output: flow: Bx2xHxW torch.Tensor
(self, im1, im2)
| 21 | # print(parameters) |
| 22 | |
| 23 | def estimate(self, im1, im2): |
| 24 | ''' |
| 25 | Input: |
| 26 | im1, im2: HxWx3 numpy (it is not necessary for im1 and im2 to be same shape, because im2 will be resized to |
| 27 | the shape of im1) |
| 28 | Output: |
| 29 | flow: Bx2xHxW torch.Tensor |
| 30 | ''' |
| 31 | h1, w1 = im1.shape[:2] |
| 32 | h2, w2 = im2.shape[:2] |
| 33 | |
| 34 | im2_resized = cv2.resize(im2, (w1, h1)) |
| 35 | im1, im2 = self.image_pair_process(im1, im2_resized) |
| 36 | padder = InputPadder(im1.shape) |
| 37 | im1, im2 = padder.pad(im1[None].cuda(), im2[None].cuda()) |
| 38 | # print(im1) |
| 39 | flow = self.inference_model(im1, im2)[0] |
| 40 | # print(flow) |
| 41 | |
| 42 | flow = padder.unpad(flow).permute(1, 2, 0).repeat(1, 1, 1, 1).permute([0, 3, 1, 2]) |
| 43 | grid = coords_grid(1, h1, w1).cuda() |
| 44 | flow += grid |
| 45 | flow[:,0] = flow[:,0] * w2 / w1 |
| 46 | flow[:,1] = flow[:,1] * h2 / h1 |
| 47 | flow -= grid |
| 48 | |
| 49 | |
| 50 | return flow |
| 51 | |
| 52 | def image_pair_process(self, img1, img2): |
| 53 | if len(img1.shape) == 2: |
no test coverage detected