initial commit

pydiffvg_tensorflow/__init__.py

@@ -0,0 +1,24 @@
+import tensorflow as tf
+try:
+    import diffvg
+except ImportError:
+    print("Warning: diffvg is not installed when you import pydiffvg_tensorflow.")
+from .device import *
+from .shape import *
+from .pixel_filter import *
+from .render_tensorflow import *
+from .image import *
+from .color import *
+import os.path
+
+print(os.path.dirname(diffvg.__file__))
+
+if tf.__cxx11_abi_flag__ == 0:
+    __data_ptr_module = tf.load_op_library(os.path.join(os.path.dirname(diffvg.__file__), 'libdiffvg_tf_data_ptr_no_cxx11_abi.so'))
+else:
+    assert(tf.__cxx11_abi_flag__ == 1)
+    __data_ptr_module = tf.load_op_library(os.path.join(os.path.dirname(diffvg.__file__), 'libdiffvg_tf_data_ptr_cxx11_abi.so'))
+
+def data_ptr(tensor):    
+    addr_as_uint64 = __data_ptr_module.data_ptr(tensor)
+    return int(addr_as_uint64)

pydiffvg_tensorflow/color.py

@@ -0,0 +1,23 @@
+import tensorflow as tf
+
+class LinearGradient:
+    def __init__(self,
+                 begin = tf.constant([0.0, 0.0]),
+                 end = tf.constant([0.0, 0.0]),
+                 offsets = tf.constant([0.0]),
+                 stop_colors = tf.constant([0.0, 0.0, 0.0, 0.0])):
+        self.begin = begin
+        self.end = end
+        self.offsets = offsets
+        self.stop_colors = stop_colors
+
+class RadialGradient:
+    def __init__(self,
+                 center = tf.constant([0.0, 0.0]),
+                 radius = tf.constant([0.0, 0.0]),
+                 offsets = tf.constant([0.0]),
+                 stop_colors = tf.constant([0.0, 0.0, 0.0, 0.0])):
+        self.center = center
+        self.radius = radius
+        self.offsets = offsets
+        self.stop_colors = stop_colors

pydiffvg_tensorflow/custom_ops/CMakeLists.txt

@@ -0,0 +1,29 @@
+cmake_minimum_required(VERSION 3.12)
+
+set(CMAKE_CXX_STANDARD 11)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+project(diffvgTFCustomOp)
+
+set(CMAKE_POSITION_INDEPENDENT_CODE ON)
+
+include_directories(SYSTEM ${TensorFlow_INCLUDE_DIR})
+
+# Compile two versions of the library
+add_library(diffvg_tf_data_ptr_cxx11_abi SHARED data_ptr.cc)
+set_target_properties(diffvg_tf_data_ptr_cxx11_abi PROPERTIES COMPILE_FLAGS -D_GLIBCXX_USE_CXX11_ABI=1)
+set_target_properties(diffvg_tf_data_ptr_cxx11_abi PROPERTIES LINK_FLAGS -D_GLIBCXX_USE_CXX11_ABI=1)
+if(APPLE)
+    # .so instead of .dylib
+    set_target_properties(diffvg_tf_data_ptr_cxx11_abi PROPERTIES SUFFIX .so)
+endif()
+target_link_libraries(diffvg_tf_data_ptr_cxx11_abi ${TensorFlow_LIBRARY})
+
+add_library(diffvg_tf_data_ptr_no_cxx11_abi SHARED data_ptr.cc)
+set_target_properties(diffvg_tf_data_ptr_no_cxx11_abi PROPERTIES COMPILE_FLAGS -D_GLIBCXX_USE_CXX11_ABI=0)
+set_target_properties(diffvg_tf_data_ptr_no_cxx11_abi PROPERTIES LINK_FLAGS -D_GLIBCXX_USE_CXX11_ABI=0)
+if(APPLE)
+    # .so instead of .dylib
+    set_target_properties(diffvg_tf_data_ptr_no_cxx11_abi PROPERTIES SUFFIX .so)
+endif()
+target_link_libraries(diffvg_tf_data_ptr_no_cxx11_abi ${TensorFlow_LIBRARY})

pydiffvg_tensorflow/custom_ops/data_ptr.cc

