Add pathreparam and smootharea emitter

Co-authored-by: Sebastien Speierer <[email protected]>

Author: loubetg

CMakeLists.txt

@@ -572,6 +572,7 @@ macro (add_file DST SRC)
 endmacro()
 
 add_file(data/srgb.coeff ${CMAKE_BINARY_DIR}/ext_build/rgb2spec/srgb.coeff rgb2spec_opt_run)
+add_file(data/vmf-hemisphere.data ${CMAKE_CURRENT_SOURCE_DIR}/src/integrators/vmf-hemisphere.data)
 
 file(COPY resources/data/ior DESTINATION ${CMAKE_CURRENT_BINARY_DIR}/dist/data)
 

docs/examples/10_inverse_rendering/invert_heightfield.py

@@ -0,0 +1,111 @@
+import os
+import time
+import enoki as ek
+
+import mitsuba
+mitsuba.set_variant('gpu_autodiff_rgb')
+
+from mitsuba.core import UInt32, Float, Thread, xml, Vector2f, Vector3f, Transform4f, ScalarTransform4f
+from mitsuba.render import SurfaceInteraction3f
+from mitsuba.python.util import traverse
+from mitsuba.python.autodiff import render, write_bitmap, Adam, SGD
+
+# Return contiguous flattened array (will be included in next enoki release)
+def ravel(buf, dim = 3):
+    idx = dim * UInt32.arange(int(len(buf) / dim))
+    if dim == 2:
+        return Vector2f(ek.gather(buf, idx), ek.gather(buf, idx + 1))
+    elif dim == 3:
+        return Vector3f(ek.gather(buf, idx), ek.gather(buf, idx + 1), ek.gather(buf, idx + 2))
+
+# Convert flat array into a vector of arrays (will be included in next enoki release)
+def unravel(source, target, dim = 3):
+    idx = UInt32.arange(ek.slices(source))
+    for i in range(dim):
+        ek.scatter(target, source[i], dim * idx + i)
+
+# Prepare output folder
+output_path = "output/invert_heightfield/"
+if not os.path.isdir(output_path):
+    os.makedirs(output_path)
+
+# Load example scene
+scene_folder = '../../../resources/data/docs/examples/invert_heightfield/'
+Thread.thread().file_resolver().append(scene_folder)
+scene = xml.load_file(scene_folder + 'heightfield.xml')
+
+params = traverse(scene)
+positions_buf = params['grid_mesh.vertex_positions_buf']
+positions_initial = ravel(positions_buf)
+normals_initial = ravel(params['grid_mesh.vertex_normals_buf'])
+vertex_count = ek.slices(positions_initial)
+
+# Create a texture with the reference displacement map
+disp_tex = xml.load_dict({
+    "type" : "bitmap",
+    "filename" : "mitsuba_coin.jpg",
+    "to_uv" : ScalarTransform4f.scale([1, -1, 1]) # texture is upside-down
+}).expand()[0]
+
+# Create a fake surface interaction with an entry per vertex on the mesh
+mesh_si = SurfaceInteraction3f.zero(vertex_count)
+mesh_si.uv = ravel(params['grid_mesh.vertex_texcoords_buf'], dim=2)
+
+# Evaluate the displacement map for the entire mesh
+disp_tex_data_ref = disp_tex.eval_1(mesh_si)
+
+# Apply displacement to mesh vertex positions and update scene (e.g. OptiX BVH)
+def apply_displacement(amplitude = 0.05):
+    new_positions = disp_tex.eval_1(mesh_si) * normals_initial * amplitude + positions_initial
+    unravel(new_positions, positions_buf)
+    params['grid_mesh.vertex_positions_buf'] = positions_buf
+    params.update()
+
+# Apply displacement before generating reference image
+apply_displacement()
+
+# Render a reference image (no derivatives used yet)
+image_ref = render(scene, spp=32)
+crop_size = scene.sensors()[0].film().crop_size()
+write_bitmap(output_path + 'out_ref.exr', image_ref, crop_size)
+print("Write " + output_path + "out_ref.exr")
+
+# Reset texture data to a constant
+disp_tex_params = traverse(disp_tex)
+disp_tex_params.keep(['data'])
+disp_tex_params['data'] = ek.full(Float, 0.25, len(disp_tex_params['data']))
+disp_tex_params.update()
+
+# Construct an Adam optimizer that will adjust the texture parameters
+opt = Adam(disp_tex_params, lr=0.005)
+
+time_a = time.time()
+
+iterations = 100
+for it in range(iterations):
+    # Perform a differentiable rendering of the scene
+    image = render(scene,
+                   optimizer=opt,
+                   spp=4,
+                   unbiased=True,
+                   pre_render_callback=apply_displacement)
+
+    write_bitmap(output_path + 'out_%03i.exr' % it, image, crop_size)
+
+    # Objective: MSE between 'image' and 'image_ref'
+    ob_val = ek.hsum(ek.sqr(image - image_ref)) / len(image)
+
+    # Back-propagate errors to input parameters
+    ek.backward(ob_val)
+
+    # Optimizer: take a gradient step -> update displacement map
+    opt.step()
+
+    # Compare iterate against ground-truth value
+    err_ref = ek.hsum(ek.sqr(disp_tex_data_ref - disp_tex.eval_1(mesh_si)))
+    print('Iteration %03i: error=%g' % (it, err_ref[0]), end='\r')
+
+time_b = time.time()
+
+print()
+print('%f ms per iteration' % (((time_b - time_a) * 1000) / iterations))

