warp an image/tensor (im2) back to im1, according to the optical flow x: [B, C, H, W] (im2) flo: [B, 2, H, W] flow
(self, x, flo)
| 97 | self.dc_conv7 = Conv2D("dc_conv7", 2, filter_size=3,stride=1,padding=1, param_attr=self.param_attr) |
| 98 | |
| 99 | def warp(self, x, flo): |
| 100 | """ |
| 101 | warp an image/tensor (im2) back to im1, according to the optical flow |
| 102 | |
| 103 | x: [B, C, H, W] (im2) |
| 104 | flo: [B, 2, H, W] flow |
| 105 | |
| 106 | """ |
| 107 | |
| 108 | B, C, H, W = x.shape |
| 109 | # mesh grid |
| 110 | xx_pd = fluid.layers.range(0, W, 1, 'float32') |
| 111 | xx_pd = fluid.layers.reshape(xx_pd, shape=[1, -1]) |
| 112 | xx_pd = fluid.layers.expand(x=xx_pd, expand_times=[H, 1]) |
| 113 | xx_pd = fluid.layers.reshape(xx_pd, shape=[1, 1, H, W]) |
| 114 | xx_pd = fluid.layers.expand(x=xx_pd, expand_times=[B, 1, 1, 1]) |
| 115 | |
| 116 | yy_pd = fluid.layers.range(0, H, 1, 'float32') |
| 117 | yy_pd = fluid.layers.reshape(yy_pd, shape=[-1, 1]) |
| 118 | yy_pd = fluid.layers.expand(x=yy_pd, expand_times=[1, W]) |
| 119 | yy_pd = fluid.layers.reshape(x=yy_pd, shape=[1, 1, H, W]) |
| 120 | yy_pd = fluid.layers.expand(x=yy_pd, expand_times=[B, 1, 1, 1]) |
| 121 | grid_pd = fluid.layers.concat(input=[xx_pd, yy_pd], axis=1) |
| 122 | flo_pd = flo |
| 123 | vgrid_pd = fluid.layers.elementwise_add(grid_pd, flo_pd) |
| 124 | vgrid_pd_0 = 2.0 * fluid.layers.slice(vgrid_pd, axes=[1], starts=[0], ends=[1]) / max(W - 1, 1) - 1.0 |
| 125 | vgrid_pd_1 = 2.0 * fluid.layers.slice(vgrid_pd, axes=[1], starts=[1], ends=[2]) / max(H - 1, 1) - 1.0 |
| 126 | vgrid_pd = fluid.layers.concat(input=[vgrid_pd_0, vgrid_pd_1], axis=1) |
| 127 | vgrid_pd = fluid.layers.transpose(vgrid_pd, [0, 2, 3, 1]) |
| 128 | output = fluid.layers.grid_sampler(name='grid_sample', x=x, grid=vgrid_pd) |
| 129 | |
| 130 | mask = fluid.layers.zeros_like(x) |
| 131 | mask = mask + 1.0 |
| 132 | mask = fluid.layers.grid_sampler(name='grid_sample', x=mask, grid=vgrid_pd) |
| 133 | mask_temp1 = fluid.layers.cast(mask < 0.9990, 'float32') |
| 134 | mask = mask * (1 - mask_temp1) |
| 135 | mask = fluid.layers.cast(mask > 0, 'float32') |
| 136 | outwarp = fluid.layers.elementwise_mul(output, mask) |
| 137 | |
| 138 | return outwarp |
| 139 | |
| 140 | def corr(self, x_1, x_2): |
| 141 | out = correlation(x_1, x_2, pad_size=self.md, kernel_size=1, max_displacement=self.md, |