bilinear sample image at coordinate x, y Args: image: a 4-D tensor of shape [B H W C] x, y: a 3-D tensor of shape [B H W], pixel index of image Return: sampled images, that output[i, j] = image[y[i], x[j]]
(image, x, y)
| 64 | |
| 65 | |
| 66 | def _sample(image, x, y): |
| 67 | """bilinear sample image at coordinate x, y |
| 68 | |
| 69 | Args: |
| 70 | image: a 4-D tensor of shape [B H W C] |
| 71 | x, y: a 3-D tensor of shape [B H W], pixel index of image |
| 72 | |
| 73 | Return: |
| 74 | sampled images, that |
| 75 | output[i, j] = image[y[i], x[j]] |
| 76 | """ |
| 77 | |
| 78 | shape = tf.shape(image) |
| 79 | batch = shape[0] |
| 80 | h = shape[1] |
| 81 | w = shape[2] |
| 82 | |
| 83 | x = tf.to_float(x) |
| 84 | y = tf.to_float(y) |
| 85 | image = tf.to_float(image) |
| 86 | x0 = tf.to_int32(tf.floor(x)) |
| 87 | y0 = tf.to_int32(tf.floor(y)) |
| 88 | x1 = x0 + 1 |
| 89 | y1 = y0 + 1 |
| 90 | |
| 91 | w00 = tf.expand_dims((tf.to_float(x1) - x) * (tf.to_float(y1) - y), -1) |
| 92 | w01 = tf.expand_dims((x - tf.to_float(x0)) * (tf.to_float(y1) - y), -1) |
| 93 | w10 = tf.expand_dims((tf.to_float(x1) - x) * (y - tf.to_float(y0)), -1) |
| 94 | w11 = tf.expand_dims((x - tf.to_float(x0)) * (y - tf.to_float(y0)), -1) |
| 95 | |
| 96 | x0 = tf.clip_by_value(x0, 0, w - 1) |
| 97 | y0 = tf.clip_by_value(y0, 0, h - 1) |
| 98 | x1 = tf.clip_by_value(x1, 0, w - 1) |
| 99 | y1 = tf.clip_by_value(y1, 0, h - 1) |
| 100 | |
| 101 | batch_idx = tf.reshape(tf.range(0, batch), [batch, 1, 1]) |
| 102 | batch_idx = tf.tile(batch_idx, [1, h, w]) |
| 103 | gather_00 = tf.stack([batch_idx, y0, x0], axis=-1) |
| 104 | gather_01 = tf.stack([batch_idx, y0, x1], axis=-1) |
| 105 | gather_10 = tf.stack([batch_idx, y1, x0], axis=-1) |
| 106 | gather_11 = tf.stack([batch_idx, y1, x1], axis=-1) |
| 107 | |
| 108 | p00 = tf.gather_nd(image, gather_00) * w00 |
| 109 | p01 = tf.gather_nd(image, gather_01) * w01 |
| 110 | p10 = tf.gather_nd(image, gather_10) * w10 |
| 111 | p11 = tf.gather_nd(image, gather_11) * w11 |
| 112 | |
| 113 | return tf.add_n([p00, p01, p10, p11]) |
| 114 | |
| 115 | |
| 116 | def _move(image, x, y): |