docs/src/inverse_rendering/advanced.rst

@@ -126,3 +126,172 @@ The solution we found optimizes the objective well (i.e. the rendered image
 matches the target), but the reconstructed texture may not match our
 expectation. In such a case, it may be advisable to introduce further
 regularization (non-negativity, smoothness, etc.).
+
+.. note::
+
+    The full Python script of this tutorial can be found in the file:
+    :file:`docs/examples/10_inverse_rendering/invert_bunny.py`.
+
+
+Heightfield optimization
+------------------------
+
+This advanced example demonstrates how to optimize a displacement map texture, which implies the
+differentiation of mesh parameters, such as vertex positions. Computing derivatives for parameters
+that affect visibility is a complex problem as it would normally make the integrants of the rendering
+equation non-differentiable. For this reason, this example requires the use of the specialized
+:ref:`pathreparam <integrator-pathreparam>` integrator, described in this
+`article <https://rgl.epfl.ch/publications/Loubet2019Reparameterizing>`_.
+
+The example scene can be found in ``resource/data/docs/examples/invert_heightfield/`` and contains a
+simple grid mesh illuminated by a rectangular light source. To avoid discontinuities around the
+area light, we use the :ref:`smootharea <emitter-smootharea>` plugin.
+
+First, we define two helper functions that we will use to transform the mesh
+parameter buffers (flatten arrays) into ``VectorXf`` type (and the other way around).
+Note that those functions will be natively supported by ``enoki`` in a futur release.
+
+.. code-block:: python
+
+    # Return contiguous flattened array (will be included in next enoki release)
+    def ravel(buf, dim = 3):
+        idx = dim * UInt32.arange(int(len(buf) / dim))
+        if dim == 2:
+            return Vector2f(ek.gather(buf, idx), ek.gather(buf, idx + 1))
+        elif dim == 3:
+            return Vector3f(ek.gather(buf, idx), ek.gather(buf, idx + 1), ek.gather(buf, idx + 2))
+
+    # Convert flat array into a vector of arrays (will be included in next enoki release)
+    def unravel(source, target, dim = 3):
+        idx = UInt32.arange(ek.slices(source))
+        for i in range(dim):
+            ek.scatter(target, source[i], dim * idx + i)
+
+
+Using those, we can now load the scene and read the initial grid mesh parameters (vertex positions, normals and texture coordinates), which we will use
+later in the script.
+
+.. code-block:: python
+
+    import enoki as ek
+    import mitsuba
+    mitsuba.set_variant('gpu_autodiff_rgb')
+
+    from mitsuba.core import UInt32, Float, Thread, xml, Vector2f, Vector3f, Transform4f
+    from mitsuba.render import SurfaceInteraction3f
+    from mitsuba.python.util import traverse
+    from mitsuba.python.autodiff import render, write_bitmap, Adam
+
+    # Load example scene
+    scene_folder = '../../../resources/data/docs/examples/invert_heightfield/'
+    Thread.thread().file_resolver().append(scene_folder)
+    scene = xml.load_file(scene_folder + 'heightfield.xml')
+
+    params = traverse(scene)
+    positions_buf = params['grid_mesh.vertex_positions_buf']
+    positions_initial = ravel(positions_buf)
+    normals_initial = ravel(params['grid_mesh.vertex_normals_buf'])
+    vertex_count = ek.slices(positions_initial)
+
+
+In this example, we implement displacement mapping directly in Python instead of using a C++ plugin.
+This showcases the flexibility of the framework, and the ability to fully control the optimization
+process. For instance, one could want to add constraints on the displacement values range, ...
+
+We first create a :ref:`Bitmap <texture-bitmap>` texture instance using
+:py:func:`mitsuba.core.xml.load_dict`, which will load the displacement map image file from disk.
+We also create a :py:class:`mitsuba.render.SurfaceInteraction3f` with one entry per vertex on the
+mesh. By properly setting the texture coordinates on this surface interaction, we can now evaluate
+the displacement map for the entire mesh in one line of code.
+
+.. code-block:: python
+
+    # Create a texture with the reference displacement map
+    disp_tex = xml.load_dict({
+        "type" : "bitmap",
+        "filename" : "mitsuba_coin.jpg"
+    }).expand()[0]
+
+    # Create a fake surface interaction with an entry per vertex on the mesh
+    mesh_si = SurfaceInteraction3f.zero(vertex_count)
+    mesh_si.uv = ravel(params['grid_mesh.vertex_texcoords_buf'], dim=2)
+
+    # Evaluate the displacement map for the entire mesh
+    disp_tex_data_ref = disp_tex.eval_1(mesh_si)
+
+Finally, we define a function that applies the displacement map onto the original mesh. This will
+be called at every iteration of the optimization loop to update the mesh data everytime the
+displacement map is refined.
+
+.. code-block:: python
+
+    # Apply displacement to mesh vertex positions and call update scene
+    def apply_displacement(amplitude = 0.05):
+        new_positions = disp_tex.eval_1(mesh_si) * normals_initial * amplitude + positions_initial
+        unravel(new_positions, positions_buf)
+        params['grid_mesh.vertex_positions_buf'] = positions_buf
+        params.update()
+
+We can now generate a reference image.
+
+.. code-block:: python
+
+    # Apply displacement before generating reference image
+    apply_displacement()
+
+    # Render a reference image (no derivatives used yet)
+    image_ref = render(scene, spp=32)
+    crop_size = scene.sensors()[0].film().crop_size()
+    write_bitmap('out_ref.exr', image_ref, crop_size)
+    print("Write out_ref.exr")
+
+Before runing the optimization loop, we need to change the displacement data to a constant value
+(here ``0.25``). This can be done using the :py:func:`mitsuba.python.util.traverse` function
+on the texture object directly. We can then create an optimizer that will adjust those texture
+parameters during the optimization process.
+
+.. code-block:: python
+
+    # Reset texture data to a constant
+    disp_tex_params = traverse(disp_tex)
+    disp_tex_params['data'] = ek.full(Float, 0.25, len(disp_tex_params['data']))
+    disp_tex_params.update()
+
+    # Construct an Adam optimizer that will adjust the texture parameters
+    disp_tex_params.keep(['data'])
+    opt = Adam(disp_tex_params, lr=0.005)
+
+The optimization loop is very similar to the previous example, to the exception that it needs to
+manually apply the displacement mapping to the mesh at every iteration.
+
+.. code-block:: python
+
+    iterations = 100
+    for it in range(iterations):
+        # Apply displacement to mesh and update scene (e.g. OptiX BVH)
+        apply_displacement()
+
+        # Perform a differentiable rendering of the scene
+        image = render(scene, optimizer=opt, spp=4)
+        write_bitmap('out_%03i.exr' % it, image, crop_size)
+
+        # Objective: MSE between 'image' and 'image_ref'
+        ob_val = ek.hsum(ek.sqr(image - image_ref)) / len(image)
+
+        # Back-propagate errors to input parameters
+        ek.backward(ob_val)
+
+        # Optimizer: take a gradient step -> update displacement map
+        opt.step()
+
+        # Compare iterate against ground-truth value
+        err_ref = ek.hsum(ek.sqr(disp_tex_data_ref - disp_tex.eval_1(mesh_si)))
+        print('Iteration %03i: error=%g' % (it, err_ref[0]), end='\r')
+
+
+Here we can see the result of the heightfield optimization:
+
+.. note::
+
+    The full Python script of this tutorial can be found in the file:
+    :file:`docs/examples/10_inverse_rendering/invert_heightfield.py`.

