Fix build and tensorflow 2.4 compatibility

csrc/CMakeLists.txt

@@ -10,7 +10,6 @@ find_library(EGL_LIBRARIES NAMES egl EGL REQUIRED)
 
 # Search for cuda headers (using the form of path that tensorflow includes them with), based on cmake-inferred nvcc, or $CUDA_HOME
 get_filename_component(NVCC_DIR ${CMAKE_CUDA_COMPILER} DIRECTORY)
-find_path(CUDA_INCLUDE_DIR NAMES cuda/include/cuda.h HINTS ${NVCC_DIR}/../.. PATHS ENV CUDA_HOME REQUIRED)
 
 # Ask tensorflow for its compile flags; one should therefore make sure cmake is run with the venv active that the op will be used in
 execute_process(COMMAND python -c "import tensorflow; print(' '.join(tensorflow.sysconfig.get_compile_flags()))" OUTPUT_VARIABLE Tensorflow_DEFAULT_COMPILE_FLAGS OUTPUT_STRIP_TRAILING_WHITESPACE)
@@ -32,13 +31,13 @@ execute_process(COMMAND python -c "import tensorflow; print(' '.join(tensorflow.
 set(Tensorflow_LINK_FLAGS "${Tensorflow_DEFAULT_LINK_FLAGS}" CACHE STRING "Tensorflow linker flags for custom ops")
 
 # in the following, we need ../external/tensorflow for cuda_config.h in tf versions with #16959 unfixed
-include_directories(SYSTEM ../external/tensorflow ${NSYNC_INCLUDE_DIR} ${CUDA_INCLUDE_DIR} ${EGL_INCLUDE_DIR} ${OPENGL_INCLUDE_DIR})
+include_directories(SYSTEM ../external/tensorflow ${NSYNC_INCLUDE_DIR} ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES} ${EGL_INCLUDE_DIR} ${OPENGL_INCLUDE_DIR})
 include_directories(${Tensorflow_INCLUDE_DIRS} ${Tensorflow_INCLUDE_DIRS}/external/nsync/public)
 
 set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -ffast-math")
 
 # in the following, NDEBUG is set to avoid issues with constexpr in absl's string_view.h
-set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -arch=sm_30 --expt-relaxed-constexpr -DNDEBUG")
+set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-relaxed-constexpr -DNDEBUG")
 
 add_library(
     rasterise SHARED

csrc/rasterise_egl.cpp

@@ -16,9 +16,9 @@
 #include "tensorflow/core/platform/stream_executor.h"
 
 // This form of include path matches what the TensorFlow headers use
-#include <cuda/include/cuda.h>
-#include <cuda/include/cuda_runtime_api.h>
-#include <cuda/include/cuda_gl_interop.h>
+#include <cuda.h>
+#include <cuda_runtime_api.h>
+#include <cuda_gl_interop.h>
 
 #define EIGEN_USE_GPU
 #include "unsupported/Eigen/CXX11/src/Tensor/TensorDeviceCuda.h"

csrc/rasterise_grad_egl.cpp

@@ -14,9 +14,9 @@
 #include "tensorflow/core/platform/stream_executor.h"
 
 // This form of include path matches what the TensorFlow headers use
-#include <cuda/include/cuda.h>
-#include <cuda/include/cuda_runtime_api.h>
-#include <cuda/include/cuda_gl_interop.h>
+#include <cuda.h>
+#include <cuda_runtime_api.h>
+#include <cuda_gl_interop.h>
 
 #define EIGEN_USE_GPU
 #include "unsupported/Eigen/CXX11/src/Tensor/TensorDeviceCuda.h"

dirt/lighting.py

@@ -26,8 +26,8 @@ def _get_face_normals(vertices, faces):
     vertices_ndim = vertices.get_shape().ndims
     vertices_by_index = tf.transpose(vertices, [vertices_ndim - 2] + list(range(vertices_ndim - 2)) + [vertices_ndim - 1])  # indexed by vertex-index, *, x/y/z
     vertices_by_face = tf.gather(vertices_by_index, faces)  # indexed by face-index, vertex-in-face, *, x/y/z
-    normals_by_face = tf.cross(vertices_by_face[:, 1] - vertices_by_face[:, 0], vertices_by_face[:, 2] - vertices_by_face[:, 0])  # indexed by face-index, *, x/y/z
-    normals_by_face /= (tf.norm(normals_by_face, axis=-1, keep_dims=True) + 1.e-12)  # ditto
+    normals_by_face = tf.linalg.cross(vertices_by_face[:, 1] - vertices_by_face[:, 0], vertices_by_face[:, 2] - vertices_by_face[:, 0])  # indexed by face-index, *, x/y/z
+    normals_by_face /= (tf.norm(normals_by_face, axis=-1, keepdims=True) + 1.e-12)  # ditto
     return normals_by_face, vertices_by_index
 
 
@@ -85,8 +85,8 @@ def vertex_normals(vertices, faces, name=None):
             dense_shape=tf.cast(tf.concat([[face_count], vbi_shape], axis=0), tf.int64)
         )  # indexed by face-index, vertex-index, *, x/y/z
 
