Experimental: Add initial wavefront/obj parser for vertices (tensorflow#1315)

This PR is an early experimental implementation of wavefront obj
parser in tensorflow-io for 3D objects.
This PR is the first step to obtain raw vertices in float32
tensor with shape of `[n, 3]`.

Additional follow up PRs will be needed to handle meshs with
different shapes (not sure if ragged tensor will be a good fit
in that case)

Some background on obj file:
Wavefront (obj) is a format widely used in 3D (another is ply)
modeling (http://paulbourke.net/dataformats/obj/). It is simple
(ASCII) with good support for many softwares. Machine learning
in 3D has been an active field with some advances such as
PolyGen (https://arxiv.org/abs/2002.10880)

Processing obj files are needed to process 3D with tensorflow.

In 3D the basic elements could be vertices or faces. This PR
tries to cover vertices first so that vertices in obj file
can be loaded into TF's graph for further processing within
graph pipeline.

Signed-off-by: Yong Tang <[email protected]>

Author: yongtang

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 = [],
+)