@@ -0,0 +1,88 @@
+// TODO: add back acknowledgement to the original author when release.
+
+#pragma warning(disable : 4003 4061 4100 4127 4242 4244 4267 4355 4365 4388 4464 4514 4574 4623 4625 4626 4647 4668 4710 4820 4946 5026 5027 5031 5039)
+
+// For windows
+#define NOMINMAX
+
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/shape_inference.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include <stdint.h>
+#include <climits>
+
+using namespace tensorflow;
+
+/* Tensorflow custom ops does not allow parameter types of list of 
+   various data types. Therefore, we can't pass a list but we have
+   to pass each objects individually. 
+
+   Consult Tensorflow source code: /tensorflow/core/framework/tensor.h
+   for what is supported by Tensorflow
+*/
+
+REGISTER_OP("DataPtr")
+    .Attr("T: {float, int32} = DT_INT32")  // To preserve backwards compatibility, you should specify a default value when adding an attr to an existing op:
+    .Input("input: T")  // Tensor
+    .Output("output: uint64")  // scalar
+    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
+      c->set_output(0, {}); // scalar
+      return Status::OK();
+    });
+
+template <typename T>
+class DataPtrOp : public OpKernel {
+ public:
+  explicit DataPtrOp(OpKernelConstruction* context) : OpKernel(context) {}
+
+  void Compute(OpKernelContext* context) override {
+    // Grab the input tensor
+    const Tensor& input_tensor = context->input(0);
+    const T *tensor = input_tensor.flat<T>().data();
+
+    // Create an output tensor
+    // NOTE: The output datatype must match the Ops definition!!!.
+    Tensor* output_tensor = NULL;
+    // Always allocate on CPU
+    AllocatorAttributes alloc_attr;
+    alloc_attr.set_on_host(true);
+    OP_REQUIRES_OK(context, 
+      context->allocate_output(0, {},  // Initialize a one-element scalar
+      &output_tensor,
+      alloc_attr)
+    );
+    auto output_flat = output_tensor->flat<uint64>();
+
+    // Cast pointer to unsigned long int
+    uintptr_t addr = (uintptr_t)tensor;
+
+    // Cast unsigned long int -> unsigned int64
+    uint64 addr_converted = addr;
+
+    output_flat(0) = addr_converted;
+  }
+};
+
+// Polymorphism: https://www.tensorflow.org/guide/extend/op#polymorphism
+REGISTER_KERNEL_BUILDER(
+  Name("DataPtr")
+  .Device(DEVICE_CPU)
+  .TypeConstraint<int32>("T"),
+  DataPtrOp<int32>);
+REGISTER_KERNEL_BUILDER(
+  Name("DataPtr")
+  .Device(DEVICE_CPU)
+  .TypeConstraint<float>("T"),
+  DataPtrOp<float>);
+REGISTER_KERNEL_BUILDER(
+  Name("DataPtr")
+  .Device(DEVICE_GPU)
+  .TypeConstraint<int32>("T")
+  .HostMemory("output"),
+  DataPtrOp<int32>);
+REGISTER_KERNEL_BUILDER(
+  Name("DataPtr")
+  .Device(DEVICE_GPU)
+  .TypeConstraint<float>("T")
+  .HostMemory("output"),
+  DataPtrOp<float>);

pydiffvg_tensorflow/device.py

@@ -0,0 +1,59 @@
+import tensorflow as tf
+
+use_gpu = tf.test.is_gpu_available(
+    cuda_only=True,
+    min_cuda_compute_capability=None
+)
+cpu_device_id = 0
+gpu_device_id = 0
+
+def get_device_name():
+    """
+        Get the current tensorflow device name we are using.
+    """
+    global use_gpu
+    global cpu_device_id
+    global gpu_device_id
+    return '/device:gpu:' + str(gpu_device_id) if use_gpu else '/device:cpu:' + str(cpu_device_id)
+
+def set_use_gpu(v: bool):
+    """
+        Set whether to use CUDA or not.
+    """
+    global use_gpu
+    use_gpu = v
+
+def get_use_gpu():
+    """
+        Get whether we are using CUDA or not.
+    """
+    global use_gpu
+    return use_gpu
+
+def set_cpu_device_id(did: int):
+    """
+        Set the cpu device id we are using.
+    """
+    global cpu_device_id
+    cpu_device_id = did
+
+def get_cpu_device_id():
+    """
+        Get the cpu device id we are using.
+    """
+    global cpu_device_id
+    return cpu_device_id
+
+def set_gpu_device_id(did: int):
+    """
+        Set the gpu device id we are using.
+    """
+    global gpu_device_id
+    gpu_device_id = did
+
+def get_gpu_device_id():
+    """
+        Get the gpu device id we are using.
+    """
+    global gpu_device_id
+    return gpu_device_id

pydiffvg_tensorflow/image.py

@@ -0,0 +1,22 @@
+import numpy as np
+import skimage
+import skimage.io
+import os
+
+def imwrite(img, filename, gamma = 2.2, normalize = False):
+    directory = os.path.dirname(filename)
+    if directory != '' and not os.path.exists(directory):
+        os.makedirs(directory)
+
+    if not isinstance(img, np.ndarray):
+        img = img.numpy()
+    if normalize:
+        img_rng = np.max(img) - np.min(img)
+        if img_rng > 0:
+            img = (img - np.min(img)) / img_rng
+    img = np.clip(img, 0.0, 1.0)
+    if img.ndim==2:
+        #repeat along the third dimension
+        img=np.expand_dims(img,2)
+    img[:, :, :3] = np.power(img[:, :, :3], 1.0/gamma)
+    skimage.io.imsave(filename, (img * 255).astype(np.uint8))

pydiffvg_tensorflow/pixel_filter.py

@@ -0,0 +1,8 @@
+import tensorflow as tf
+
+class PixelFilter:
+    def __init__(self,
+                 type,
+                 radius = tf.constant(0.5)):
+        self.type = type
+        self.radius = radius

pydiffvg_tensorflow/render_tensorflow.py

