Correlation Cost Volume computation. This is a fallback Python-only implementation, specialized just for FlowNet2. It takes a lot of memory and is slow. If you know to compile a custom op yourself, it's better to use the cuda implementation here: https://github.com/PatWie/tens
(ina, inb,
kernel_size, max_displacement,
stride_1, stride_2,
pad, data_format)
| 36 | |
| 37 | |
| 38 | def correlation(ina, inb, |
| 39 | kernel_size, max_displacement, |
| 40 | stride_1, stride_2, |
| 41 | pad, data_format): |
| 42 | """ |
| 43 | Correlation Cost Volume computation. |
| 44 | |
| 45 | This is a fallback Python-only implementation, specialized just for FlowNet2. |
| 46 | It takes a lot of memory and is slow. |
| 47 | |
| 48 | If you know to compile a custom op yourself, it's better to use the cuda implementation here: |
| 49 | https://github.com/PatWie/tensorflow-recipes/tree/master/OpticalFlow/user_ops |
| 50 | """ |
| 51 | assert pad == max_displacement |
| 52 | assert kernel_size == 1 |
| 53 | assert data_format == 'NCHW' |
| 54 | assert max_displacement % stride_2 == 0 |
| 55 | assert stride_1 == 1 |
| 56 | |
| 57 | D = int(max_displacement / stride_2 * 2) + 1 # D^2 == number of correlations per spatial location |
| 58 | |
| 59 | b, c, h, w = ina.shape.as_list() |
| 60 | |
| 61 | inb = tf.pad(inb, [[0, 0], [0, 0], [pad, pad], [pad, pad]]) |
| 62 | |
| 63 | res = [] |
| 64 | for k1 in range(0, D): |
| 65 | start_h = k1 * stride_2 |
| 66 | for k2 in range(0, D): |
| 67 | start_w = k2 * stride_2 |
| 68 | s = tf.slice(inb, [0, 0, start_h, start_w], [-1, -1, h, w]) |
| 69 | ans = tf.reduce_mean(ina * s, axis=1, keepdims=True) |
| 70 | res.append(ans) |
| 71 | res = tf.concat(res, axis=1) # ND^2HW |
| 72 | return res |
| 73 | |
| 74 | |
| 75 | def resample(img, flow): |