Combining Mitsuba with Python for Photorealistic Renders

In the past, I have used Paraview for visualising voxel models and finite element meshes. Paraview is geared towards scientific visualisation built towards ray-tracing capability for volumes and meshes using OSPRay. Unfortunately, despite options to automate their preparation of models using packages such as PyVista, it often required additional manual image editing. I wanted to find an alternative that semi-automates some of the model visualisation and also improve the quality of the visualisation, whilst providing some additional ‘glossy’ pictures for this website.

Mitsuba 3, is a cross-platform open source photorealistic render originating Wenzel Jakob. This alongside some other photo-realistic rendering codes that are available and can operate within Python. For most purposes these codes are targeted for `academi’c education and for research use for developing accurate physical light rendering models. Some academic papers have used this also for rendering their 3D mode with a better aesthetic quality compared to mainstream mesh editing programs or visualisation software.

Mitsuba is straightforward to setup and install within a standard Python environment (pip install Mitsuba). Mitsuba’s documentation is reasonable to follow in terms of documentation of their plugins and API, however, the number of examples and python excerpts was a little lacking beyond those packaged bundled ‘scene.xml’ files.

For reference, below is an excerpt that can be used to assist with rendering some objects. This involves creating a scene definition using a Python Dict object. Be aware that additional ‘realistic’ material BSDF models are available, such as metal, plastic and transparent/translucent materials. The scene is relatively simple, consisting of a circular disc (z=0) to cast shadows upon to. Careful seleciton of light emitters (area lights) should be provide good illumination in the scene, otherwise convergence and noise artefacts can become present in the render. Crucially, in this script the mesh provided is rescaled and positioned above the disc plane based on the Z-height.

  • Note: an average of the bounding box could be taken. This is done in order to simplfify the alignment and position of the perspective camera in the scene and make it independent of the geometry provided.
import os
import numpy as np
import matplotlib.pyplot as plt
import mitsuba as mi
from mitsuba import ScalarTransform4f as T

import trimesh
from trimesh.transformations import rotation_matrix

mi.set_variant('scalar_rgb')

# Load the mesh file here
myMesh = trimesh.load('bracket.stl')

# Scale the mesh to approximately one unit based on the height
sf = 1.
myMesh.apply_scale(sf/myMesh.extents[2])
myMesh = myMesh.apply_transform(rotation_matrix(np.deg2rad(90), [1.0,0.0,0]))

# Translate the mesh so that it's centroid is at the origin and rests on the ground plane
myMesh.apply_translation([-myMesh.bounds[0,0] - myMesh.extents[0] / 2.0,
                          -myMesh.bounds[0,1] - myMesh.extents[1] / 2.0,
                          -myMesh.bounds[0,2]])

# Fix the mesh normals for the mesh
myMesh.fix_normals()

# Write the mesh to an external file (Wavefront .obj)
with open('mesh.obj', 'w') as f:
    f.write(trimesh.exchange.export.export_obj(myMesh,include_normals=True ))

#Create a sensor that is used for rendering the scene
def load_sensor(r, phi, theta):
    # Apply two rotations to convert from spherical coordinates to world 3D coordinates.
    origin = T.rotate([0, 0, 1], phi).rotate([0, 1, 0], theta) @ mi.ScalarPoint3f([0, 0, r])

    return mi.load_dict({
        'type': 'perspective',
        'fov': 40.0,
        'to_world': T.look_at(
            origin=origin,
            target=[0, 0, myMesh.extents[2]/2],
            up=[0, 0, 1]
        ),
        'sampler': {
            'type': 'independent',
            'sample_count': 16
        },
        'film': {
            'type': 'hdrfilm',
            'width': 1024,
            'height': 768,
            'rfilter': {
                'type': 'tent',
            },
            'pixel_format': 'rgb',
        },
    })

# Scene parameters
relativeLightHeight = 5.0

# A scene dictionary contains the description of the rendering scene.
scene2 = mi.load_dict({
    'type': 'scene',
    # The keys below correspond to object IDs and can be chosen arbitrarily
    'integrator': {'type': 'path'},

    'mesh': {
        'type': 'obj',
        'filename': 'mesh.obj',
        'face_normals': True, # This prevents smoothing of sharp-corners by discarding surface-normals. Useful for engineering CAD.
        'bsdf': {
            'type': 'diffuse',
            'reflectance': {
                'type': 'rgb',
                'value': [0.1, 0.27, 0.86]
            }
        }
    },

    # A general emitter is used for illuminating the entire scene (renders the background white)
    'light': {'type': 'constant', 'radiance': 1.0},
    'areaLight': {
        'type': 'rectangle',
        # The height of the light can be adjusted below
        'to_world': T.translate([0,0.0,myMesh.bounds[1,2] + relativeLightHeight]).scale(1.0).rotate([1,0,0], 5.0),
        'flip_normals': True,
        'emitter': {
            'type': 'area',
            'radiance': {
                'type': 'spectrum',
                'value': 25.0,
            }
        }
    },

    'floor': {
        'type': 'disk',
        'to_world': T.scale(3).translate([0.0,0.0,0.0]),
        'material': {
            'type': 'diffuse',
            'reflectance': {'type': 'rgb', 'value': 0.75},
        }
    }
})

sensor_count = 1

radius = 8
phis = [70.0]
theta = 60.0

sensors = [load_sensor(radius, phi, theta) for phi in phis]

"""
Render the Scene
The render samples are specified in spp
"""
image = mi.render(scene2, sensor=sensors[0], spp=256)

# Write the output
mi.util.write_bitmap("my_first_render.png", image)
mi.util.write_bitmap("my_first_render.exr", image)

# Display the output in an Image
plt.imshow(image** (1.0 / 2.2))
plt.axis('off')

A sample output is produced below of topology optimised bracket used with the above script. In the scene a flat disc (Z=0) is illuminated by a global ‘constant’ light emitter and an area emitter to provide soft shadows.

3D Printing Topology Optimised Component Rendered using Mitsuba
Render of a topology optimised bracket using the Mitsuba 3 Renderer.

On a laptop, this took approximately 30s to render a 1024×768 to 128 samples per pixel (SPP). Renders are relatively quick to generate on modern multi-core computer systems using just their CPUS.

Rendering Meshes with Vertex Colours

Plotting just geometry in a single colour isn’t very interesting, especially when we have results or extra data stored within the mesh. Mitsuba has the option to render colours assigned to each vertex.These can be extracted relatively easily from Finite Element meshes, or other functions. This normally would be trivial, but a subtle trick required is adapting the script to export to the .ply format and separately assign the vertex colour attribute within the scene structure. This can be done by setting the option within the BSDF diffuse material reflectance property, as documented

'bsdf': {
    'type': 'diffuse',
    'reflectance': {
        'type': 'mesh_attribute',
        'name': 'vertex_color'
    }

Trimesh, to date, exports 4-component (RGBA) vertex-colour attribute to the .ply format. The 4-component vertex-colour attribute, however, is unfortunately incompatible with Mitsuba, therefore the data array must be attached separately to the loaded scene, using the traverse function. This can be done by accessing the ‘mesh’ object within the declared Mitsuba scene.

Note: Mitusba uses DrJit for its data representation, but this directly interoperates with numpy arrays. For updating the internal buffer data in the scene, a flat data structure must be supplied to the vertex attribute property.

myMitsubaMesh = scene2.shapes()[2] # Access the .ply mesh object in the loaded scene

# Add a separate 3 component vertex colour attribute with the same number of vertices as the mesh

myMitsubaMesh.add_attribute('vertex_color', 3, [0] * myMesh.vertices.shape[0])
N = myMesh.vertices.shape[0]

# Use Mitsuba traverse function to modify data in the scene graph/structure
meshParams = mi.traverse(myMitsubaMesh)

# Generate a colour mapping based solely on the z-corrdiante of the mesh
vertColor = trimesh.visual.color.interpolate(myMesh.vertices[:,2], 'Paired') [:,:3] / 255.0 

# Update the vertex colour data buffer/array in mitsuba associated with the .ply mesh
meshParams["vertex_color"] = vertColor.ravel()
meshParams.update()

The full example excerpt is presented below

import os
import numpy as np
import mitsuba as mi
import matplotlib.pyplot as plt

import drjit as dr
mi.set_variant('scalar_rgb')

from mitsuba import ScalarTransform4f as T

import trimesh
from trimesh.transformations import rotation_matrix

# Load the mesh file here
myMesh = trimesh.load('bracket.stl')

# Scale the mesh to approximately one unit based on the height
sf = 1.
myMesh.apply_scale(sf/myMesh.extents[2])
myMesh = myMesh.apply_transform(rotation_matrix(np.deg2rad(90), [1.0,0.0,0]))

# Translate the mesh so that it's centroid is at the origin and rests on the ground plane
myMesh.apply_translation([-myMesh.bounds[0,0] - myMesh.extents[0] / 2.0,
                          -myMesh.bounds[0,1] - myMesh.extents[1] / 2.0,
                          -myMesh.bounds[0,2]])

# Fix the mesh normals for the mesh
myMesh.fix_normals()

# Write the mesh to an external file (Wavefront .obj)
with open('mesh.ply', 'wb') as f:
    f.write(trimesh.exchange.export.export_ply(myMesh))

#Create a sensor that is used for rendering the scene
def load_sensor(r, phi, theta):
    # Apply two rotations to convert from spherical coordinates to world 3D coordinates.
    origin = T.rotate([0, 0, 1], phi).rotate([0, 1, 0], theta) @ mi.ScalarPoint3f([0, 0, r])

    return mi.load_dict({
        'type': 'perspective',
        'fov': 40.0,
        'to_world': T.look_at(
            origin=origin,
            target=[0, 0, myMesh.extents[2]/2],
            up=[0, 0, 1]
        ),
        'sampler': {
            'type': 'independent',
            'sample_count': 16
        },
        'film': {
            'type': 'hdrfilm',
            'width': 1024,
            'height': 768,
            'rfilter': {
                'type': 'tent',
            },
            'pixel_format': 'rgb',
        },
    })

# Scene parameters
relativeLightHeight = 5.0

# A scene dictionary contains the description of the rendering scene.
scene2 = mi.load_dict({
    'type': 'scene',
    # The keys below correspond to object IDs and can be chosen arbitrarily
    'integrator': {'type': 'path'},

    'mesh': {
        'type': 'ply',
        'filename': 'mesh.ply',
        'face_normals': True,
        'bsdf': {
            'type': 'diffuse',
            'reflectance': {
                'type': 'mesh_attribute',
                'name': 'vertex_color'
            }
        }
    },

    # A general emitter is used for illuminating the entire scene (renders the background white)
    'light': {'type': 'constant', 'radiance': 1.0},
    'areaLight': {
        'type': 'rectangle',
        # The height of the light can be adjusted below
        'to_world': T.translate([0,0.0,myMesh.bounds[1,2] + relativeLightHeight]).scale(1.0).rotate([1,0,0], 5),
        'flip_normals': True,
        'emitter': {
            'type': 'area',
            'radiance': {
                'type': 'spectrum',
                'value': 25.0,
            }
        }
    },

    'floor': {
        'type': 'disk',
        'to_world': T.scale(3).translate([0.0,0.0,0.0]),
        'material': {
            'type': 'diffuse',
            'reflectance': {'type': 'rgb', 'value': 0.75},
        }
    }
})


myMitsubaMesh = scene2.shapes()[2]
myMitsubaMesh.add_attribute('vertex_color', 3, [0] * myMesh.vertices.shape[0])
N = myMesh.vertices.shape[0]

#vertex_colors = dr.zeros(mi.Float, 3 * N)
#vertex_colors +=
#vertex_colors /= 255

meshParams = mi.traverse(myMitsubaMesh)
vertColor = trimesh.visual.color.interpolate(myMesh.vertices[:,2], 'Paired') [:,:3] / 255.0 #paired / plasma
meshParams["vertex_color"] = vertColor.ravel()
meshParams.update()

sensor_count = 1

radius = 8
phis = [70.0]
theta = 60.0

sensors = [load_sensor(radius, phi, theta) for phi in phis]

"""
Render the Scene
The render samples are specified in spp
"""
image = mi.render(scene2, sensor=sensors[0], spp=256)

# Write the output
mi.util.write_bitmap("my_first_render.png", image)
mi.util.write_bitmap("my_first_render.exr", image)

plt.imshow(image** (1.0 / 2.2))
plt.axis('off')

The output of this is shown below, which uses a stratified colour map to render the z-value position. On closer inspect it can be observed that the interpolation of the Z position is interpolated across each triangle, so the precise isolevel boundaries are not exact. For most purposes, using a high resolution mesh, this would not cause concern.

Rendering of a topology optimised bracket with Z component isolevels generated and attached as vertex_color attribute

Rendering Volumetric Textures

Mitusba 3, provides the opportunity to both render volumes and interestingly apply volumes interpolated as volumetric surface textures. In this situation, a 3-component (RGB) voxel grid can be generated and intersecting values corresponding with intersecting mesh surface are used as the colour information.

Render of a topology optimised bracket with a volume field attached as a surface texture

Mitsuba uses its own simple .vol format for storing voxel grid information, although a convenient Python function does not exist for this. The definition of the current file format used is presented in the table. This is relatively simple to generate and export using Python’s inbuilt file handing functions. More efficient approaches would simply write the Numpy array directly to the file by correctly ordering the data in the array with the channel data (last axis) moving the fastest across the flattened array.

Position [Bytes]Content
1-3ASCII Bytes ’V’, ’O’, and ’L’
4File format version number (currently 3)
5-8Encoding identified (32-bit integer). 1 = Float32
9-12Number of cells along the X axis (32 bit integer)
13-16Number of cells along the Y axis (32 bit integer)
17-20Number of cells along the Z axis (32 bit integer)
21-24Number of channels (32 bit integer, supported values: 1, 3 or 6)
25-48Axis-aligned bounding box of the data stored in single precision (order: xmin, ymin, zmin, xmax, ymax, zmax)
49-*Binary data of the volume stored in the specified encoding. The data are ordered so that the following C-style indexing operation makes sense after the file has been loaded into memory: data[((zpos*yres + ypos)*xres +xpos)*channels + chan] where (xpos, ypos, zpos, chan) denotes the lookup location.
Structure of the .vol file format used by Mitsuba (version 3.0)

Below is an excerpt of the function that can export a numpy 4D (m \times n \times p \times 3 ) array with three colour channels to this file format.


def writeVol(filename, vol: np.ndarray, bbox):

    def int8(val) -> bytes:
        return struct.pack('b', val)

    def int32(val) -> bytes:
        return struct.pack('i', val)

    def float(val) -> bytes:
        return struct.pack('f', val)

    with open(filename, 'wb') as f:
        f.write('VOL'.encode('ascii'))
        f.write(int8(3))
        f.write(int32(1)) # 1 = float type
        f.write(int32(vol.shape[0]) ) # X grid size
        f.write(int32(vol.shape[1]))  # Y grid size
        f.write(int32(vol.shape[2]))  # Z grid size

        f.write(int32(vol.shape[3])) # Number of channels

        # Write the bounding Box of the grid
        # Values [x0,y0,z0, x1,y1,z1]
 
        f.write(bbox.astype(np.float32).tobytes())

        for k in range(vol.shape[2]):
            for j in range(vol.shape[1]):
                for i in range(vol.shape[0]):
                    for m in range(vol.shape[3]):
                        f.write(float(vol[i,j,k,m]))

An excerpt is presented below for generating a TPMS lattice U field. A volume is generated with an equal number of unit cells across each dimension covering the bounding box of the mesh. The values of the U field are transformed using matplotlib colourmap via trimesh.visual.color.interpolate as before and these are exported as the Mitsuba volume format using the inline function writeVol() . Later in the scene definition, this is transformed to align with the original mesh by using the ‘to_world‘ attribute

import os
import numpy as np
import mitsuba as mi
import drjit as dr
import struct
mi.set_variant('scalar_rgb')

from mitsuba import ScalarTransform4f as T

import matplotlib.pyplot as plt

import trimesh
from trimesh.transformations import rotation_matrix

def writeVol(filename, vol: np.ndarray, bbox):

    def int8(val) -> bytes:
        return struct.pack('b', val)

    def int32(val) -> bytes:
        return struct.pack('i', val)

    def float(val) -> bytes:
        return struct.pack('f', val)

    vol.shape[0]

    with open(filename, 'wb') as f:
        f.write('VOL'.encode('ascii'))
        f.write(int8(3))
        f.write(int32(1)) # 1 = float type
        f.write(int32(vol.shape[0]) ) # X grid size
        f.write(int32(vol.shape[1]))  # Y grid size
        f.write(int32(vol.shape[2]))  # Z grid size

        f.write(int32(vol.shape[3])) # Number of channels

        # Write the bounding Box of the grid
        # Values [x0,y0,z0, x1,y1,z1]
        print('bounding box', bbox)
        f.write(bbox.astype(np.float32).tobytes())

        for k in range(vol.shape[2]):
            for j in range(vol.shape[1]):
                for i in range(vol.shape[0]):
                    for m in range(vol.shape[3]):
                        f.write(float(vol[i,j,k,m]))


# Load the mesh file here
myMesh = trimesh.load('bracket.stl')
myMesh = myMesh.apply_transform(rotation_matrix(np.deg2rad(90), [1.0,0.0,0]))
#myMesh.apply_scale([2.5,1,3.0])

# Resolution of the Lattice Grid (note if this is set too low, Mitsuba has render issues)...
res = 0.3499

# Create a gyroid field
Lx = myMesh.extents[0]
Ly = myMesh.extents[1]
Lz = myMesh.extents[2]

""" Number of lattice unit cells"""
cellLength = 5.0

kx = Lx/cellLength
ky = Ly/cellLength
kz = Lz/cellLength

""" Create the computational grid - note np operates with k(z) numerical indexing unlike the default matlab equivalent"""
x,y,z = np.meshgrid(np.arange(0.0, Lx, res),
                    np.arange(0.0, Ly, res),
                    np.arange(0.0, Lz, res))


""" 
Calculating the Gyroid TPMS
"""
Tg = 0.7

U = ( np.cos(kx*2*np.pi*(x/Lx))*np.sin(ky*2*np.pi*(y/Ly))
    + np.cos(ky*2*np.pi*(y/Ly))*np.sin(kz*2*np.pi*(z/Lz))
    + np.cos(kz*2*np.pi*(z/Lz))*np.sin(kx*2*np.pi*(x/Lx)) )**2 - Tg**2

vol = trimesh.visual.color.interpolate(U, 'plasma').reshape(list(U.shape) + [4])[:,:,:,:3] / 256.0

# Delete the temporary variables
del x,y,z, U

# Scale the mesh to approximately one unit based on the height
sf = 2.5
myMesh.apply_scale(sf/myMesh.extents[2])

# Translate the mesh so that it's centroid is at the origin and rests on the ground plane
myMesh.apply_translation([-myMesh.bounds[0,0] - myMesh.extents[0] / 2.0,
                          -myMesh.bounds[0,1] - myMesh.extents[1] / 2.0,
                          -myMesh.bounds[0,2]])


# Fix the mesh normals for the mesh
myMesh.fix_normals()

# Write the volume
bounds = myMesh.bounds.reshape(-1,1).copy()
bounds = np.array([0,0,0,1,1,1])

# Write out the volume
writeVol('out.vol', vol, bounds)


# Write the mesh to an external file (Wavefront .obj)
with open('mesh.ply', 'wb') as f:
    f.write(trimesh.exchange.export.export_ply(myMesh))

#Create a sensor that is used for rendering the scene
def load_sensor(r, phi, theta):
    # Apply two rotations to convert from spherical coordinates to world 3D coordinates.
    origin = T.rotate([0, 0, 1], phi).rotate([0, 1, 0], theta) @ mi.ScalarPoint3f([0, 0, r])

    return mi.load_dict({
        'type': 'perspective',
        'fov': 40.0,
        'to_world': T.look_at(
            origin=origin,
            target=[0, 0, myMesh.extents[2]/2],
            up=[0, 0, 1]
        ),
        'sampler': {
            'type': 'independent',
            'sample_count': 30,

        },
        'film': {
            'type': 'hdrfilm',
            'width': 1024,
            'height': 768,

            'pixel_format': 'rgb',
        },
    })

# Scene parameters
relativeLightHeight = 5.0

# A scene dictionary contains the description of the rendering scene.
scene2 = mi.load_dict({
    'type': 'scene',
    # The keys below correspond to object IDs and can be chosen arbitrarily
    'integrator': {'type': 'path'},


    'mesh': {
        'type': 'ply',
        'filename': 'mesh.ply',
        'face_normals': True,
        'bsdf': {
            'type': 'diffuse',
            'reflectance': {
                'type': 'volume',
                'volume': {
                    'to_world': T.translate([-myMesh.extents[0]/2.0,-myMesh.extents[1]/2.0,0.0]).scale([myMesh.extents[0], myMesh.extents[1], myMesh.extents[2]]),
                    'type': 'gridvolume',
                    'filename': 'out.vol',
                }
            }
        }
    },

    # A general emitter is used for illuminating the entire scene (renders the background white)
    'light': {'type': 'constant', 'radiance': 1.0},
    'areaLight': {
        'type': 'rectangle',
        # The height of the light can be adjusted below
        'to_world': T.translate([0,0.0,myMesh.bounds[1,2] + relativeLightHeight]).scale(1.0).rotate([1,0,0], 5),
        'flip_normals': True,
        'emitter': {
            'type': 'area',
            'radiance': {
                'type': 'spectrum',
                'value': 25.0,
            }
        }
    },

    'floor': {
        'type': 'disk',
        'to_world': T.scale(3).translate([0.0,0.0,0.0]),
        'material': {
            'type': 'diffuse',
            'reflectance': {'type': 'rgb', 'value': 0.75},
        }
    }
})


sensor_count = 1

radius = 8
phis = [70.0]
theta = 60.0

sensors = [load_sensor(radius, phi, theta) for phi in phis]

"""
Render the Scene
The render samples are specified in spp
"""
image = mi.render(scene2, sensor=sensors[0], spp=256)

# Write the output
mi.util.write_bitmap("my_first_render.png", image)
mi.util.write_bitmap("my_first_render.exr", image)

plt.imshow(image** (1.0 / 2.2))
plt.axis('off')

Conclusions

Hopefully, these excerpts and explanations can assist those who wish to render some nice images of their models directly within Python, without the cumbersome additional steps required using external render programs.