@@ -0,0 +1,649 @@
+import tensorflow as tf
+import diffvg
+import pydiffvg_tensorflow as pydiffvg
+import time
+from enum import IntEnum
+import warnings
+
+print_timing = False
+__EMPTY_TENSOR = tf.constant([])
+
+def is_empty_tensor(tensor):
+    return  tf.equal(tf.size(tensor), 0)
+
+def set_print_timing(val):
+    global print_timing
+    print_timing=val
+
+class OutputType(IntEnum):
+    color = 1
+    sdf = 2
+
+class ShapeType:
+    __shapetypes = [
+        diffvg.ShapeType.circle,
+        diffvg.ShapeType.ellipse,
+        diffvg.ShapeType.path,
+        diffvg.ShapeType.rect
+    ]
+
+    @staticmethod
+    def asTensor(type):
+        for i in range(len(ShapeType.__shapetypes)):
+            if ShapeType.__shapetypes[i] == type:
+                return tf.constant(i)
+
+    @staticmethod
+    def asShapeType(index: tf.Tensor):
+        if is_empty_tensor(index):
+            return None
+        try:
+            type = ShapeType.__shapetypes[index]
+        except IndexError:
+            print(f'{index} is out of range: [0, {len(ShapeType.__shapetypes)})')
+            import sys
+            sys.exit()
+        else:
+            return type
+
+class ColorType:
+    __colortypes = [
+        diffvg.ColorType.constant,
+        diffvg.ColorType.linear_gradient,
+        diffvg.ColorType.radial_gradient
+    ]
+
+    @staticmethod
+    def asTensor(type):
+        for i in range(len(ColorType.__colortypes)):
+            if ColorType.__colortypes[i] == type:
+                return tf.constant(i)
+
+    @staticmethod
+    def asColorType(index: tf.Tensor):
+        if is_empty_tensor(index):
+            return None
+        try:
+            type = ColorType.__colortypes[index]
+        except IndexError:
+            print(f'{index} is out of range: [0, {len(ColorType.__colortypes)})')
+            import sys
+            sys.exit()
+        else:
+            return type
+
+class FilterType:
+    __filtertypes = [
+        diffvg.FilterType.box,
+        diffvg.FilterType.tent,
+        diffvg.FilterType.hann
+    ]
+
+    @staticmethod
+    def asTensor(type):
+        for i in range(len(FilterType.__filtertypes)):
+            if FilterType.__filtertypes[i] == type:
+                return tf.constant(i)    
+
+    @staticmethod
+    def asFilterType(index: tf.Tensor):
+        if is_empty_tensor(index):
+            return None
+        try:
+            type = FilterType.__filtertypes[index]
+        except IndexError:
+            print(f'{index} is out of range: [0, {len(FilterType.__filtertypes)})')
+            import sys
+            sys.exit()
+        else:
+            return type
+
+def serialize_scene(canvas_width,
+                    canvas_height,
+                    shapes,
+                    shape_groups,
+                    filter = pydiffvg.PixelFilter(type = diffvg.FilterType.box,
+                                                  radius = tf.constant(0.5)),
+                    output_type = OutputType.color,
+                    use_prefiltering = False):
+    """
+        Given a list of shapes, convert them to a linear list of argument,
+        so that we can use it in TF.
+    """
+    with tf.device('/device:cpu:' + str(pydiffvg.get_cpu_device_id())):
+        num_shapes = len(shapes)
+        num_shape_groups = len(shape_groups)
+        args = []
+        args.append(tf.constant(canvas_width))
+        args.append(tf.constant(canvas_height))
+        args.append(tf.constant(num_shapes))
+        args.append(tf.constant(num_shape_groups))
+        args.append(tf.constant(output_type))
+        args.append(tf.constant(use_prefiltering))
+        for shape in shapes:
+            if isinstance(shape, pydiffvg.Circle):
+                args.append(ShapeType.asTensor(diffvg.ShapeType.circle))
+                args.append(tf.identity(shape.radius))
+                args.append(tf.identity(shape.center))
+            elif isinstance(shape, pydiffvg.Ellipse):
+                args.append(ShapeType.asTensor(diffvg.ShapeType.ellipse))
+                args.append(tf.identity(shape.radius))
+                args.append(tf.identity(shape.center))
+            elif isinstance(shape, pydiffvg.Path):
+                assert(shape.points.shape[1] == 2)
+                args.append(ShapeType.asTensor(diffvg.ShapeType.path))
+                args.append(tf.identity(shape.num_control_points, type=tf.int32))
+                args.append(tf.identity(shape.points))
+                args.append(tf.constant(shape.is_closed))
+            elif isinstance(shape, pydiffvg.Polygon):
+                assert(shape.points.shape[1] == 2)
+                args.append(ShapeType.asTensor(diffvg.ShapeType.path))
+                if shape.is_closed:
+                    args.append(tf.zeros(shape.points.shape[0], dtype = tf.int32))
+                else:
+                    args.append(tf.zeros(shape.points.shape[0] - 1, dtype = tf.int32))
+                args.append(tf.identity(shape.points))
+                args.append(tf.constant(shape.is_closed))
+            elif isinstance(shape, pydiffvg.Rect):
+                args.append(ShapeType.asTensor(diffvg.ShapeType.rect))
+                args.append(tf.identity(shape.p_min))
+                args.append(tf.identity(shape.p_max))
+            else:
+                assert(False)
+            args.append(tf.identity(shape.stroke_width))
+
+        for shape_group in shape_groups:
+            args.append(tf.identity(shape_group.shape_ids))
+            # Fill color
+            if shape_group.fill_color is None:
+                args.append(__EMPTY_TENSOR)
+            elif tf.is_tensor(shape_group.fill_color):
+                args.append(ColorType.asTensor(diffvg.ColorType.constant))
+                args.append(tf.identity(shape_group.fill_color))
+            elif isinstance(shape_group.fill_color, pydiffvg.LinearGradient):
+                args.append(ColorType.asTensor(diffvg.ColorType.linear_gradient))
+                args.append(tf.identity(shape_group.fill_color.begin))
+                args.append(tf.identity(shape_group.fill_color.end))
+                args.append(tf.identity(shape_group.fill_color.offsets))
+                args.append(tf.identity(shape_group.fill_color.stop_colors))
+            elif isinstance(shape_group.fill_color, pydiffvg.RadialGradient):
+                args.append(ColorType.asTensor(diffvg.ColorType.radial_gradient))
+                args.append(tf.identity(shape_group.fill_color.center))
+                args.append(tf.identity(shape_group.fill_color.radius))
+                args.append(tf.identity(shape_group.fill_color.offsets))
+                args.append(tf.identity(shape_group.fill_color.stop_colors))
+
+            if shape_group.fill_color is not None:
+                # go through the underlying shapes and check if they are all closed
+                for shape_id in shape_group.shape_ids:
+                    if isinstance(shapes[shape_id], pydiffvg.Path):
+                        if not shapes[shape_id].is_closed:
+                            warnings.warn("Detected non-closed paths with fill color. This might causes unexpected results.", Warning)
+
+            # Stroke color
+            if shape_group.stroke_color is None:
+                args.append(__EMPTY_TENSOR)
+            elif tf.is_tensor(shape_group.stroke_color):
+                args.append(tf.constant(0))
+                args.append(tf.identity(shape_group.stroke_color))
+            elif isinstance(shape_group.stroke_color, pydiffvg.LinearGradient):
+                args.append(ColorType.asTensor(diffvg.ColorType.linear_gradient))
+                args.append(tf.identity(shape_group.stroke_color.begin))
+                args.append(tf.identity(shape_group.stroke_color.end))
+                args.append(tf.identity(shape_group.stroke_color.offsets))
+                args.append(tf.identity(shape_group.stroke_color.stop_colors))
+            elif isinstance(shape_group.stroke_color, pydiffvg.RadialGradient):
+                args.append(ColorType.asTensor(diffvg.ColorType.radial_gradient))
+                args.append(tf.identity(shape_group.stroke_color.center))
+                args.append(tf.identity(shape_group.stroke_color.radius))
+                args.append(tf.identity(shape_group.stroke_color.offsets))
+                args.append(tf.identity(shape_group.stroke_color.stop_colors))
+            args.append(tf.constant(shape_group.use_even_odd_rule))
+            # Transformation
+            args.append(tf.identity(shape_group.shape_to_canvas))
+        args.append(FilterType.asTensor(filter.type))
+        args.append(tf.constant(filter.radius))
+    return args
+
+class Context: pass
+
+def forward(width,
+            height,
+            num_samples_x,
+            num_samples_y,
+            seed,
+            *args):
+    """
+        Forward rendering pass: given a serialized scene and output an image.
+    """
+    # Unpack arguments
+    with tf.device('/device:cpu:' + str(pydiffvg.get_cpu_device_id())):
+        current_index = 0
+        canvas_width = int(args[current_index])
+        current_index += 1
+        canvas_height = int(args[current_index])
+        current_index += 1
+        num_shapes = int(args[current_index])
+        current_index += 1
+        num_shape_groups = int(args[current_index])
+        current_index += 1
+        output_type = OutputType(int(args[current_index]))
+        current_index += 1
+        use_prefiltering = bool(args[current_index])
+        current_index += 1
+        shapes = []
+        shape_groups = []
+        shape_contents = [] # Important to avoid GC deleting the shapes
+        color_contents = [] # Same as above
+        for shape_id in range(num_shapes):
+            shape_type = ShapeType.asShapeType(args[current_index])
+            current_index += 1
+            if shape_type == diffvg.ShapeType.circle:
+                radius = args[current_index]
+                current_index += 1
+                center = args[current_index]
+                current_index += 1
+                shape = diffvg.Circle(float(radius),
+                                      diffvg.Vector2f(float(center[0]), float(center[1])))
+            elif shape_type == diffvg.ShapeType.ellipse:
+                radius = args[current_index]
+                current_index += 1
+                center = args[current_index]
+                current_index += 1
+                shape = diffvg.Ellipse(diffvg.Vector2f(float(radius[0]), float(radius[1])),
+                                       diffvg.Vector2f(float(center[0]), float(center[1])))
+            elif shape_type == diffvg.ShapeType.path:
+                num_control_points = args[current_index]
+                current_index += 1
+                points = args[current_index]
+                current_index += 1
+                is_closed = args[current_index]
+                current_index += 1
+                shape = diffvg.Path(diffvg.int_ptr(pydiffvg.data_ptr(num_control_points)),
+                                    diffvg.float_ptr(pydiffvg.data_ptr(points)),
+                                    num_control_points.shape[0],
+                                    points.shape[0],
+                                    is_closed)
+            elif shape_type == diffvg.ShapeType.rect:
+                p_min = args[current_index]
+                current_index += 1
+                p_max = args[current_index]
+                current_index += 1
+                shape = diffvg.Rect(diffvg.Vector2f(float(p_min[0]), float(p_min[1])),
+                                    diffvg.Vector2f(float(p_max[0]), float(p_max[1])))
+            else:
+                assert(False)
+            stroke_width = args[current_index]
+            current_index += 1
+            shapes.append(diffvg.Shape(\
+                shape_type, shape.get_ptr(), float(stroke_width)))
+            shape_contents.append(shape)
+
+        for shape_group_id in range(num_shape_groups):
+            shape_ids = args[current_index]
+            current_index += 1
+            fill_color_type = ColorType.asColorType(args[current_index])
+            current_index += 1
+            if fill_color_type == diffvg.ColorType.constant:
+                color = args[current_index]
+                current_index += 1
+                fill_color = diffvg.Constant(\
+                    diffvg.Vector4f(color[0], color[1], color[2], color[3]))
+            elif fill_color_type == diffvg.ColorType.linear_gradient:
+                beg = args[current_index]
+                current_index += 1
+                end = args[current_index]
+                current_index += 1
+                offsets = args[current_index]
+                current_index += 1
+                stop_colors = args[current_index]
+                current_index += 1
+                assert(offsets.shape[0] == stop_colors.shape[0])
+                fill_color = diffvg.LinearGradient(diffvg.Vector2f(float(beg[0]), float(beg[1])),
+                                                   diffvg.Vector2f(float(end[0]), float(end[1])),
+                                                   offsets.shape[0],
+                                                   diffvg.float_ptr(pydiffvg.data_ptr(offsets)),
+                                                   diffvg.float_ptr(pydiffvg.data_ptr(stop_colors)))
+            elif fill_color_type == diffvg.ColorType.radial_gradient:
+                center = args[current_index]
+                current_index += 1
+                radius = args[current_index]
+                current_index += 1
+                offsets = args[current_index]
+                current_index += 1
+                stop_colors = args[current_index]
+                current_index += 1
+                assert(offsets.shape[0] == stop_colors.shape[0])
+                fill_color = diffvg.RadialGradient(diffvg.Vector2f(float(center[0]), float(center[1])),
+                                                   diffvg.Vector2f(float(radius[0]), float(radius[1])),
+                                                   offsets.shape[0],
+                                                   diffvg.float_ptr(pydiffvg.data_ptr(offsets)),
+                                                   diffvg.float_ptr(pydiffvg.data_ptr(stop_colors)))
+            elif fill_color_type is None:
+                fill_color = None
+            else:
+                assert(False)
+
+            stroke_color_type = ColorType.asColorType(args[current_index])
+            current_index += 1
+            if stroke_color_type == diffvg.ColorType.constant:
+                color = args[current_index]
+                current_index += 1
+                stroke_color = diffvg.Constant(\
+                    diffvg.Vector4f(float(color[0]),
+                                    float(color[1]),
+                                    float(color[2]),
+                                    float(color[3])))
+            elif stroke_color_type == diffvg.ColorType.linear_gradient:
+                beg = args[current_index]
+                current_index += 1
+                end = args[current_index]
+                current_index += 1
+                offsets = args[current_index]
+                current_index += 1
+                stop_colors = args[current_index]
+                current_index += 1
+                assert(offsets.shape[0] == stop_colors.shape[0])
+                stroke_color = diffvg.LinearGradient(\
+                    diffvg.Vector2f(float(beg[0]), float(beg[1])),
+                    diffvg.Vector2f(float(end[0]), float(end[1])),
+                    offsets.shape[0],
+                    diffvg.float_ptr(pydiffvg.data_ptr(offsets)),
+                    diffvg.float_ptr(stop_colors.data_ptr()))
+            elif stroke_color_type == diffvg.ColorType.radial_gradient:
+                center = args[current_index]
+                current_index += 1
+                radius = args[current_index]
+                current_index += 1
+                offsets = args[current_index]
+                current_index += 1
+                stop_colors = args[current_index]
+                current_index += 1
+                assert(offsets.shape[0] == stop_colors.shape[0])
+                stroke_color = diffvg.RadialGradient(\
+                    diffvg.Vector2f(float(center[0]), float(center[1])),
+                    diffvg.Vector2f(float(radius[0]), float(radius[1])),
+                    offsets.shape[0],
+                    diffvg.float_ptr(pydiffvg.data_ptr(offsets)),
+                    diffvg.float_ptr(pydiffvg.data_ptr(stop_colors)))
+            elif stroke_color_type is None:
+                stroke_color = None
+            else:
+                assert(False)
+            use_even_odd_rule = bool(args[current_index])
+            current_index += 1
+            shape_to_canvas = args[current_index]
+            current_index += 1
+
+            if fill_color is not None:
+                color_contents.append(fill_color)
+            if stroke_color is not None:
+                color_contents.append(stroke_color)
+            shape_groups.append(diffvg.ShapeGroup(\
+                diffvg.int_ptr(pydiffvg.data_ptr(shape_ids)),
+                shape_ids.shape[0],
+                diffvg.ColorType.constant if fill_color_type is None else fill_color_type,
+                diffvg.void_ptr(0) if fill_color is None else fill_color.get_ptr(),
+                diffvg.ColorType.constant if stroke_color_type is None else stroke_color_type,
+                diffvg.void_ptr(0) if stroke_color is None else stroke_color.get_ptr(),
+                use_even_odd_rule,
+                diffvg.float_ptr(pydiffvg.data_ptr(shape_to_canvas))))
+
+        filter_type = FilterType.asFilterType(args[current_index])
+        current_index += 1
+        filter_radius = args[current_index]
+        current_index += 1
+        filt = diffvg.Filter(filter_type, filter_radius)
+
+    device_name = pydiffvg.get_device_name()
+    device_spec = tf.DeviceSpec.from_string(device_name)
+    use_gpu = device_spec.device_type == 'GPU'
+    gpu_index = device_spec.device_index if device_spec.device_index is not None else 0
+
+    start = time.time()
+    scene = diffvg.Scene(canvas_width,
+                         canvas_height,
+                         shapes,
+                         shape_groups,
+                         filt,
+                         use_gpu,
+                         gpu_index)
+    time_elapsed = time.time() - start
+    global print_timing
+    if print_timing:
+        print('Scene construction, time: %.5f s' % time_elapsed)
+
+    with tf.device(device_name):
+        if output_type == OutputType.color:
+            rendered_image = tf.zeros((int(height), int(width), 4), dtype = tf.float32)
+        else:
+            assert(output_type == OutputType.sdf)
+            rendered_image = tf.zeros((int(height), int(width), 1), dtype = tf.float32)
+
+        start = time.time()
+        diffvg.render(scene,
+                      diffvg.float_ptr(pydiffvg.data_ptr(rendered_image) if output_type == OutputType.color else 0),
+                      diffvg.float_ptr(pydiffvg.data_ptr(rendered_image) if output_type == OutputType.sdf else 0),
+                      width,
+                      height,
+                      int(num_samples_x),
+                      int(num_samples_y),
+                      seed,
+                      diffvg.float_ptr(0), # d_render_image
+                      diffvg.float_ptr(0), # d_render_sdf
+                      diffvg.float_ptr(0), # d_translation
+                      use_prefiltering)
+        time_elapsed = time.time() - start
+        if print_timing:
+            print('Forward pass, time: %.5f s' % time_elapsed)
+
+    ctx = Context()
+    ctx.scene = scene
+    ctx.shape_contents = shape_contents
+    ctx.color_contents = color_contents
+    ctx.filter = filt
+    ctx.width = width
+    ctx.height = height
+    ctx.num_samples_x = num_samples_x
+    ctx.num_samples_y = num_samples_y
+    ctx.seed = seed
+    ctx.output_type = output_type
+    ctx.use_prefiltering = use_prefiltering
+    return rendered_image, ctx
+
+@tf.custom_gradient
+def render(*x):
+    """
+        The main TensorFlow interface of C++ diffvg.
+    """
+    assert(tf.executing_eagerly())
+    if pydiffvg.get_use_gpu() and os.environ['TF_FORCE_GPU_ALLOW_GROWTH'] != 'true':
+        print('******************** WARNING ********************')
+        print('Tensorflow by default allocates all GPU memory,')
+        print('causing huge amount of page faults when rendering.')
+        print('Please set the environment variable TF_FORCE_GPU_ALLOW_GROWTH to true,')
+        print('so that Tensorflow allocates memory on demand.')
+        print('*************************************************')
+
+    width = x[0]
+    height = x[1]
+    num_samples_x = x[2]
+    num_samples_y = x[3]
+    seed = x[4]
+    args = x[5:]
+    img, ctx = forward(width, height, num_samples_x, num_samples_y, seed, *args)
+
+    def backward(grad_img):
+        scene = ctx.scene
+        width = ctx.width
+        height = ctx.height
+        num_samples_x = ctx.num_samples_x
+        num_samples_y = ctx.num_samples_y
+        seed = ctx.seed
+        output_type = ctx.output_type
+        use_prefiltering = ctx.use_prefiltering
+
+        start = time.time()
+        with tf.device(pydiffvg.get_device_name()):
+            diffvg.render(scene,
+                          diffvg.float_ptr(0), # render_image
+                          diffvg.float_ptr(0), # render_sdf
+                          width,
+                          height,
+                          num_samples_x,
+                          num_samples_y,
+                          seed,
+                          diffvg.float_ptr(pydiffvg.data_ptr(grad_img) if output_type == OutputType.color else 0),
+                          diffvg.float_ptr(pydiffvg.data_ptr(grad_img) if output_type == OutputType.sdf else 0),
+                          diffvg.float_ptr(0), # d_translation
+                          use_prefiltering)
+        time_elapsed = time.time() - start
+        global print_timing
+        if print_timing:
+            print('Backward pass, time: %.5f s' % time_elapsed)
+
+        with tf.device('/device:cpu:' + str(pydiffvg.get_cpu_device_id())):
+            d_args = []
+            d_args.append(None) # width
+            d_args.append(None) # height
+            d_args.append(None) # num_samples_x
+            d_args.append(None) # num_samples_y
+            d_args.append(None) # seed
+            d_args.append(None) # canvas_width
+            d_args.append(None) # canvas_height
+            d_args.append(None) # num_shapes
+            d_args.append(None) # num_shape_groups
+            d_args.append(None) # output_type
+            d_args.append(None) # use_prefiltering
+            for shape_id in range(scene.num_shapes):
+                d_args.append(None) # type
+                d_shape = scene.get_d_shape(shape_id)
+                if d_shape.type == diffvg.ShapeType.circle:
+                    d_circle = d_shape.as_circle()
+                    radius = tf.constant(d_circle.radius)
+                    d_args.append(radius)
+                    c = d_circle.center
+                    c = tf.constant((c.x, c.y))
+                    d_args.append(c)
+                elif d_shape.type == diffvg.ShapeType.ellipse:
+                    d_ellipse = d_shape.as_ellipse()
+                    r = d_ellipse.radius
+                    r = tf.constant((d_ellipse.radius.x, d_ellipse.radius.y))
+                    d_args.append(r)
+                    c = d_ellipse.center
+                    c = tf.constant((c.x, c.y))
+                    d_args.append(c)
+                elif d_shape.type == diffvg.ShapeType.path:
+                    d_path = d_shape.as_path()
+                    points = tf.zeros((d_path.num_points, 2), dtype=tf.float32)
+                    d_path.copy_to(diffvg.float_ptr(points.data_ptr()))
+                    d_args.append(None) # num_control_points
+                    d_args.append(points)
+                    d_args.append(None) # is_closed
+                elif d_shape.type == diffvg.ShapeType.rect:
+                    d_rect = d_shape.as_rect()
+                    p_min = tf.constant((d_rect.p_min.x, d_rect.p_min.y))
+                    p_max = tf.constant((d_rect.p_max.x, d_rect.p_max.y))
+                    d_args.append(p_min)
+                    d_args.append(p_max)
+                else:
+                    assert(False)
+                w = tf.constant((d_shape.stroke_width))
+                d_args.append(w)
+
+            for group_id in range(scene.num_shape_groups):
+                d_shape_group = scene.get_d_shape_group(group_id)
+                d_args.append(None) # shape_ids
+                d_args.append(None) # fill_color_type
+                if d_shape_group.has_fill_color():
+                    if d_shape_group.fill_color_type == diffvg.ColorType.constant:
+                        d_constant = d_shape_group.fill_color_as_constant()
+                        c = d_constant.color
+                        d_args.append(tf.constant((c.x, c.y, c.z, c.w)))
+                    elif d_shape_group.fill_color_type == diffvg.ColorType.linear_gradient:
+                        d_linear_gradient = d_shape_group.fill_color_as_linear_gradient()
+                        beg = d_linear_gradient.begin
+                        d_args.append(tf.constant((beg.x, beg.y)))
+                        end = d_linear_gradient.end
+                        d_args.append(tf.constant((end.x, end.y)))
+                        offsets = tf.zeros((d_linear_gradient.num_stops), dtype=tf.float32)
+                        stop_colors = tf.zeros((d_linear_gradient.num_stops, 4), dtype=tf.float32)
+                        # HACK: tensorflow's eager mode uses a cache to store scalar
+                        #       constants to avoid memory copy. If we pass scalar tensors
+                        #       into the C++ code and modify them, we would corrupt the
+                        #       cache, causing incorrect result in future scalar constant
+                        #       creations. Thus we force tensorflow to copy by plusing a zero.
+                        # (also see https://github.com/tensorflow/tensorflow/issues/11186
+                        #  for more discussion regarding copying tensors)
+                        if offsets.shape.num_elements() == 1:
+                            offsets = offsets + 0
+                        d_linear_gradient.copy_to(\
+                            diffvg.float_ptr(pydiffvg.data_ptr(offsets)),
+                            diffvg.float_ptr(pydiffvg.data_ptr(stop_colors)))
+                        d_args.append(offsets)
+                        d_args.append(stop_colors)
+                    elif d_shape_group.fill_color_type == diffvg.ColorType.radial_gradient:
+                        d_radial_gradient = d_shape_group.fill_color_as_radial_gradient()
+                        center = d_radial_gradient.center
+                        d_args.append(tf.constant((center.x, center.y)))
+                        radius = d_radial_gradient.radius
+                        d_args.append(tf.constant((radius.x, radius.y)))
+                        offsets = tf.zeros((d_radial_gradient.num_stops))
+                        if offsets.shape.num_elements() == 1:
+                            offsets = offsets + 0
+                        stop_colors = tf.zeros((d_radial_gradient.num_stops, 4))
+                        d_radial_gradient.copy_to(\
+                            diffvg.float_ptr(pydiffvg.data_ptr(offsets)),
+                            diffvg.float_ptr(pydiffvg.data_ptr(stop_colors)))
+                        d_args.append(offsets)
+                        d_args.append(stop_colors)
+                    else:
+                        assert(False)
+                d_args.append(None) # stroke_color_type
+                if d_shape_group.has_stroke_color():
+                    if d_shape_group.stroke_color_type == diffvg.ColorType.constant:
+                        d_constant = d_shape_group.stroke_color_as_constant()
+                        c = d_constant.color
+                        d_args.append(tf.constant((c.x, c.y, c.z, c.w)))
+                    elif d_shape_group.stroke_color_type == diffvg.ColorType.linear_gradient:
+                        d_linear_gradient = d_shape_group.stroke_color_as_linear_gradient()
+                        beg = d_linear_gradient.begin
+                        d_args.append(tf.constant((beg.x, beg.y)))
+                        end = d_linear_gradient.end
+                        d_args.append(tf.constant((end.x, end.y)))
+                        offsets = tf.zeros((d_linear_gradient.num_stops))
+                        stop_colors = tf.zeros((d_linear_gradient.num_stops, 4))
+                        if offsets.shape.num_elements() == 1:
+                            offsets = offsets + 0
+                        d_linear_gradient.copy_to(\
+                            diffvg.float_ptr(pydiffvg.data_ptr(offsets)),
+                            diffvg.float_ptr(pydiffvg.data_ptr(stop_colors)))
+                        d_args.append(offsets)
+                        d_args.append(stop_colors)
+                    elif d_shape_group.fill_color_type == diffvg.ColorType.radial_gradient:
+                        d_radial_gradient = d_shape_group.stroke_color_as_radial_gradient()
+                        center = d_radial_gradient.center
+                        d_args.append(tf.constant((center.x, center.y)))
+                        radius = d_radial_gradient.radius
+                        d_args.append(tf.constant((radius.x, radius.y)))
+                        offsets = tf.zeros((d_radial_gradient.num_stops))
+                        stop_colors = tf.zeros((d_radial_gradient.num_stops, 4))
+                        if offsets.shape.num_elements() == 1:
+                            offsets = offsets + 0
+                        d_radial_gradient.copy_to(\
+                            diffvg.float_ptr(pydiffvg.data_ptr(offsets)),
+                            diffvg.float_ptr(pydiffvg.data_ptr(stop_colors)))
+                        d_args.append(offsets)
+                        d_args.append(stop_colors)
+                    else:
+                        assert(False)
+                d_args.append(None) # use_even_odd_rule
+                d_shape_to_canvas = tf.zeros((3, 3), dtype = tf.float32)
+                d_shape_group.copy_to(diffvg.float_ptr(pydiffvg.data_ptr(d_shape_to_canvas)))
+                d_args.append(d_shape_to_canvas)
+            d_args.append(None) # filter_type
+            d_args.append(tf.constant(scene.get_d_filter_radius()))
+
+        return d_args
+
+    return img, backward

