Improve shape inference with `DecodeAndCropJpeg` (#13750)

* Improve shape inference with `DecodeAndCropJpeg`

While working on improving shape inference for several
other ops in 13561 and 13193, I noticed that `DecodeAndCropJpeg`
does not inference shape even though crop size might have already
be provided. In that case the shape will be `[h, w, channel]`
and `h`, `w` is part of the `crop_window`.

This fix updates the shape function in `DecodeAndCropJpeg`
for improving shape inference.

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

* Add test cases to cover shape inference for `DecodeAndCropJpeg`

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

* Address failed unit tests

Signed-off-by: Yong Tang <yong.tang.github@outlook.com>

2 files changed, 33 insertions, 4 deletions

diff --git a/tensorflow/core/ops/image_ops.cc b/tensorflow/core/ops/image_ops.cc
index 66765a3333..89c9da81c5 100644
--- a/tensorflow/core/ops/image_ops.cc
+++ b/tensorflow/core/ops/image_ops.cc
@@ -453,7 +453,36 @@ REGISTER_OP("DecodeAndCropJpeg")
     .Attr("acceptable_fraction: float = 1.0")
     .Attr("dct_method: string = ''")
     .Output("image: uint8")
-    .SetShapeFn(DecodeImageShapeFn)
+    .SetShapeFn([](InferenceContext* c) {
+      ShapeHandle unused;
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &unused));
+      DimensionHandle channels_dim = c->UnknownDim();
+      DimensionHandle h = c->UnknownDim();
+      DimensionHandle w = c->UnknownDim();
+
+      int32 channels;
+      TF_RETURN_IF_ERROR(c->GetAttr("channels", &channels));
+      if (channels != 0) {
+        if (channels < 0) {
+          return errors::InvalidArgument("channels must be non-negative, got ",
+                                         channels);
+        }
+        channels_dim = c->MakeDim(channels);
+      }
+
+      DimensionHandle unused_dim;
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &unused));
+      TF_RETURN_IF_ERROR(c->WithValue(c->Dim(unused, 0), 4, &unused_dim));
+
+      const Tensor* crop_window = c->input_tensor(1);
+      if (crop_window != nullptr) {
+        auto crop_window_vec = crop_window->vec<int32>();
+        h = c->MakeDim(crop_window_vec(2));
+        w = c->MakeDim(crop_window_vec(3));
+      }
+      c->set_output(0, c->MakeShape({h, w, channels_dim}));
+      return Status::OK();
+    })
     .Doc(strings::StrCat(R"doc(
 Decode and Crop a JPEG-encoded image to a uint8 tensor.
 )doc",
diff --git a/tensorflow/core/ops/image_ops_test.cc b/tensorflow/core/ops/image_ops_test.cc
index c34b11a15e..5f0b391b0d 100644
--- a/tensorflow/core/ops/image_ops_test.cc
+++ b/tensorflow/core/ops/image_ops_test.cc
@@ -105,7 +105,7 @@ TEST(ImageOpsTest, DecodeAndCropJpeg_ShapeFn) {
                    .Input({"img", 0, DT_STRING})
                    .Input({"crop_window", 1, DT_INT32})
                    .Finalize(&op.node_def));
-  INFER_OK(op, "[];[]", "[?,?,?]");
+  INFER_OK(op, "[];[?]", "[?,?,?]");
 
   // Set the channel, so that part of output shape is known.
   TF_ASSERT_OK(NodeDefBuilder("test", op_name)
@@ -113,7 +113,7 @@ TEST(ImageOpsTest, DecodeAndCropJpeg_ShapeFn) {
                    .Input({"crop_window", 1, DT_INT32})
                    .Attr("channels", 4)
                    .Finalize(&op.node_def));
-  INFER_OK(op, "[];[]", "[?,?,4]");
+  INFER_OK(op, "[];[?]", "[?,?,4]");
 
   // Negative channel value is rejected.
   TF_ASSERT_OK(NodeDefBuilder("test", op_name)
@@ -139,7 +139,7 @@ TEST(ImageOpsTest, DecodeAndCropJpeg_InvalidCropWindow) {
                    .Input({"img", 0, DT_STRING})
                    .Input({"crop_window", 1, DT_INT32})
                    .Finalize(&op.node_def));
-  INFER_OK(op, "[];[]", "[?,?,?]");
+  INFER_OK(op, "[];[?]", "[?,?,?]");
 }
 
 TEST(ImageOpsTest, EncodeImage_ShapeFn) {