{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Custom Python plugin" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Overview\n", "\n", "In Mitsuba it easy to add custom code for BSDFs, integrators, emitters, sensors, and more. This tutorial will show you how to create custom plugins in Python and register them for use.\n", "\n", "To illustrate this, we are going to implement a new *tinted dielectric* BSDF, that behaves much like a regular [dielectric BSDF][1] but adds a colorful tint to the reflections at grazing angle. We will then register this new BSDF and use it to render a simple scene.\n", "\n", "

\n", "\n", "🚀 **You will learn how to:**\n", " \n", "

Create a custom BSDF plugin in Python
Register a custom plugin to the system
Use a registered custom plugin

\n", " \n", "

\n", "\n", "[1]: https://mitsuba.readthedocs.io/en/latest/src/generated/plugins_bsdfs.html#smooth-dielectric-material-dielectric" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Setup\n", "\n", "Custom Mitsuba plugins written in Python work best with Just-In-Time (JIT) variants. This is because it would be pretty inefficient to execute Python BSDF code for millions of *scalar* light paths. JIT variants, on the other hand, only execute a few calls to those methods on arrays containing millions of entries at once, mitigating the overhead coming from the the Python layer. In this example we will therefore stick with the `llvm_ad_rgb` variant." ] }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [ "import drjit as dr\n", "import mitsuba as mi\n", "\n", "mi.set_variant('llvm_ad_rgb')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Implementation\n", "\n", "As mentioned in the tutorial overview, we are going to implement a tinted dielectric BSDF in this tutorial. This code is very similar to the actual [C++ implementation][0] of the [dielectric BSDF][1], so we will not look at it in great detail.\n", "\n", "First, our BSDF Python class `MyBSDF` needs to inherit from [BSDF][2]. This allows us to override the constructor method as well as [sample()][3], [eval()][4] and [pdf()][5].\n", "\n", "The constructor takes a [Properties][6] object as an argument, which can be used to read parameters defined in the XML scene description or passed in the [load_dict()][7] dictionary. Here we read the index of refraction ratio `eta` as well as the tint color `tint` from `props`. In the constructor, we also properly set the BSDF members used in other methods like `m_flags` and `m_components`.\n", "\n", "Similarly to the regular dielectric BSDF, the `eval()` and `pdf()` methods of our custom BSDF should always return `0.0` and never be called as it is a degenerate BSDF described by a Dirac delta distribution.\n", "\n", "Regarding the `sample()` method, apart from the computation of the tinted reflection value `value_r`, the rest of code should be identical to the [C++ implementation][0] of `dielectric`.\n", "\n", "Note that it is also possible to override the [to_string()][8] method which is called in any printing/logging routine.\n", "\n", "Finally, we override the implementation of the [traverse()][9] and [parameters_changed()][10] methods to expose the `tint` parameter via the *traverse* mechanism. This will allow us to edit this parameter after the BSDF is instanciated.\n", "\n", "[0]: https://github.com/mitsuba-renderer/mitsuba3/blob/master/src/bsdfs/dielectric.cpp \n", "[1]: https://mitsuba.readthedocs.io/en/latest/src/generated/plugins_bsdfs.html#smooth-dielectric-material-dielectric\n", "[2]: https://mitsuba.readthedocs.io/en/latest/src/api_reference.html#mitsuba.BSDF\n", "[3]: https://mitsuba.readthedocs.io/en/latest/src/api_reference.html#mitsuba.BSDF\n", "[4]: https://mitsuba.readthedocs.io/en/latest/src/api_reference.html#mitsuba.BSDF\n", "[5]: https://mitsuba.readthedocs.io/en/latest/src/api_reference.html#mitsuba.BSDF\n", "[6]: https://mitsuba.readthedocs.io/en/latest/src/api_reference.html#mitsuba.Properties\n", "[7]: https://mitsuba.readthedocs.io/en/latest/src/api_reference.html#mitsuba.load_dict\n", "[8]: https://mitsuba.readthedocs.io/en/latest/src/api_reference.html#mitsuba.Object.to_string\n", "[9]: https://mitsuba.readthedocs.io/en/latest/src/api_reference.html#mitsuba.Object.traverse\n", "[10]: https://mitsuba.readthedocs.io/en/latest/src/api_reference.html#mitsuba.Object.parameters_changed" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "class MyBSDF(mi.BSDF):\n", " def __init__(self, props):\n", " mi.BSDF.__init__(self, props)\n", "\n", " # Read 'eta' and 'tint' properties from `props`\n", " self.eta = 1.33\n", " if props.has_property('eta'):\n", " self.eta = props['eta']\n", "\n", " self.tint = props['tint']\n", "\n", " # Set the BSDF flags\n", " reflection_flags = mi.BSDFFlags.DeltaReflection | mi.BSDFFlags.FrontSide | mi.BSDFFlags.BackSide\n", " transmission_flags = mi.BSDFFlags.DeltaTransmission | mi.BSDFFlags.FrontSide | mi.BSDFFlags.BackSide\n", " self.m_components = [reflection_flags, transmission_flags]\n", " self.m_flags = reflection_flags | transmission_flags\n", "\n", " def sample(self, ctx, si, sample1, sample2, active):\n", " # Compute Fresnel terms\n", " cos_theta_i = mi.Frame3f.cos_theta(si.wi)\n", " r_i, cos_theta_t, eta_it, eta_ti = mi.fresnel(cos_theta_i, self.eta)\n", " t_i = dr.maximum(1.0 - r_i, 0.0)\n", "\n", " # Pick between reflection and transmission\n", " selected_r = (sample1 <= r_i) & active\n", "\n", " # Fill up the BSDFSample struct\n", " bs = mi.BSDFSample3f()\n", " bs.pdf = dr.select(selected_r, r_i, t_i)\n", " bs.sampled_component = dr.select(selected_r, mi.UInt32(0), mi.UInt32(1))\n", " bs.sampled_type = dr.select(selected_r, mi.UInt32(+mi.BSDFFlags.DeltaReflection),\n", " mi.UInt32(+mi.BSDFFlags.DeltaTransmission))\n", " bs.wo = dr.select(selected_r,\n", " mi.reflect(si.wi),\n", " mi.refract(si.wi, cos_theta_t, eta_ti))\n", " bs.eta = dr.select(selected_r, 1.0, eta_it)\n", "\n", " # For reflection, tint based on the incident angle (more tint at grazing angle)\n", " value_r = dr.lerp(mi.Color3f(self.tint), mi.Color3f(1.0), dr.clamp(cos_theta_i, 0.0, 1.0))\n", "\n", " # For transmission, radiance must be scaled to account for the solid angle compression\n", " value_t = mi.Color3f(1.0) * dr.sqr(eta_ti)\n", "\n", " value = dr.select(selected_r, value_r, value_t)\n", " \n", " return (bs, value)\n", "\n", " def eval(self, ctx, si, wo, active):\n", " return 0.0\n", "\n", " def pdf(self, ctx, si, wo, active):\n", " return 0.0\n", "\n", " def eval_pdf(self, ctx, si, wo, active):\n", " return 0.0, 0.0\n", "\n", " def traverse(self, callback):\n", " callback.put_parameter('tint', self.tint, mi.ParamFlags.Differentiable)\n", "\n", " def parameters_changed(self, keys):\n", " print(\"🏝️ there is nothing to do here 🏝️\")\n", "\n", " def to_string(self):\n", " return ('MyBSDF[\\n'\n", " ' eta=%s,\\n'\n", " ' tint=%s,\\n'\n", " ']' % (self.eta, self.tint))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Plugin registration\n", "\n", "There's only one more thing to do before we can use our custom BSDF in scenes. We need to register it in the system so it can be used. This can be done by calling the [register_bsdf()][1] function and specifying the name to be used to instantiate this plugin. The function takes a *constructor lambda* function as the second parameter.\n", "\n", "\n", "

\n", "\n", "📑 **Note**\n", " \n", "Similar functions exist for other types of plugins, e.g. \n", " \n", "

register_integrator
register_emitter
register_sensor
register_film
register_mesh
register_texture
register_volume
register_phasefunction
register_medium
register_sampler

\n", "\n", "

\n", "\n", "[1]: https://mitsuba.readthedocs.io/en/latest/src/api_reference.html#mitsuba.register_bsdf" ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [], "source": [ "mi.register_bsdf(\"mybsdf\", lambda props: MyBSDF(props))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Plugin instantiation" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "You can now use this plugin like you would with any other BSDF plugin and set the appropriate properties of the BSDF expected in its constructor in the XML or `dict` representation." ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "MyBSDF[\n", " eta=1.33,\n", " tint=[0.2, 0.9, 0.2],\n", "]" ] }, "execution_count": 4, "metadata": {}, "output_type": "execute_result" } ], "source": [ "my_bsdf = mi.load_dict({\n", " 'type' : 'mybsdf',\n", " 'tint' : [0.2, 0.9, 0.2],\n", " 'eta' : 1.33\n", "})\n", "\n", "my_bsdf" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Rendering\n", "\n", "Finally, let's use our custom BSDF in an actual scene and render it to see how our tinted BSDF looks like." ] }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [ { "data": { "text/html": [ "

" ], "text/plain": [ "Bitmap[\n", " pixel_format = rgb,\n", " component_format = float32,\n", " size = [768, 576],\n", " srgb_gamma = 1,\n", " struct = Struct<12>[\n", " float32 R; // @0, gamma, premultiplied alpha\n", " float32 G; // @4, gamma, premultiplied alpha\n", " float32 B; // @8, gamma, premultiplied alpha\n", " ],\n", " data = [ 5.06 MiB of image data ]\n", "]" ] }, "execution_count": 5, "metadata": {}, "output_type": "execute_result" } ], "source": [ "scene = mi.load_dict({\n", " 'type': 'scene',\n", " 'integrator': { \n", " 'type': 'path' \n", " },\n", " 'light': { \n", " 'type': 'constant',\n", " 'radiance': 0.99,\n", " },\n", " 'sphere' : {\n", " 'type': 'sphere',\n", " 'bsdf': my_bsdf\n", " },\n", " 'sensor': {\n", " 'type': 'perspective',\n", " 'to_world': mi.ScalarTransform4f.look_at(origin=[0, -5, 5], \n", " target=[0, 0, 0], \n", " up=[0, 0, 1]),\n", " }\n", "})\n", "\n", "image = mi.render(scene)\n", " \n", "mi.Bitmap(image).convert(srgb_gamma=True)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Edit parameters\n", "\n", "As expected, it is possible to access our custom BSDF's parameters using the [traverse][1] mechanism.\n", "\n", "[1]: https://mitsuba.readthedocs.io/en/latest/src/api_reference.html#mitsuba.traverse" ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "SceneParameters[\n", " -------------------------------------------------------------------------------\n", " Name Flags Type Parent\n", " -------------------------------------------------------------------------------\n", " light.radiance.value ∂ Float UniformSpectrum\n", " sensor.near_clip float PerspectiveCamera\n", " sensor.far_clip float PerspectiveCamera\n", " sensor.shutter_open float PerspectiveCamera\n", " sensor.shutter_open_time float PerspectiveCamera\n", " sensor.film.size ScalarVector2u HDRFilm\n", " sensor.film.crop_size ScalarVector2u HDRFilm\n", " sensor.film.crop_offset ScalarPoint2u HDRFilm\n", " sensor.x_fov Float PerspectiveCamera\n", " sensor.to_world Transform4f PerspectiveCamera\n", " sphere.to_world Transform4f Sphere\n", " sphere.bsdf.tint ∂ Array3f64 BSDF\n", "]" ] }, "execution_count": 6, "metadata": {}, "output_type": "execute_result" } ], "source": [ "params = mi.traverse(scene)\n", "params" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We can then update the `tint` value:" ] }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "🏝️ there is nothing to do here 🏝️\n" ] } ], "source": [ "key = 'sphere.bsdf.tint'\n", "params[key] = mi.ScalarColor3f(0.9, 0.2, 0.2)\n", "params.update();" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "When re-rendering this scene, we now see that the new tint is indeed used!" ] }, { "cell_type": "code", "execution_count": 8, "metadata": { "tags": [] }, "outputs": [ { "data": { "text/html": [ "