pydiffvg_tensorflow/shape.py

@@ -0,0 +1,53 @@
+import tensorflow as tf
+import math
+
+class Circle:
+    def __init__(self, radius, center, stroke_width = tf.constant(1.0), id = ''):
+        self.radius = radius
+        self.center = center
+        self.stroke_width = stroke_width
+        self.id = id
+
+class Ellipse:
+    def __init__(self, radius, center, stroke_width = tf.constant(1.0), id = ''):
+        self.radius = radius
+        self.center = center
+        self.stroke_width = stroke_width
+        self.id = id
+
+class Path:
+    def __init__(self, num_control_points, points, is_closed, stroke_width = tf.constant(1.0), id = ''):
+        self.num_control_points = num_control_points
+        self.points = points
+        self.is_closed = is_closed
+        self.stroke_width = stroke_width
+        self.id = id
+
+class Polygon:
+    def __init__(self, points, is_closed, stroke_width = tf.constant(1.0), id = ''):
+        self.points = points
+        self.is_closed = is_closed
+        self.stroke_width = stroke_width
+        self.id = id
+
+class Rect:
+    def __init__(self, p_min, p_max, stroke_width = tf.constant(1.0), id = ''):
+        self.p_min = p_min
+        self.p_max = p_max
+        self.stroke_width = stroke_width
+        self.id = id
+
+class ShapeGroup:
+    def __init__(self,
+                 shape_ids,
+                 fill_color,
+                 use_even_odd_rule = True,
+                 stroke_color = None,
+                 shape_to_canvas = tf.eye(3),
+                 id = ''):
+        self.shape_ids = shape_ids
+        self.fill_color = fill_color
+        self.use_even_odd_rule = use_even_odd_rule
+        self.stroke_color = stroke_color
+        self.shape_to_canvas = shape_to_canvas
+        self.id = id