| 66 | self.block_tea = IFBlock(10+4, c=90) |
| 67 | |
| 68 | def forward(self, x, scale_list=[4, 2, 1], training=False): |
| 69 | if training == False: |
| 70 | channel = x.shape[1] // 2 |
| 71 | img0 = x[:, :channel] |
| 72 | img1 = x[:, channel:] |
| 73 | flow_list = [] |
| 74 | merged = [] |
| 75 | mask_list = [] |
| 76 | warped_img0 = img0 |
| 77 | warped_img1 = img1 |
| 78 | flow = (x[:, :4]).detach() * 0 |
| 79 | mask = (x[:, :1]).detach() * 0 |
| 80 | block = [self.block0, self.block1, self.block2] |
| 81 | for i in range(3): |
| 82 | f0, m0 = block[i](torch.cat((warped_img0[:, :3], warped_img1[:, :3], mask), 1), flow, scale=scale_list[i]) |
| 83 | f1, m1 = block[i](torch.cat((warped_img1[:, :3], warped_img0[:, :3], -mask), 1), torch.cat((flow[:, 2:4], flow[:, :2]), 1), scale=scale_list[i]) |
| 84 | flow = flow + (f0 + torch.cat((f1[:, 2:4], f1[:, :2]), 1)) / 2 |
| 85 | mask = mask + (m0 + (-m1)) / 2 |
| 86 | mask_list.append(mask) |
| 87 | flow_list.append(flow) |
| 88 | warped_img0 = warp(img0, flow[:, :2], device=x.device) |
| 89 | warped_img1 = warp(img1, flow[:, 2:4], device=x.device) |
| 90 | merged.append((warped_img0, warped_img1)) |
| 91 | ''' |
| 92 | c0 = self.contextnet(img0, flow[:, :2]) |
| 93 | c1 = self.contextnet(img1, flow[:, 2:4]) |
| 94 | tmp = self.unet(img0, img1, warped_img0, warped_img1, mask, flow, c0, c1) |
| 95 | res = tmp[:, 1:4] * 2 - 1 |
| 96 | ''' |
| 97 | for i in range(3): |
| 98 | mask_list[i] = torch.sigmoid(mask_list[i]) |
| 99 | merged[i] = merged[i][0] * mask_list[i] + merged[i][1] * (1 - mask_list[i]) |
| 100 | return flow_list, mask_list[2], merged |
| 101 | |
| 102 | @staticmethod |
| 103 | def state_dict_converter(): |