Experimental: Add initial wavefront/obj parser for vertices

WORKSPACE

@@ -1110,3 +1110,14 @@ http_archive(
         "https://github.com/mongodb/mongo-c-driver/releases/download/1.16.2/mongo-c-driver-1.16.2.tar.gz",
     ],
 )
+
+http_archive(
+    name = "tinyobjloader",
+    build_file = "//third_party:tinyobjloader.BUILD",
+    sha256 = "b8c972dfbbcef33d55554e7c9031abe7040795b67778ad3660a50afa7df6ec56",
+    strip_prefix = "tinyobjloader-2.0.0rc8",
+    urls = [
+        "https://storage.googleapis.com/mirror.tensorflow.org/github.com/tinyobjloader/tinyobjloader/archive/v2.0.0rc8.tar.gz",
+        "https://github.com/tinyobjloader/tinyobjloader/archive/v2.0.0rc8.tar.gz",
+    ],
+)

tensorflow_io/core/BUILD

@@ -695,6 +695,22 @@ cc_library(
     alwayslink = 1,
 )
 
+cc_library(
+    name = "obj_ops",
+    srcs = [
+        "kernels/obj_kernels.cc",
+        "ops/obj_ops.cc",
+    ],
+    copts = tf_io_copts(),
+    linkstatic = True,
+    deps = [
+        "@local_config_tf//:libtensorflow_framework",
+        "@local_config_tf//:tf_header_lib",
+        "@tinyobjloader",
+    ],
+    alwayslink = 1,
+)
+
 cc_binary(
     name = "python/ops/libtensorflow_io.so",
     copts = tf_io_copts(),
@@ -717,6 +733,7 @@ cc_binary(
         "//tensorflow_io/core:parquet_ops",
         "//tensorflow_io/core:pcap_ops",
         "//tensorflow_io/core:pulsar_ops",
+        "//tensorflow_io/core:obj_ops",
         "//tensorflow_io/core:operation_ops",
         "//tensorflow_io/core:pubsub_ops",
         "//tensorflow_io/core:serialization_ops",

tensorflow_io/core/kernels/obj_kernels.cc

@@ -0,0 +1,70 @@
+/* Copyright 2021 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/core/platform/logging.h"
+#include "tiny_obj_loader.h"
+
+namespace tensorflow {
+namespace io {
+namespace {
+
+class DecodeObjOp : public OpKernel {
+ public:
+  explicit DecodeObjOp(OpKernelConstruction* context) : OpKernel(context) {}
+
+  void Compute(OpKernelContext* context) override {
+    const Tensor* input_tensor;
+    OP_REQUIRES_OK(context, context->input("input", &input_tensor));
+    OP_REQUIRES(context, TensorShapeUtils::IsScalar(input_tensor->shape()),
+                errors::InvalidArgument("input must be scalar, got shape ",
+                                        input_tensor->shape().DebugString()));
+    const tstring& input = input_tensor->scalar<tstring>()();
+
+    tinyobj::ObjReader reader;
+
+    if (!reader.ParseFromString(input.c_str(), "")) {
+      OP_REQUIRES(
+          context, false,
+          errors::Internal("Unable to read obj file: ", reader.Error()));
+    }
+
+    if (!reader.Warning().empty()) {
+      LOG(WARNING) << "TinyObjReader: " << reader.Warning();
+    }
+
+    auto& attrib = reader.GetAttrib();
+
+    int64 count = attrib.vertices.size() / 3;
+
+    Tensor* output_tensor = nullptr;
+    OP_REQUIRES_OK(context, context->allocate_output(0, TensorShape({count, 3}),
+                                                     &output_tensor));
+    // Loop over attrib.vertices:
+    for (int64 i = 0; i < count; i++) {
+      tinyobj::real_t x = attrib.vertices[i * 3 + 0];
+      tinyobj::real_t y = attrib.vertices[i * 3 + 1];
+      tinyobj::real_t z = attrib.vertices[i * 3 + 2];
+      output_tensor->tensor<float, 2>()(i, 0) = x;
+      output_tensor->tensor<float, 2>()(i, 1) = y;
+      output_tensor->tensor<float, 2>()(i, 2) = z;
+    }
+  }
+};
+REGISTER_KERNEL_BUILDER(Name("IO>DecodeObj").Device(DEVICE_CPU), DecodeObjOp);
+
+}  // namespace
+}  // namespace io
+}  // namespace tensorflow

tensorflow_io/core/ops/obj_ops.cc

@@ -0,0 +1,36 @@
+/* Copyright 2021 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/common_shape_fns.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/shape_inference.h"
+
+namespace tensorflow {
+namespace io {
+namespace {
+
+REGISTER_OP("IO>DecodeObj")
+    .Input("input: string")
+    .Output("output: float32")
+    .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(), 3}));
+      return Status::OK();
+    });
+
+}  // namespace
+}  // namespace io
+}  // namespace tensorflow

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

@@ -27,4 +27,5 @@
     decode_yuy2,
     decode_avif,
     decode_jp2,
+    decode_obj,
 )

tensorflow_io/core/python/experimental/image_ops.py

@@ -208,3 +208,18 @@ def decode_jp2(contents, dtype=tf.uint8, name=None):
       A `Tensor` of type `uint8` and shape of `[height, width, 3]` (RGB).
     """
     return core_ops.io_decode_jpeg2k(contents, dtype=dtype, name=name)
+
+
+def decode_obj(contents, name=None):
+    """
+    Decode a Wavefront (obj) file into a float32 tensor.
+
+    Args:
+      contents: A 0-dimensional Tensor of type string, i.e the
+        content of the Wavefront (.obj) file.
+      name: A name for the operation (optional).
+
+    Returns:
+      A `Tensor` of type `float32` and shape of `[n, 3]` for vertices.
+    """
+    return core_ops.io_decode_obj(contents, name=name)

tests/test_obj.py

@@ -0,0 +1,37 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""Test Wavefront OBJ"""
+
+import os
+import numpy as np
+import pytest
+
+import tensorflow as tf
+import tensorflow_io as tfio
+
+
+def test_decode_obj():
+    """Test case for decode obj"""
+    filename = os.path.join(
+        os.path.dirname(os.path.abspath(__file__)), "test_obj", "sample.obj",
+    )
+    filename = "file://" + filename
+
+    obj = tfio.experimental.image.decode_obj(tf.io.read_file(filename))
+    expected = np.array(
+        [[-0.5, 0.0, 0.4], [-0.5, 0.0, -0.8], [-0.5, 1.0, -0.8], [-0.5, 1.0, 0.4]],
+        dtype=np.float32,
+    )
+    assert np.array_equal(obj, expected)

tests/test_obj/sample.obj

@@ -0,0 +1,6 @@
+# Simple Wavefront file
+v -0.500000 0.000000 0.400000
+v -0.500000 0.000000 -0.800000
+v -0.500000 1.000000 -0.800000
+v -0.500000 1.000000 0.400000
+f -4 -3 -2 -1

third_party/tinyobjloader.BUILD

@@ -0,0 +1,14 @@
+package(default_visibility = ["//visibility:public"])
+
+licenses(["notice"])  # MIT license
+
+cc_library(
+    name = "tinyobjloader",
+    srcs = [
+        "tiny_obj_loader.cc",
+    ],
+    hdrs = [
+        "tiny_obj_loader.h",
+    ],
+    copts = [],
+)