From 3498584e06f7681bdd6b03f4a806a16db701bb3d Mon Sep 17 00:00:00 2001 From: Wenzel Jakob Date: Tue, 3 Mar 2020 10:07:26 +0100 Subject: [PATCH] environment map optimization support --- .../10_inverse_rendering/invert_bunny.py | 68 +++++++++++++++++++ include/mitsuba/render/bsdf.h | 2 +- include/mitsuba/render/endpoint.h | 4 ++ resources/data | 2 +- src/emitters/envmap.cpp | 34 +++++++++- src/librender/bsdf.cpp | 2 + src/librender/endpoint.cpp | 4 +- src/python/python/autodiff.py | 9 ++- 8 files changed, 118 insertions(+), 7 deletions(-) create mode 100644 docs/examples/10_inverse_rendering/invert_bunny.py diff --git a/docs/examples/10_inverse_rendering/invert_bunny.py b/docs/examples/10_inverse_rendering/invert_bunny.py new file mode 100644 index 000000000..968902dd7 --- /dev/null +++ b/docs/examples/10_inverse_rendering/invert_bunny.py @@ -0,0 +1,68 @@ +# Simple inverse rendering example: render a cornell box reference image, then +# then replace one of the scene parameters and try to recover it using +# differentiable rendering and gradient-based optimization. + +import enoki as ek +import mitsuba +mitsuba.set_variant('gpu_autodiff_rgb') + +from mitsuba.core import Float +from mitsuba.core.xml import load_file +from mitsuba.python.util import traverse +from mitsuba.python.autodiff import render, write_bitmap, Adam +import time + +# Load example scene +scene = load_file('bunny.xml') + +# Find differentiable scene parameters +params = traverse(scene) + +# Make a backup copy +param_res = params['my_envmap.resolution'] +param_ref = Float(params['my_envmap.data']) + +# Discard all parameters except for one we want to differentiate +params.keep(['my_envmap.data']) + +# Render a reference image (no derivatives used yet) +image_ref = render(scene, spp=1) +crop_size = scene.sensors()[0].film().crop_size() +write_bitmap('out_ref.png', image_ref, crop_size) + +# Change to a uniform white lighting environment +params['my_envmap.data'] = ek.full(Float, 1.0, len(params['my_envmap.data'])) +params.update() + +# Construct an Adam optimizer that will adjust the parameters 'params' +# opt = Adam(params, lr=.02) +opt = Adam(params, lr=.02) +print(opt) + +time_a = time.time() + +iterations = 100 +for it in range(iterations): + # Perform a differentiable rendering of the scene + image = render(scene, optimizer=opt, unbiased=True, spp=1) + write_bitmap('out_%03i.png' % it, image, crop_size) + write_bitmap('envmap_%03i.png' % it, params['my_envmap.data'], + (param_res[1], param_res[0])) + + # 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 + opt.step() + + # Compare iterate against ground-truth value + err_ref = ek.hsum(ek.sqr(param_ref - params['my_envmap.data'])) + 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)) diff --git a/include/mitsuba/render/bsdf.h b/include/mitsuba/render/bsdf.h index 6da92da6d..b1569f03b 100644 --- a/include/mitsuba/render/bsdf.h +++ b/include/mitsuba/render/bsdf.h @@ -432,7 +432,7 @@ class MTS_EXPORT_RENDER BSDF : public Object { } /// Return a string identifier - std::string id() const override { return m_id; } + std::string id() const override; /// Return a human-readable representation of the BSDF std::string to_string() const override = 0; diff --git a/include/mitsuba/render/endpoint.h b/include/mitsuba/render/endpoint.h index c60a52963..e190acf61 100644 --- a/include/mitsuba/render/endpoint.h +++ b/include/mitsuba/render/endpoint.h @@ -216,6 +216,9 @@ class MTS_EXPORT_RENDER Endpoint : public Object { */ virtual void set_scene(const Scene *scene); + /// Return a string identifier + std::string id() const override; + //! @} // ============================================================= @@ -233,6 +236,7 @@ class MTS_EXPORT_RENDER Endpoint : public Object { Shape *m_shape = nullptr; bool m_needs_sample_2 = true; bool m_needs_sample_3 = true; + std::string m_id; }; MTS_EXTERN_CLASS_RENDER(Endpoint) diff --git a/resources/data b/resources/data index 3b832036d..ea66366b5 160000 --- a/resources/data +++ b/resources/data @@ -1 +1 @@ -Subproject commit 3b832036d9006b059a76c94c1ea02b42612a8e36 +Subproject commit ea66366b5e313298a6971894a8244e17c42f5e68 diff --git a/src/emitters/envmap.cpp b/src/emitters/envmap.cpp index 21b3ff84e..ff428bb57 100644 --- a/src/emitters/envmap.cpp +++ b/src/emitters/envmap.cpp @@ -109,7 +109,7 @@ class EnvironmentMapEmitter final : public Emitter { } *lum_ptr++ = lum * sin_theta; - store_unaligned(ptr, coeff); + store(ptr, coeff); ptr += 4; } } @@ -124,7 +124,37 @@ class EnvironmentMapEmitter final : public Emitter { } void parameters_changed() override { - // TODO update warp for better importance sampling when data has changed + m_data.managed(); + + std::unique_ptr luminance(new ScalarFloat[hprod(m_resolution)]); + + ScalarFloat *ptr = (ScalarFloat *) m_data.data(), + *lum_ptr = (ScalarFloat *) luminance.get(); + + + for (size_t y = 0; y < m_resolution.y(); ++y) { + ScalarFloat sin_theta = + std::sin(y / ScalarFloat(m_resolution.y() - 1) * math::Pi); + + for (size_t x = 0; x < m_resolution.x(); ++x) { + ScalarVector4f coeff = load(ptr); + ScalarFloat lum; + + if constexpr (is_monochromatic_v) { + lum = coeff.x(); + } else if constexpr (is_rgb_v) { + lum = mitsuba::luminance(ScalarColor3f(head<3>(coeff))); + } else { + static_assert(is_spectral_v); + lum = srgb_model_mean(head<3>(coeff)) * coeff.w(); + } + + *lum_ptr++ = lum * sin_theta; + ptr += 4; + } + } + + m_warp = Warp(luminance.get(), m_resolution); } void set_scene(const Scene *scene) override { diff --git a/src/librender/bsdf.cpp b/src/librender/bsdf.cpp index 309275101..bda8face5 100644 --- a/src/librender/bsdf.cpp +++ b/src/librender/bsdf.cpp @@ -13,6 +13,8 @@ MTS_VARIANT Spectrum BSDF::eval_null_transmission( return 0.f; } +MTS_VARIANT std::string BSDF::id() const { return m_id; } + template std::string type_mask_to_string(Index type_mask) { std::ostringstream oss; diff --git a/src/librender/endpoint.cpp b/src/librender/endpoint.cpp index 1e1228931..5b79e9f54 100644 --- a/src/librender/endpoint.cpp +++ b/src/librender/endpoint.cpp @@ -5,7 +5,7 @@ NAMESPACE_BEGIN(mitsuba) -MTS_VARIANT Endpoint::Endpoint(const Properties &props) { +MTS_VARIANT Endpoint::Endpoint(const Properties &props) : m_id(props.id()) { m_world_transform = props.animated_transform("to_world", ScalarTransform4f()).get(); } @@ -48,6 +48,8 @@ MTS_VARIANT Spectrum Endpoint::eval(const SurfaceInteraction3f NotImplementedError("eval"); } +MTS_VARIANT std::string Endpoint::id() const { return m_id; } + MTS_IMPLEMENT_CLASS_VARIANT(Endpoint, Object) MTS_INSTANTIATE_CLASS(Endpoint) NAMESPACE_END(mitsuba) diff --git a/src/python/python/autodiff.py b/src/python/python/autodiff.py index 22ba34b2d..c00e36b7d 100644 --- a/src/python/python/autodiff.py +++ b/src/python/python/autodiff.py @@ -91,7 +91,7 @@ def _render_helper(scene, spp=None, sensor_index=0): return values / (weight + 1e-8) -def write_bitmap(filename, data, resolution): +def write_bitmap(filename, data, resolution, write_async=True): """ Write the linearized RGB image in `data` to a PNG/EXR/.. file with resolution `resolution`. @@ -110,7 +110,12 @@ def write_bitmap(filename, data, resolution): filename.endswith('.jpeg'): bitmap = bitmap.convert(Bitmap.PixelFormat.RGB, Struct.Type.UInt8, True) - bitmap.write_async(filename, quality=0 if filename.endswith('png') else -1) + quality = 0 if filename.endswith('png') else -1 + + if write_async: + bitmap.write_async(filename, quality=quality) + else: + bitmap.write(filename, quality=quality) def render(scene,