Add decode_yuy2 (YUYV/YUY2/YUYV422) support

This PR adds support of decode_yuy2 to allow convert yuy2 to rgb.
This is useful for Video Capture in Linux (Video4Linux2) as the default
format is YUYV.

This PR fixes tensorflow#825

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

Author: yongtang

tensorflow_io/core/BUILD

@@ -168,6 +168,7 @@ cc_library(
         "kernels/image_pnm_kernels.cc",
         "kernels/image_tiff_kernels.cc",
         "kernels/image_webp_kernels.cc",
+        "kernels/image_yuy2_kernels.cc",
         "ops/image_ops.cc",
     ],
     copts = tf_io_copts(),

tensorflow_io/core/kernels/image_yuy2_kernels.cc

@@ -0,0 +1,79 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow_io/core/kernels/io_stream.h"
+
+#include "libyuv/convert_argb.h"
+#include "libyuv/convert_from_argb.h"
+
+namespace tensorflow {
+namespace io {
+namespace {
+
+class DecodeYUY2Op : public OpKernel {
+ public:
+  explicit DecodeYUY2Op(OpKernelConstruction* context) : OpKernel(context) {
+    env_ = context->env();
+  }
+
+  void Compute(OpKernelContext* context) override {
+    const Tensor* input_tensor;
+    OP_REQUIRES_OK(context, context->input("input", &input_tensor));
+
+    const Tensor* size_tensor;
+    OP_REQUIRES_OK(context, context->input("size", &size_tensor));
+
+    const tstring& input = input_tensor->scalar<tstring>()();
+
+    int64 channels = 3;
+    int64 height = size_tensor->flat<int32>()(0);
+    int64 width = size_tensor->flat<int32>()(1);
+
+    Tensor* image_tensor = nullptr;
+    OP_REQUIRES_OK(
+        context, context->allocate_output(
+                     0, TensorShape({height, width, channels}), &image_tensor));
+
+    string buffer;
+    buffer.resize(width * height * 4);
+    uint8* argb = (uint8*)&buffer[0];
+    uint8* yuy2 = (uint8*)&input[0];
+    uint32 yuy2_stride = width * 2;
+    uint32 argb_stride = width * 4;
+    int status =
+        libyuv::YUY2ToARGB(yuy2, yuy2_stride, argb, argb_stride, width, height);
+    OP_REQUIRES(
+        context, (status == 0),
+        errors::InvalidArgument("unable to convert yuy2 to argb: ", status));
+
+    uint8* rgb = image_tensor->flat<uint8>().data();
+    uint32 rgb_stride = width * 3;
+    status =
+        libyuv::ARGBToRAW(argb, argb_stride, rgb, rgb_stride, width, height);
+    OP_REQUIRES(
+        context, (status == 0),
+        errors::InvalidArgument("unable to convert argb to rgb: ", status));
+  }
+
+ private:
+  mutex mu_;
+  Env* env_ GUARDED_BY(mu_);
+};
+REGISTER_KERNEL_BUILDER(Name("IO>DecodeYUY2").Device(DEVICE_CPU), DecodeYUY2Op);
+
+}  // namespace
+}  // namespace io
+}  // namespace tensorflow

tensorflow_io/core/ops/image_ops.cc

@@ -183,6 +183,17 @@ REGISTER_OP("IO>DecodeNV12")
       return Status::OK();
     });
 
+REGISTER_OP("IO>DecodeYUY2")
+    .Input("input: string")
+    .Input("size: int32")
+    .Output("image: uint8")
+    .SetShapeFn([](shape_inference::InferenceContext* c) {
+      shape_inference::ShapeHandle unused;
+      TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &unused));
+      c->set_output(0, c->MakeShape({c->UnknownDim(), c->UnknownDim(), 3}));
+      return Status::OK();
+    });
+
 }  // namespace
 }  // namespace io
 }  // namespace tensorflow

tensorflow_io/core/python/api/experimental/image.py

@@ -24,4 +24,5 @@
     decode_pnm,
     decode_hdr,
     decode_nv12,
+    decode_yuy2,
 )

tensorflow_io/core/python/experimental/image_ops.py

@@ -137,7 +137,7 @@ def decode_pnm(contents, dtype=tf.uint8, name=None):
 
 def decode_hdr(contents, name=None):
     """
-    Decode a HDR-encoded image to a uint8 tensor.
+    Decode a HDR-encoded image to a tf.float tensor.
 
     Args:
       contents: A `Tensor` of type `string`. 0-D.  The HDR-encoded image.
@@ -163,3 +163,19 @@ def decode_nv12(contents, size, name=None):
       A `Tensor` of type `uint8` and shape of `[height, width, 3]` (RGB).
     """
     return core_ops.io_decode_nv12(contents, size=size, name=name)
+
+
+def decode_yuy2(contents, size, name=None):
+    """
+    Decode a YUY2-encoded image to a uint8 tensor.
+
+    Args:
+      contents: A `Tensor` of type `string`. 0-D.  The YUY2-encoded image.
+      size: A 1-D int32 Tensor of 2 elements: height, width. The size
+        for the images.
+      name: A name for the operation (optional).
+
+    Returns:
+      A `Tensor` of type `uint8` and shape of `[height, width, 3]` (RGB).
+    """
+    return core_ops.io_decode_yuy2(contents, size=size, name=name)

tests/test_image/Jelly-Beans.yuv tests/test_image/Jelly-Beans.nv12

@@ -294,10 +294,10 @@ def test_decode_hdr():
 def test_decode_nv12():
     """Test case for decode_nv12"""
     filename = os.path.join(
-        os.path.dirname(os.path.abspath(__file__)), "test_image", "Jelly-Beans.yuv"
+        os.path.dirname(os.path.abspath(__file__)), "test_image", "Jelly-Beans.nv12"
     )
     png_filename = os.path.join(
-        os.path.dirname(os.path.abspath(__file__)), "test_image", "Jelly-Beans.png"
+        os.path.dirname(os.path.abspath(__file__)), "test_image", "Jelly-Beans.nv12.png"
     )
     png = tf.image.decode_png(tf.io.read_file(png_filename))
 
@@ -308,5 +308,22 @@ def test_decode_nv12():
     assert np.all(rgb == png)
 
 
+def test_decode_yuy2():
+    """Test case for decode_yuy2"""
+    filename = os.path.join(
+        os.path.dirname(os.path.abspath(__file__)), "test_image", "Jelly-Beans.yuy2"
+    )
+    png_filename = os.path.join(
+        os.path.dirname(os.path.abspath(__file__)), "test_image", "Jelly-Beans.yuy2.png"
+    )
+    png = tf.image.decode_png(tf.io.read_file(png_filename))
+
+    contents = tf.io.read_file(filename)
+    rgb = tfio.experimental.image.decode_yuy2(contents, size=[256, 256])
+    assert rgb.dtype == tf.uint8
+    assert rgb.shape == [256, 256, 3]
+    assert np.all(rgb == png)
+
+
 if __name__ == "__main__":
     test.main()

tests/test_image/Jelly-Beans.png tests/test_image/Jelly-Beans.nv12.png

@@ -13,6 +13,7 @@ cc_library(
         "source/scale_*.cc",
     ]) + [
         "source/convert_argb.cc",
+        "source/convert_from_argb.cc",
         "source/cpu_id.cc",
         "source/planar_functions.cc",
     ],