-        summed_normals_by_vertex = tf.sparse_reduce_sum(normals_by_face_and_vertex, axis=0)  # indexed by vertex-index, *, x/y/z
-        renormalised_normals_by_vertex = summed_normals_by_vertex / (tf.norm(summed_normals_by_vertex, axis=-1, keep_dims=True) + 1.e-12)  # ditto
+        summed_normals_by_vertex = tf.compat.v1.sparse_reduce_sum(normals_by_face_and_vertex, axis=0)  # indexed by vertex-index, *, x/y/z
+        renormalised_normals_by_vertex = summed_normals_by_vertex / (tf.norm(summed_normals_by_vertex, axis=-1, keepdims=True) + 1.e-12)  # ditto
 
         result = tf.transpose(renormalised_normals_by_vertex, list(range(1, vertices_ndim - 1)) + [0, vertices_ndim - 1])
         result.set_shape(vertices.get_shape())
@@ -277,8 +277,8 @@ def specular_directional(vertex_positions, vertex_normals, vertex_reflectivities
         reflected_directions = -vertices_to_light_direction + 2. * tf.matmul(vertex_normals, vertices_to_light_direction[..., tf.newaxis]) * vertex_normals  # indexed by *, vertex-index, x/y/z
         vertex_to_camera_displacements = camera_position[..., tf.newaxis, :] - vertex_positions  # indexed by *, vertex-index, x/y/z
         cosines = tf.reduce_sum(
-            (vertex_to_camera_displacements / tf.norm(vertex_to_camera_displacements, axis=-1, keep_dims=True) + 1.e-12) * reflected_directions,
-            axis=-1, keep_dims=True
+            (vertex_to_camera_displacements / tf.norm(vertex_to_camera_displacements, axis=-1, keepdims=True) + 1.e-12) * reflected_directions,
+            axis=-1, keepdims=True
         )  # indexed by *, vertex-index, singleton
         if double_sided:
             cosines = tf.abs(cosines)
@@ -333,7 +333,7 @@ def diffuse_point(vertex_positions, vertex_normals, vertex_colors, light_positio
         light_color = tf.convert_to_tensor(light_color, name='light_color')
 
         relative_positions = vertex_positions - light_position[..., tf.newaxis, :]  # indexed by *, vertex-index, x/y/z
-        incident_directions = relative_positions / (tf.norm(relative_positions, axis=-1, keep_dims=True) + 1.e-12)  # ditto
+        incident_directions = relative_positions / (tf.norm(relative_positions, axis=-1, keepdims=True) + 1.e-12)  # ditto
         cosines = tf.reduce_sum(vertex_normals * incident_directions, axis=-1)  # indexed by *, vertex-index
         if double_sided:
             cosines = tf.abs(cosines)

samples/deferred.py

@@ -6,7 +6,8 @@
 import dirt
 import dirt.matrices as matrices
 import dirt.lighting as lighting
-
+from PIL import Image
+import numpy as np
 
 frame_width, frame_height = 640, 480
 
@@ -76,7 +77,7 @@ def shader_fn(gbuffer, view_matrix, light_direction):
         diffuse_contribution = tf.reshape(diffuse_contribution, [frame_height, frame_width, 3])
 
         # Calculate a white specular (Phong) lighting component
-        camera_position_world = tf.matrix_inverse(view_matrix)[3, :3]
+        camera_position_world = tf.linalg.inv(view_matrix)[3, :3]
         specular_contribution = lighting.specular_directional(
             tf.reshape(positions, [-1, 3]),
             tf.reshape(normals, [-1, 3]),
@@ -116,15 +117,13 @@ def shader_fn(gbuffer, view_matrix, light_direction):
         shader_additional_inputs=[view_matrix, light_direction]
     )
 
-    save_pixels = tf.write_file(
-        'deferred.jpg',
-        tf.image.encode_jpeg(tf.cast(pixels * 255, tf.uint8))
-    )
+    pixels = tf.cast(pixels * 255, tf.uint8)
 
-    session = tf.Session(config=tf.ConfigProto(gpu_options=tf.GPUOptions(allow_growth=True)))
+    session = tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(gpu_options=tf.compat.v1.GPUOptions(allow_growth=True)))
     with session.as_default():
-
-        save_pixels.run()
+        image = pixels
+        img = Image.fromarray( np.asarray(image))
+        img.save("test_def.png")
 
 
 if __name__ == '__main__':

samples/simple.py

@@ -7,6 +7,7 @@
 import dirt
 import dirt.matrices as matrices
 import dirt.lighting as lighting
+tf.compat.v1.disable_eager_execution()
 
 
 frame_width, frame_height = 640, 480
@@ -24,7 +25,7 @@ def build_cube():
 
 
 def unit(vector):
-    return tf.convert_to_tensor(vector) / tf.norm(vector)
+    return tf.convert_to_tensor(value=vector) / tf.norm(tensor=vector)
 
 
 def main():
@@ -73,7 +74,7 @@ def main():
         width=frame_width, height=frame_height, channels=3
     )
 
