[B, H, W, C]->[B, H, W, c*k1*k2]
(x, patch=3, stride=1)
| 336 | |
| 337 | |
| 338 | def _make_vector(x, patch=3, stride=1): |
| 339 | """[B, H, W, C]->[B, H, W, c*k1*k2]""" |
| 340 | k1, k2 = to_list(patch, 2) |
| 341 | h, w = tf.shape(x)[1], tf.shape(x)[2] |
| 342 | padded_x = tf.pad(x, [[0, 0], [k1 // 2] * 2, [k2 // 2] * 2, [0, 0]]) |
| 343 | vec = [] |
| 344 | for i in range(k1): |
| 345 | for j in range(k2): |
| 346 | vec.append(padded_x[:, i:i + h:stride, j:j + w:stride, :]) |
| 347 | return tf.concat(vec, axis=-1) |
| 348 | |
| 349 | |
| 350 | def _make_displacement(x, patch=3, max_dis=1, stride1=1, stride2=1): |