1 /* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #include "tensorflow/core/framework/common_shape_fns.h"
17 #include "tensorflow/core/framework/op.h"
18 #include "tensorflow/core/framework/shape_inference.h"
19
20 namespace tensorflow {
21
22 using shape_inference::DimensionHandle;
23 using shape_inference::InferenceContext;
24 using shape_inference::ShapeHandle;
25
26 namespace {
27
28 // Sets output[0] to shape [batch_dim,height,width,channel_dim], where
29 // height and width come from the size_tensor.
SetOutputToSizedImage(InferenceContext * c,DimensionHandle batch_dim,int size_input_idx,DimensionHandle channel_dim)30 Status SetOutputToSizedImage(InferenceContext* c, DimensionHandle batch_dim,
31 int size_input_idx, DimensionHandle channel_dim) {
32 // Verify shape of size input.
33 ShapeHandle size;
34 TF_RETURN_IF_ERROR(c->WithRank(c->input(size_input_idx), 1, &size));
35 DimensionHandle unused;
36 TF_RETURN_IF_ERROR(c->WithValue(c->Dim(size, 0), 2, &unused));
37
38 // Get size values from the size tensor.
39 const Tensor* size_tensor = c->input_tensor(size_input_idx);
40 DimensionHandle width;
41 DimensionHandle height;
42 if (size_tensor == nullptr) {
43 width = c->UnknownDim();
44 height = c->UnknownDim();
45 } else {
46 // TODO(petewarden) - Remove once we have constant evaluation in C++ only.
47 if (size_tensor->dtype() != DT_INT32) {
48 return errors::InvalidArgument(
49 "Bad size input type for SetOutputToSizedImage: Expected DT_INT32 "
50 "but got ",
51 DataTypeString(size_tensor->dtype()), " for input #", size_input_idx,
52 " in ", c->DebugString());
53 }
54 auto vec = size_tensor->vec<int32>();
55 height = c->MakeDim(vec(0));
56 width = c->MakeDim(vec(1));
57 }
58 c->set_output(0, c->MakeShape({batch_dim, height, width, channel_dim}));
59 return Status::OK();
60 }
61
62 // TODO(qyu): Move this to core/framework/common_shape_fns.h
ResizeShapeFn(InferenceContext * c)63 Status ResizeShapeFn(InferenceContext* c) {
64 ShapeHandle input;
65 TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 4, &input));
66 return SetOutputToSizedImage(c, c->Dim(input, 0), 2 /* size_input_idx */,
67 c->Dim(input, 3));
68 }
69
70 static const char kImageProjectiveTransformDoc[] = R"doc(
71 Applies the given transform to each of the images.
72
73 Input `image` is a `Tensor` in NHWC format (where the axes are image in batch,
74 rows, columns, and channels. Input `transforms` is a num_images x 8 or 1 x 8
75 matrix, where each row corresponds to a 3 x 3 projective transformation matrix,
76 with the last entry assumed to be 1. If there is one row, the same
77 transformation will be applied to all images.
78
79 If one row of `transforms` is `[a0, a1, a2, b0, b1, b2, c0, c1]`, then it maps
80 the *output* point `(x, y)` to a transformed *input* point
81 `(x', y') = ((a0 x + a1 y + a2) / k, (b0 x + b1 y + b2) / k)`, where
82 `k = c0 x + c1 y + 1`. If the transformed point lays outside of the input
83 image, the output pixel is set to 0.
84
85 images: 4D `Tensor`, input image(s) in NHWC format.
86 transforms: 2D `Tensor`, projective transform(s) to apply to the image(s).
87
88 transformed_images: 4D `Tensor`, image(s) in NHWC format, generated by applying
89 the `transforms` to the `images`. Satisfies the description above.
90 )doc";
91
92 } // namespace
93
94 // TODO(ringwalt): Add a "fill_mode" attr with "constant", "mirror", etc.
95 // TODO(ringwalt): Add a "fill_constant" argument for constant mode (default 0).
96 REGISTER_OP("ImageProjectiveTransform")
97 .Input("images: dtype")
98 .Input("transforms: float32")
99 .Attr("dtype: {uint8, int32, int64, float16, float32, float64}")
100 .Attr("interpolation: string")
101 .Output("transformed_images: dtype")
102 // Output shape is identical to input images.
__anone7bf74500202(InferenceContext* c) 103 .SetShapeFn([](InferenceContext* c) {
104 c->set_output(0, c->input(0));
105 return Status::OK();
106 })
107 .Doc(kImageProjectiveTransformDoc);
108
109 // V2 op supports output_shape.
110 REGISTER_OP("ImageProjectiveTransformV2")
111 .Input("images: dtype")
112 .Input("transforms: float32")
113 .Input("output_shape: int32")
114 .Attr("dtype: {uint8, int32, int64, float16, float32, float64}")
115 .Attr("interpolation: string")
116 .Output("transformed_images: dtype")
117 .SetShapeFn(ResizeShapeFn)
118 .Doc(kImageProjectiveTransformDoc);
119
120 REGISTER_OP("BipartiteMatch")
121 .Input("distance_mat: float")
122 .Input("num_valid_rows: float")
123 .Attr("top_k: int = -1")
124 .Output("row_to_col_match_indices: int32")
125 .Output("col_to_row_match_indices: int32")
126 .SetIsStateful()
__anone7bf74500302(InferenceContext* c) 127 .SetShapeFn([](InferenceContext* c) {
128 ShapeHandle input;
129 TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 2, &input));
130 c->set_output(0, c->MakeShape({c->Dim(input, 0)}));
131 c->set_output(1, c->MakeShape({c->Dim(input, 1)}));
132 return Status::OK();
133 })
134 .Doc(R"doc(
135 Find bipartite matching based on a given distance matrix.
136
137 A greedy bi-partite matching algorithm is used to obtain the matching with the
138 (greedy) minimum distance.
139
140 distance_mat: A 2-D float tensor of shape `[num_rows, num_columns]`. It is a
141 pair-wise distance matrix between the entities represented by each row and
142 each column. It is an asymmetric matrix. The smaller the distance is, the more
143 similar the pairs are. The bipartite matching is to minimize the distances.
144 num_valid_rows: A scalar or a 1-D tensor with one element describing the
145 number of valid rows of distance_mat to consider for the bipartite matching.
146 If set to be negative, then all rows from `distance_mat` are used.
147 top_k: A scalar that specifies the number of top-k matches to retrieve.
148 If set to be negative, then is set according to the maximum number of
149 matches from `distance_mat`.
150 row_to_col_match_indices: A vector of length num_rows, which is the number of
151 rows of the input `distance_matrix`.
152 If `row_to_col_match_indices[i]` is not -1, row i is matched to column
153 `row_to_col_match_indices[i]`.
154 col_to_row_match_indices: A vector of length num_columns, which is the number
155 of columns of the input distance matrix.
156 If `col_to_row_match_indices[j]` is not -1, column j is matched to row
157 `col_to_row_match_indices[j]`.
158 )doc");
159
160 REGISTER_OP("ImageConnectedComponents")
161 .Input("image: dtype")
162 .Output("components: int64")
163 .Attr(
164 "dtype: {int64, int32, uint16, int16, uint8, int8, half, float, "
165 "double, bool, string}")
__anone7bf74500402(InferenceContext* c) 166 .SetShapeFn([](InferenceContext* c) {
167 return shape_inference::UnchangedShape(c);
168 })
169 .Doc(R"doc(
170 Find the connected components of image(s).
171
172 For each image (along the 0th axis), all connected components of adjacent pixels
173 with the same non-zero value are detected and given unique ids.
174
175 The returned `components` tensor has 0s for the zero pixels of `images`, and
176 arbitrary nonzero ids for the connected components of nonzero values. Ids are
177 unique across all of the images, and are in row-major order by the first pixel
178 in the component.
179
180 Uses union-find with union by rank but not path compression, giving a runtime of
181 `O(n log n)`. See:
182 https://en.wikipedia.org/wiki/Disjoint-set_data_structure#Time_Complexity
183
184 image: Image(s) with shape (N, H, W).
185 components: Component ids for each pixel in "image". Same shape as "image". Zero
186 pixels all have an output of 0, and all components of adjacent pixels with
187 the same value are given consecutive ids, starting from 1.
188 )doc");
189
190 } // namespace tensorflow
191