forked from alex04072000/ObstructionRemoval
-
Notifications
You must be signed in to change notification settings - Fork 0
/
warp_utils.py
202 lines (161 loc) · 8.93 KB
/
warp_utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
"""
core_warp.py
Warp a level of image1's feature pyramid using the up-sampled flow at level+1 of image2's pyramid.
The official TF implementation requires that batch image height and width be known during graph building.
Our implementation below doesn't have the same requirements. We care about this flexibility because we sometimes train
on small image patches (e.g., (384, 448)) but do online validation on larger sizes (e.g., (436, 1024)).
Written by The TensorFlow Authors, modifications by Phil Ferriere
Licensed under the MIT License (see LICENSE for details)
Based on:
- https://github.com/tensorflow/tensorflow/blob/r1.10/tensorflow/contrib/image/python/ops/dense_image_warp.py
Written by The TensorFlow Authors, Copyright 2018 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License 2.0
TF Ref:
- tf.contrib.image.dense_image_warp
https://www.tensorflow.org/api_docs/python/tf/contrib/image/dense_image_warp
Notes:
- A list of CUDA-accelerated implementations:
https://github.com/sampepose/flownet2-tf/tree/master/src/ops/correlation
https://github.com/jgorgenucsd/corr_tf
https://github.com/simonmeister/UnFlow/blob/8e74f2b33138ab72d775bf1c3a9256105677834e/ops/correlation_op.cu.cc
- Will there be a cost-volume implementation in tf.contrib at some point?
See https://github.com/tensorflow/tensorflow/pull/21392
"""
from __future__ import absolute_import, division, print_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
def _interpolate_bilinear(grid,
query_points,
name='interpolate_bilinear',
indexing='ij'):
"""Similar to Matlab's interp2 function.
Finds values for query points on a grid using bilinear interpolation.
Args:
grid: a 4-D float `Tensor` of shape `[batch, height, width, channels]`.
query_points: a 3-D float `Tensor` of N points with shape `[batch, N, 2]`.
name: a name for the operation (optional).
indexing: whether the query points are specified as row and column (ij),
or Cartesian coordinates (xy).
Returns:
values: a 3-D `Tensor` with shape `[batch, N, channels]`
Raises:
ValueError: if the indexing mode is invalid, or if the shape of the inputs
invalid.
"""
if indexing != 'ij' and indexing != 'xy':
raise ValueError('Indexing mode must be \'ij\' or \'xy\'')
with ops.name_scope(name):
grid = ops.convert_to_tensor(grid)
query_points = ops.convert_to_tensor(query_points)
shape = array_ops.unstack(array_ops.shape(grid))
if len(shape) != 4:
msg = 'Grid must be 4 dimensional. Received: '
raise ValueError(msg + str(shape))
batch_size, height, width, channels = shape
query_type = query_points.dtype
query_shape = array_ops.unstack(array_ops.shape(query_points))
grid_type = grid.dtype
if len(query_shape) != 3:
msg = ('Query points must be 3 dimensional. Received: ')
raise ValueError(msg + str(query_shape))
_, num_queries, _ = query_shape
alphas = []
floors = []
ceils = []
index_order = [0, 1] if indexing == 'ij' else [1, 0]
unstacked_query_points = array_ops.unstack(query_points, axis=2)
for dim in index_order:
with ops.name_scope('dim-' + str(dim)):
queries = unstacked_query_points[dim]
size_in_indexing_dimension = shape[dim + 1]
# max_floor is size_in_indexing_dimension - 2 so that max_floor + 1
# is still a valid index into the grid.
max_floor = math_ops.cast(size_in_indexing_dimension - 2, query_type)
min_floor = constant_op.constant(0.0, dtype=query_type)
floor = math_ops.minimum(
math_ops.maximum(min_floor, math_ops.floor(queries)), max_floor)
int_floor = math_ops.cast(floor, dtypes.int32)
floors.append(int_floor)
ceil = int_floor + 1
ceils.append(ceil)
# alpha has the same type as the grid, as we will directly use alpha
# when taking linear combinations of pixel values from the image.
alpha = math_ops.cast(queries - floor, grid_type)
min_alpha = constant_op.constant(0.0, dtype=grid_type)
max_alpha = constant_op.constant(1.0, dtype=grid_type)
alpha = math_ops.minimum(math_ops.maximum(min_alpha, alpha), max_alpha)
# Expand alpha to [b, n, 1] so we can use broadcasting
# (since the alpha values don't depend on the channel).
alpha = array_ops.expand_dims(alpha, 2)
alphas.append(alpha)
flattened_grid = array_ops.reshape(grid,
[batch_size * height * width, channels])
batch_offsets = array_ops.reshape(
math_ops.range(batch_size) * height * width, [batch_size, 1])
# This wraps array_ops.gather. We reshape the image data such that the
# batch, y, and x coordinates are pulled into the first dimension.
# Then we gather. Finally, we reshape the output back. It's possible this
# code would be made simpler by using array_ops.gather_nd.
def gather(y_coords, x_coords, name):
with ops.name_scope('gather-' + name):
linear_coordinates = batch_offsets + y_coords * width + x_coords
gathered_values = array_ops.gather(flattened_grid, linear_coordinates)
return array_ops.reshape(gathered_values,
[batch_size, num_queries, channels])
# grab the pixel values in the 4 corners around each query point
top_left = gather(floors[0], floors[1], 'top_left')
top_right = gather(floors[0], ceils[1], 'top_right')
bottom_left = gather(ceils[0], floors[1], 'bottom_left')
bottom_right = gather(ceils[0], ceils[1], 'bottom_right')
# now, do the actual interpolation
with ops.name_scope('interpolate'):
interp_top = alphas[1] * (top_right - top_left) + top_left
interp_bottom = alphas[1] * (bottom_right - bottom_left) + bottom_left
interp = alphas[0] * (interp_bottom - interp_top) + interp_top
return interp
def dense_image_warp(image, flow, name='dense_image_warp'):
"""Image warping using per-pixel flow vectors.
Apply a non-linear warp to the image, where the warp is specified by a dense
flow field of offset vectors that define the correspondences of pixel values
in the output image back to locations in the source image. Specifically, the
pixel value at output[b, j, i, c] is
images[b, j - flow[b, j, i, 0], i - flow[b, j, i, 1], c].
The locations specified by this formula do not necessarily map to an int
index. Therefore, the pixel value is obtained by bilinear
interpolation of the 4 nearest pixels around
(b, j - flow[b, j, i, 0], i - flow[b, j, i, 1]). For locations outside
of the image, we use the nearest pixel values at the image boundary.
Args:
image: 4-D float `Tensor` with shape `[batch, height, width, channels]`.
flow: A 4-D float `Tensor` with shape `[batch, height, width, 2]`.
name: A name for the operation (optional).
Note that image and flow can be of type tf.half, tf.float32, or tf.float64,
and do not necessarily have to be the same type.
Returns:
A 4-D float `Tensor` with shape`[batch, height, width, channels]`
and same type as input image.
Raises:
ValueError: if height < 2 or width < 2 or the inputs have the wrong number
of dimensions.
"""
with ops.name_scope(name):
batch_size, height, width, channels = array_ops.unstack(array_ops.shape(image))
# The flow is defined on the image grid. Turn the flow into a list of query
# points in the grid space.
grid_x, grid_y = array_ops.meshgrid(
math_ops.range(width), math_ops.range(height))
stacked_grid = math_ops.cast(
array_ops.stack([grid_y, grid_x], axis=2), flow.dtype)
batched_grid = array_ops.expand_dims(stacked_grid, axis=0)
query_points_on_grid = batched_grid - flow
query_points_flattened = array_ops.reshape(query_points_on_grid,
[batch_size, height * width, 2])
# Compute values at the query points, then reshape the result back to the
# image grid.
interpolated = _interpolate_bilinear(image, query_points_flattened)
interpolated = array_ops.reshape(interpolated,
[batch_size, height, width, channels])
return interpolated