Differential Rendering (Deb98) for Mitsuba 2

differential_rendering.py

# Differential rendering (Deb98) script for mitsuba2

import mitsuba
mitsuba.set_variant('scalar_rgb')
from mitsuba.core import Thread, Float, Color3f
from mitsuba.core.xml import load_file
from mitsuba.python.util import traverse

import pyexr
import numpy as np

# depth_with_obj    : xml scene with depth integrator, with virtual object
# depth_without_obj : xml scene with depth integrator, without virual object
# path_with_obj     : xml scene with path integrator, with virtual object
# path_without_obj  : xml scene with path integrator, without virtual object

# resolution of rendered output
WIDTH = 1024
HEIGHT = 768

def render_scene(path):
    exr_path = path[:-3]+"exr"

    scene = load_file(path)
    sensor = scene.sensors()[0]
    integrator = scene.integrator()
    integrator.render(scene, sensor)

    film = sensor.film()
    film.set_destination_file(exr_path)
    film.develop()

    return pyexr.read(exr_path)


if __name__ == '__main__':

    Thread.thread().file_resolver().append('.')

    # Load real photo
    photo = pyexr.read("photo.exr")

    # Render 4 scenes
    depth_with_obj = render_scene("depth_with_obj.xml")
    depth_without_obj = render_scene("depth_without_obj.xml")
    path_with_obj = render_scene("path_with_obj.xml")
    path_without_obj = render_scene("path_without_obj.xml")

    # Extract masks
    depth_diff = (depth_without_obj - depth_with_obj)[:,:,0]
    alpha_path = path_with_obj[:,:,3]
    # Its pixel value is 1 if object
    mask_obj = np.ndarray((HEIGHT, WIDTH))
    # Its pixel value is 1 if local model
    mask_local_obj = np.ndarray((HEIGHT, WIDTH))

    for w in range(WIDTH):
        for h in range(HEIGHT):
            if depth_diff[h][w] != 0:
                mask_obj[h][w] = 1

    # Not used in differential rendering actually
    pyexr.write("mask_obj.exr", mask_obj)

    for w in range(WIDTH):
        for h in range(HEIGHT):
            if alpha_path[h][w] != 0 and mask_obj[h][w] == 0:
                mask_local_obj[h][w] = 1
            else:
                mask_local_obj[h][w] = 0

    pyexr.write("mask_local_obj.exr", mask_local_obj)


    # Differential rendering
    path_diff = path_with_obj - path_without_obj

    # Replace pixel value of photo to object where mask_obj is 1 except alpha value
    for w in range(WIDTH):
        for h in range(HEIGHT):
            if mask_obj[h][w] == 1:
                photo[h][w] = path_with_obj[h][w]

    # Add difference of local model to photo where mask_local_obj is 1
    for w in range(WIDTH):
        for h in range(HEIGHT):
            if mask_local_obj[h][w] == 1:
                photo[h][w][:3] += path_diff[h][w][:3]

    pyexr.write("result.exr", photo)