docs/src/inverse_rendering/diff_render.rst

@@ -322,7 +322,7 @@ overly large value.
 .. note::
 
     The full Python script of this tutorial can be found in the file:
-    :file:`docs/examples/10_diff_render/invert_cbox.py`.
+    :file:`docs/examples/10_inverse_rendering/invert_cbox.py`.
 
 
 Forward-mode differentiation
@@ -395,4 +395,4 @@ respect to albedo, the red color disappears.
 .. note::
 
     The full Python script of this tutorial can be found in the file:
-    :file:`docs/examples/10_diff_render/forward_diff.py`.
+    :file:`docs/examples/10_inverse_rendering/forward_diff.py`.

include/mitsuba/core/warp.h

@@ -554,9 +554,9 @@ MTS_INLINE Value square_to_beckmann_pdf(const Vector<Value, 3> &m,
 // =======================================================================
 
 /// Warp a uniformly distributed square sample to a von Mises Fisher distribution
-template <typename Value>
+template <typename Value, typename KappaValue = scalar_t<Value>>
 MTS_INLINE Vector<Value, 3> square_to_von_mises_fisher(const Point<Value, 2> &sample,
-                                                       const scalar_t<Value> &kappa) {
+                                                       const KappaValue &kappa) {
 #if 0
     // Approach 1: warping method based on standard disk mapping
 
@@ -595,8 +595,8 @@ MTS_INLINE Vector<Value, 3> square_to_von_mises_fisher(const Point<Value, 2> &sa
 }
 
 /// Inverse of the mapping \ref von_mises_fisher_to_square
-template <typename Value>
-MTS_INLINE Point<Value, 2> von_mises_fisher_to_square(const Vector<Value, 3> &v, scalar_t<Value> kappa) {
+template <typename Value, typename KappaValue = scalar_t<Value>>
+MTS_INLINE Point<Value, 2> von_mises_fisher_to_square(const Vector<Value, 3> &v, KappaValue kappa) {
     Value expm2k = exp(-2.f * kappa),
           t      = exp((v.z() - 1.f) * kappa),
           sy     = (expm2k - t) / (expm2k - 1.f),
@@ -607,17 +607,15 @@ MTS_INLINE Point<Value, 2> von_mises_fisher_to_square(const Vector<Value, 3> &v,
 }
 
 /// Probability density of \ref square_to_von_mises_fisher()
-template <typename Value>
-MTS_INLINE Value square_to_von_mises_fisher_pdf(const Vector<Value, 3> &v, scalar_t<Value> kappa) {
+template <typename Value, typename KappaValue = scalar_t<Value>>
+MTS_INLINE Value square_to_von_mises_fisher_pdf(const Vector<Value, 3> &v, KappaValue kappa) {
     /* Stable algorithm for evaluating the von Mises Fisher distribution
        https://www.mitsuba-renderer.org/~wenzel/files/vmf.pdf */
 
-    assert(kappa >= 0);
-    if (unlikely(kappa == 0))
-        return math::InvFourPi<Value>;
-    else
-        return exp(kappa * (v.z() - 1.f)) * (kappa * math::InvTwoPi<Value>) /
+    Value result = exp(kappa * (v.z() - 1.f)) * (kappa * math::InvTwoPi<Value>) /
                (1.f - exp(-2.f * kappa));
+    masked(result, eq(kappa, 0.f)) = math::InvFourPi<Value>;
+    return result;
 }
 
 // =======================================================================

include/mitsuba/render/bsdf.h

@@ -220,6 +220,9 @@ template <typename Float, typename Spectrum> struct BSDFSample3 {
     /// Stores the component index that was sampled by \ref BSDF::sample()
     UInt32 sampled_component;
 
+    /// Roughness of the sampled material (used in \ref DiffPathIntegrator)
+    Float sampled_roughness;
+
     //! @}
     // =============================================================
 
@@ -242,13 +245,13 @@ template <typename Float, typename Spectrum> struct BSDFSample3 {
      */
     BSDFSample3(const Vector3f &wo)
         : wo(wo), pdf(0.f), eta(1.f), sampled_type(0),
-          sampled_component(uint32_t(-1)) { }
+          sampled_component(uint32_t(-1)), sampled_roughness(math::Infinity<Float>) { }
 
 
     //! @}
     // =============================================================
 
-    ENOKI_STRUCT(BSDFSample3, wo, pdf, eta, sampled_type, sampled_component);
+    ENOKI_STRUCT(BSDFSample3, wo, pdf, eta, sampled_type, sampled_component, sampled_roughness);
 };
 
 
@@ -509,7 +512,7 @@ NAMESPACE_END(mitsuba)
 // -----------------------------------------------------------------------
 
 ENOKI_STRUCT_SUPPORT(mitsuba::BSDFSample3, wo, pdf, eta,
-                     sampled_type, sampled_component)
+                     sampled_type, sampled_component, sampled_roughness)
 
 //! @}
 // -----------------------------------------------------------------------