-    session = tf.Session()
+    session = tf.compat.v1.Session()
     with session.as_default():
 
         pixels_eval = pixels.eval()

samples/textured.py

@@ -4,6 +4,7 @@
 import os
 import numpy as np
 import tensorflow as tf
+from PIL import Image
 
 import dirt
 import dirt.matrices as matrices
@@ -84,9 +85,11 @@ def add_quad(vertices, uvs):
 
     cube_vertices_object = np.asarray(cube_vertices_object, np.float32)
     cube_uvs = np.asarray(cube_uvs, np.float32)
-
+    image  = Image.open(os.path.dirname(__file__) + '/cat.jpg')
+    image.load()
+    data = np.asarray( image, dtype=np.uint8 )
     # Load the texture image
-    texture = tf.cast(tf.image.decode_jpeg(tf.read_file(os.path.dirname(__file__) + '/cat.jpg')), tf.float32) / 255.
+    texture = tf.cast(data, tf.float32) / 255.
 
     # Convert vertices to homogeneous coordinates
     cube_vertices_object = tf.concat([
@@ -121,7 +124,7 @@ def shader_fn(gbuffer, texture, light_direction):
         normals = gbuffer[:, :, 3:]
 
         # Sample the texture at locations corresponding to each pixel; this defines the unlit material color at each point
-        unlit_colors = sample_texture(texture, uvs_to_pixel_indices(uvs, tf.shape(texture)[:2]))
+        unlit_colors = sample_texture(texture, uvs_to_pixel_indices(uvs, tf.shape(input=texture)[:2]))
 
         # Calculate a simple grey ambient lighting component
         ambient_contribution = unlit_colors * [0.4, 0.4, 0.4]
@@ -157,15 +160,14 @@ def shader_fn(gbuffer, texture, light_direction):
         shader_additional_inputs=[texture, light_direction]
     )
 
-    save_pixels = tf.write_file(
-        'textured.jpg',
-        tf.image.encode_jpeg(tf.cast(pixels * 255, tf.uint8))
-    )
+    pixels = tf.cast(pixels * 255, tf.uint8)
 
-    session = tf.Session(config=tf.ConfigProto(gpu_options=tf.GPUOptions(allow_growth=True)))
+    session = tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(gpu_options=tf.compat.v1.GPUOptions(allow_growth=True)))
     with session.as_default():
 
-        save_pixels.run()
+        image = pixels
+        img = Image.fromarray( np.asarray(image))
+        img.save("test_tex.png")
 
 
 if __name__ == '__main__':