From 0c2797320064159bac8a06bb0cc03ac0d867e7e8 Mon Sep 17 00:00:00 2001 From: Yong Tang Date: Thu, 25 Feb 2021 17:27:43 -0800 Subject: [PATCH] Experimental: Add initial wavefront/obj parser for vertices 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 --- WORKSPACE | 11 +++ tensorflow_io/core/BUILD | 17 +++++ tensorflow_io/core/kernels/obj_kernels.cc | 70 +++++++++++++++++++ tensorflow_io/core/ops/obj_ops.cc | 36 ++++++++++ .../core/python/api/experimental/image.py | 1 + .../core/python/experimental/image_ops.py | 15 ++++ tests/test_obj.py | 37 ++++++++++ tests/test_obj/sample.obj | 6 ++ third_party/tinyobjloader.BUILD | 14 ++++ 9 files changed, 207 insertions(+) create mode 100644 tensorflow_io/core/kernels/obj_kernels.cc create mode 100644 tensorflow_io/core/ops/obj_ops.cc create mode 100644 tests/test_obj.py create mode 100644 tests/test_obj/sample.obj create mode 100644 third_party/tinyobjloader.BUILD diff --git a/WORKSPACE b/WORKSPACE index 12238181b..d5f7ab095 100644 --- a/WORKSPACE +++ b/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", + ], +) diff --git a/tensorflow_io/core/BUILD b/tensorflow_io/core/BUILD index 01958c5ee..41141c2b1 100644 --- a/tensorflow_io/core/BUILD +++ b/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", diff --git a/tensorflow_io/core/kernels/obj_kernels.cc b/tensorflow_io/core/kernels/obj_kernels.cc new file mode 100644 index 000000000..e619431e1 --- /dev/null +++ b/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()(); + + 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()(i, 0) = x; + output_tensor->tensor()(i, 1) = y; + output_tensor->tensor()(i, 2) = z; + } + } +}; +REGISTER_KERNEL_BUILDER(Name("IO>DecodeObj").Device(DEVICE_CPU), DecodeObjOp); + +} // namespace +} // namespace io +} // namespace tensorflow diff --git a/tensorflow_io/core/ops/obj_ops.cc b/tensorflow_io/core/ops/obj_ops.cc new file mode 100644 index 000000000..e3c45653a --- /dev/null +++ b/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 diff --git a/tensorflow_io/core/python/api/experimental/image.py b/tensorflow_io/core/python/api/experimental/image.py index 0d9c07e74..643b7b9c0 100644 --- a/tensorflow_io/core/python/api/experimental/image.py +++ b/tensorflow_io/core/python/api/experimental/image.py @@ -27,4 +27,5 @@ decode_yuy2, decode_avif, decode_jp2, + decode_obj, ) diff --git a/tensorflow_io/core/python/experimental/image_ops.py b/tensorflow_io/core/python/experimental/image_ops.py index b399de102..ebde7e6ae 100644 --- a/tensorflow_io/core/python/experimental/image_ops.py +++ b/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) diff --git a/tests/test_obj.py b/tests/test_obj.py new file mode 100644 index 000000000..ff89ef2b8 --- /dev/null +++ b/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) diff --git a/tests/test_obj/sample.obj b/tests/test_obj/sample.obj new file mode 100644 index 000000000..da8b327ff --- /dev/null +++ b/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 diff --git a/third_party/tinyobjloader.BUILD b/third_party/tinyobjloader.BUILD new file mode 100644 index 000000000..0e9f74df4 --- /dev/null +++ b/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 = [], +)