Created stub for the optix denoiser

include/mitsuba/render/bsdf.h

@@ -462,6 +462,9 @@ class MI_EXPORT_LIB BSDF : public Object {
     /// Set a string identifier
     void set_id(const std::string& id) override { m_id = id; };
 
+    /// Return the diffuse reflectance value (if any)
+    virtual Spectrum get_diffuse_reflectance(const SurfaceInteraction3f &si, Mask active = true) const;
+
     /// Return a human-readable representation of the BSDF
     std::string to_string() const override = 0;
 
@@ -541,6 +544,7 @@ DRJIT_VCALL_TEMPLATE_BEGIN(mitsuba::BSDF)
     DRJIT_VCALL_METHOD(eval_null_transmission)
     DRJIT_VCALL_METHOD(pdf)
     DRJIT_VCALL_METHOD(eval_pdf)
+    DRJIT_VCALL_METHOD(get_diffuse_reflectance)
     DRJIT_VCALL_GETTER(flags, uint32_t)
     auto needs_differentials() const {
         return has_flag(flags(), mitsuba::BSDFFlags::NeedsDifferentials);

include/mitsuba/render/denoiser.h

@@ -0,0 +1,128 @@
+// #prsagma once
+
+#include <mitsuba/core/bitmap.h>
+#include <mitsuba/core/platform.h>
+#include <mitsuba/render/optix_api.h>
+#include <enoki-jit/optix.h>
+#include <enoki/array.h>
+#include <iostream>
+#include <exception>
+
+#define optix_check(call) if(call != 0) {std::cerr << "Optix call failed" << std::endl; return Float(0.0);}                  
+
+NAMESPACE_BEGIN(mitsuba)
+
+template <typename Float>
+Float denoise(const Bitmap& noisy, Bitmap* albedo, Bitmap* normals) {
+    if constexpr (ek::is_array_v<Float>) {
+        OptixDeviceContext context = jit_optix_context();
+
+        OptixDenoiser denoiser = nullptr;  
+        OptixDenoiserSizes sizes = {};
+        OptixDenoiserParams params = {};
+        OptixDenoiserOptions options = {};
+        OptixDenoiserModelKind modelKind = albedo == nullptr ? OPTIX_DENOISER_MODEL_KIND_HDR : OPTIX_DENOISER_MODEL_KIND_AOV;
+        options.inputKind = albedo == nullptr ? OPTIX_DENOISER_INPUT_RGB : OPTIX_DENOISER_INPUT_RGB_ALBEDO_NORMAL;
+
+        uint32_t scratch_size = 0;
+        uint32_t state_size = 0;
+
+        OptixImage2D inputs[3] = {};
+        OptixImage2D output;
+
+        optix_check(optixDenoiserCreate(context, &options, &denoiser)); 
+        optix_check(optixDenoiserSetModel(denoiser, modelKind, nullptr, 0));
+        optix_check(optixDenoiserComputeMemoryResources(denoiser, noisy.width(), noisy.height(), &sizes));
+        scratch_size = static_cast<uint32_t>(sizes.withoutOverlapScratchSizeInBytes);
+        state_size = static_cast<uint32_t>(sizes.stateSizeInBytes);
+        // const uint64_t frame_size = noisy.pixel_count() * noisy.bytes_per_pixel();
+
+        std::cout << scratch_size << std::endl;
+        std::cout << state_size << std::endl;
+
+        Float intensity_float = ek::zero<Float>(1);
+        Float scratch_float = ek::zero<Float>(scratch_size / sizeof(float));
+        Float state_float = ek::zero<Float>(state_size / sizeof(float));
+        Float output_data = ek::zero<Float>(noisy.pixel_count() * noisy.channel_count());
+        ek::eval(intensity_float, scratch_float, state_float, output_data);
+        ek::sync_thread();
+
+        CUdeviceptr intensity = 0;
+        CUdeviceptr scratch = 0;
+        CUdeviceptr state = 0;
+        intensity = intensity_float.data();
+        scratch = scratch_float.data();
+        state = state_float.data();
+        output.data = output_data.data();
+
+        output.width = noisy.width();
+        output.height = noisy.height();
+        output.format = OPTIX_PIXEL_FORMAT_FLOAT4;
+        output.rowStrideInBytes = noisy.width() * noisy.bytes_per_pixel();
+        output.pixelStrideInBytes = noisy.bytes_per_pixel();
+        DynamicBuffer<Float> noisy_data = ek::load<DynamicBuffer<Float>>(noisy.data(), noisy.pixel_count() * noisy.channel_count());
+        inputs[0].data = noisy_data.data();
+        inputs[0].width = noisy.width();
+        inputs[0].height = noisy.height();
+        inputs[0].format = OPTIX_PIXEL_FORMAT_FLOAT4;
+        inputs[0].rowStrideInBytes = noisy.width() * noisy.bytes_per_pixel();
+        inputs[0].pixelStrideInBytes = noisy.bytes_per_pixel();
+
+        // std::cout << noisy.pixel_format() << std::endl;
+        // std::cout << noisy.has_alpha() << std::endl;
+        // std::cout << noisy.bytes_per_pixel() << std::endl;
+        std::cout << noisy_data << std::endl;
+
+
+        unsigned int nb_channels = 1;
+        if(albedo != nullptr) {
+
+            DynamicBuffer<Float> albedo_data = ek::load<DynamicBuffer<Float>>(albedo->data(), albedo->pixel_count() * albedo->channel_count());
+            inputs[1].data = albedo_data.data();
+            inputs[1].width = albedo->width();
+            inputs[1].height = albedo->height();
+            inputs[1].format = OPTIX_PIXEL_FORMAT_FLOAT4;
+            inputs[1].rowStrideInBytes = albedo->width() * albedo->bytes_per_pixel();
+            inputs[1].pixelStrideInBytes = albedo->bytes_per_pixel();
+            nb_channels++;
+        }
+        if(normals != nullptr) {
+            DynamicBuffer<Float> normals_data = ek::load<DynamicBuffer<Float>>(normals->data(), normals->pixel_count() * normals->channel_count());
+            inputs[2].data = normals_data.data();
+            inputs[2].width = normals->width();
+            inputs[2].height = normals->height();
+            inputs[2].format = OPTIX_PIXEL_FORMAT_FLOAT4;
+            inputs[2].rowStrideInBytes = normals->width() * normals->bytes_per_pixel();
+            inputs[2].pixelStrideInBytes = normals->bytes_per_pixel();
+            nb_channels++;
+        }
+
+        optix_check(optixDenoiserSetup(denoiser, 0, noisy.width(), noisy.height(), state, state_size, scratch, scratch_size));
+
+        params.denoiseAlpha = 0;
+        params.hdrIntensity = intensity;
+        params.blendFactor  = 0.0f;
+
+        optix_check(optixDenoiserComputeIntensity(denoiser, 0, inputs, intensity, scratch, scratch_size)); 
+        optix_check(optixDenoiserInvoke(denoiser, 0, &params, state, state_size, inputs, nb_channels, 
+                                        0, 0, &output, scratch, scratch_size));
+
+        Float denoised_data = ek::map<Float>(output.data, noisy.pixel_count() * noisy.channel_count());
+        // Float denoised_data = ek::map<Float>(output.data, 114688);
+
+        ek::eval();
+        ek::sync_thread();
+
+        std::cout << denoised_data << std::endl;
+        optix_check(optixDenoiserDestroy(denoiser));
+        return denoised_data;
+    }
+    return Float(0.0);
+}
+
+template <typename Float>
+Float denoise (const Bitmap& noisy) {
+    return denoise<Float>(noisy, nullptr, nullptr);
+}
+
+NAMESPACE_END(mitsuba)

include/mitsuba/render/optix_api.h

@@ -231,6 +231,120 @@ struct OptixShaderBindingTable {
     unsigned int callablesRecordCount;
 };
 
+/// Various sizes related to the denoiser.
+///
+/// \see #optixDenoiserComputeMemoryResources()
+struct OptixDenoiserSizes
+{
+    size_t       stateSizeInBytes;
+    size_t       withOverlapScratchSizeInBytes;
+    size_t       withoutOverlapScratchSizeInBytes;
+    unsigned int overlapWindowSizeInPixels;
+};
+
+/// Various parameters used by the denoiser
+///
+/// \see #optixDenoiserInvoke()
+/// \see #optixDenoiserComputeIntensity()
+/// \see #optixDenoiserComputeAverageColor()
+struct OptixDenoiserParams
+{
+    /// if set to nonzero value, denoise alpha channel (if present) in first inputLayer image
+    unsigned int denoiseAlpha;
+
+    /// average log intensity of input image (default null pointer). points to a single float.
+    /// with the default (null pointer) denoised results will not be optimal for very dark or
+    /// bright input images.
+    CUdeviceptr  hdrIntensity;
+
+    /// blend factor.
+    /// If set to 0 the output is 100% of the denoised input. If set to 1, the output is 100% of
+    /// the unmodified input. Values between 0 and 1 will linearly interpolate between the denoised
+    /// and unmodified input.
+    float        blendFactor;
+
+    /// this parameter is used when the OPTIX_DENOISER_MODEL_KIND_AOV model kind is set.
+    /// average log color of input image, separate for RGB channels (default null pointer).
+    /// points to three floats. with the default (null pointer) denoised results will not be
+    /// optimal.
+    CUdeviceptr  hdrAverageColor;
+};
+
+/// Input kinds used by the denoiser.
+///
+/// RGB(A) values less than zero will be clamped to zero.
+/// Albedo values must be in the range [0..1] (values less than zero will be clamped to zero).
+/// The normals must be transformed into screen space. The z component is not used.
+/// \see #OptixDenoiserOptions::inputKind
+enum OptixDenoiserInputKind
+{
+    OPTIX_DENOISER_INPUT_RGB               = 0x2301,
+    OPTIX_DENOISER_INPUT_RGB_ALBEDO        = 0x2302,
+    OPTIX_DENOISER_INPUT_RGB_ALBEDO_NORMAL = 0x2303,
+};
+
+/// Model kind used by the denoiser.
+///
+/// \see #optixDenoiserSetModel()
+enum OptixDenoiserModelKind
+{
+    /// Use the model provided by the associated pointer.  See the programming guide for a
+    /// description of how to format the data.
+    OPTIX_DENOISER_MODEL_KIND_USER = 0x2321,
+
+    /// Use the built-in model appropriate for low dynamic range input.
+    OPTIX_DENOISER_MODEL_KIND_LDR = 0x2322,
+
+    /// Use the built-in model appropriate for high dynamic range input.
+    OPTIX_DENOISER_MODEL_KIND_HDR = 0x2323,
+
+    /// Use the built-in model appropriate for high dynamic range input and support for AOVs
+    OPTIX_DENOISER_MODEL_KIND_AOV = 0x2324,
+
+};
+
+/// Options used by the denoiser
+///
+/// \see #optixDenoiserCreate()
+struct OptixDenoiserOptions
+{
+    /// The kind of denoiser input.
+    OptixDenoiserInputKind inputKind;
+};
+
+/// Pixel formats used by the denoiser.
+///
+/// \see #OptixImage2D::format
+enum OptixPixelFormat
+{
+    OPTIX_PIXEL_FORMAT_HALF3  = 0x2201,  ///< three halfs, RGB
+    OPTIX_PIXEL_FORMAT_HALF4  = 0x2202,  ///< four halfs, RGBA
+    OPTIX_PIXEL_FORMAT_FLOAT3 = 0x2203,  ///< three floats, RGB
+    OPTIX_PIXEL_FORMAT_FLOAT4 = 0x2204,  ///< four floats, RGBA
+    OPTIX_PIXEL_FORMAT_UCHAR3 = 0x2205,  ///< three unsigned chars, RGB
+    OPTIX_PIXEL_FORMAT_UCHAR4 = 0x2206   ///< four unsigned chars, RGBA
+};
+
+/// Image descriptor used by the denoiser.
+///
+/// \see #optixDenoiserInvoke(), #optixDenoiserComputeIntensity()
+struct OptixImage2D
+{
+    /// Pointer to the actual pixel data.
+    CUdeviceptr data;
+    /// Width of the image (in pixels)
+    unsigned int width;
+    /// Height of the image (in pixels)
+    unsigned int height;
+    /// Stride between subsequent rows of the image (in bytes).
+    unsigned int rowStrideInBytes;
+    /// Stride between subsequent pixels of the image (in bytes).
+    /// For now, only 0 or the value that corresponds to a dense packing of pixels (no gaps) is supported.
+    unsigned int pixelStrideInBytes;
+    /// Pixel format.
+    OptixPixelFormat format;
+};
+
 // =====================================================
 //             Commonly used OptiX functions
 // =====================================================

src/bsdfs/diffuse.cpp

@@ -178,6 +178,11 @@ class SmoothDiffuse final : public BSDF<Float, Spectrum> {
 
         return { depolarizer<Spectrum>(value) & active, dr::select(active, pdf, 0.f) };
     }
+
+    /// Return the diffuse reflectance value (if any)
+    Spectrum get_diffuse_reflectance(const SurfaceInteraction3f &si, Mask active) const override {
+        return m_reflectance->eval(si, active);
+    }
 
     std::string to_string() const override {
         std::ostringstream oss;

src/bsdfs/plastic.cpp

@@ -360,6 +360,10 @@ class SmoothPlastic final : public BSDF<Float, Spectrum> {
                  dr::select(active, hemi_pdf * prob_diffuse, 0.f) };
     }
 
+    Spectrum get_diffuse_reflectance(const SurfaceInteraction3f &si, Mask active) const override {
+        return m_diffuse_reflectance->eval(si, active);
+    }
+
     std::string to_string() const override {
         std::ostringstream oss;
         oss << "SmoothPlastic[" << std::endl

src/bsdfs/roughplastic.cpp

@@ -501,6 +501,11 @@ class RoughPlastic final : public BSDF<Float, Spectrum> {
         return { depolarizer<Spectrum>(value) & active, pdf };
     }
 
+    /// Return the diffuse reflectance value (if any)
+    Spectrum get_diffuse_reflectance(const SurfaceInteraction3f &si, Mask active) const override {
+        return m_diffuse_reflectance->eval(si, active);
+    }
+
     std::string to_string() const override {
         std::ostringstream oss;
         oss << "RoughPlastic[" << std::endl

src/render/bsdf.cpp

@@ -21,6 +21,13 @@ MI_VARIANT Spectrum BSDF<Float, Spectrum>::eval_null_transmission(
     return 0.f;
 }
 
+MI_VARIANT Spectrum BSDF<Float, Spectrum>::get_diffuse_reflectance(
+    const SurfaceInteraction3f &si, Mask active) const {
+    Vector3f placeholder = Vector3f(0.0f, 0.0f, 1.0f);
+    BSDFContext ctx;
+    return eval(ctx, si, placeholder, active) * 3.14159265358979323846; //M_PI
+}
+
 template <typename Index>
 std::string type_mask_to_string(Index type_mask) {
     std::ostringstream oss;

src/render/optix_api.cpp

@@ -27,6 +27,13 @@ void optix_initialize() {
     L(optixProgramGroupCreate);
     L(optixProgramGroupDestroy)
     L(optixSbtRecordPackHeader);
+    L(optixDenoiserCreate);
+    L(optixDenoiserSetModel);
+    L(optixDenoiserComputeMemoryResources);
+    L(optixDenoiserSetup);
+    L(optixDenoiserComputeIntensity);
+    L(optixDenoiserInvoke);
+    L(optixDenoiserDestroy);
 
     #undef L
 }

src/render/python/CMakeLists.txt

@@ -2,6 +2,7 @@ set(RENDER_PY_V_SRC
   ${CMAKE_CURRENT_SOURCE_DIR}/bsdf_v.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/emitter_v.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/endpoint_v.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/denoiser_v.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/film_v.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/fresnel_v.cpp
   ${CMAKE_CURRENT_SOURCE_DIR}/imageblock_v.cpp

src/render/python/bsdf_v.cpp

@@ -60,6 +60,11 @@ MI_VARIANT class PyBSDF : public BSDF<Float, Spectrum> {
         PYBIND11_OVERRIDE_PURE(Return, BSDF, eval_pdf, ctx, si, wo, active);
     }
 
+    Spectrum get_diffuse_reflectance(const SurfaceInteraction3f &si,
+              Mask active) const override {
+        PYBIND11_OVERRIDE_PURE(Spectrum, BSDF, get_diffuse_reflectance, si, active);
+    }
+
     std::string to_string() const override {
         PYBIND11_OVERRIDE_PURE(std::string, BSDF, to_string,);
     }
@@ -96,6 +101,10 @@ template <typename Ptr, typename Cls> void bind_bsdf_generic(Cls &cls) {
              [](Ptr bsdf, const SurfaceInteraction3f &si, Mask active) {
                  return bsdf->eval_null_transmission(si, active);
              }, "si"_a, "active"_a = true, D(BSDF, eval_null_transmission))
+        .def("get_diffuse_reflectance",
+             [](Ptr bsdf, const SurfaceInteraction3f &si, Mask active) {
+                 return bsdf->get_diffuse_reflectance(si, active);
+             }, "si"_a, "active"_a = true)
         .def("flags", [](Ptr bsdf) { return bsdf->flags(); }, D(BSDF, flags))
         .def("needs_differentials",
              [](Ptr bsdf) { return bsdf->needs_differentials(); },

src/render/python/denoiser_v.cpp

@@ -0,0 +1,12 @@
+#include <mitsuba/render/denoiser.h>
+#include <mitsuba/python/python.h>
+
+MI_PY_EXPORT(denoiser) {
+    MI_PY_IMPORT_TYPES()
+    m.def("denoise",
+        py::overload_cast<const Bitmap &>(&denoise<Float>),
+        "noisy"_a, D(denoise))
+    .def("denoise",
+        py::overload_cast<const Bitmap &, Bitmap *, Bitmap *>(&denoise<Float>),
+        "noisy"_a, "albedo"_a, "normals"_a, D(denoise, 2));
+}