Inverse rendering: Load new scenes inbetween optimizer steps

[maxfrei750]

I'd like to carry out an optimization, loading a new scene for each optimization step, but keeping the parameters which are being optimized. E.g. continuously optimize a material using a new scene file for each optimization step. Is there a way to do this?

[njroussel]

I have never done this myself, but I believe that it is possible.

If you take a look at the Object Pose Estimation tutorial, we have a set of parameters which are not scene-dependent. Specifically, we "inject" these parameters inside of the scene with the apply_transformation function. In the tutorial, during the optimization loop, this injection is done over and over for the same params object, but nothing is stopping you from using a new set of parameters from a new scene. You will just need to carefully figure out the name of the scene parameters in each new scene file that you're trying optimizing.

For performance reasons, try your best to destroy every possible variable that is linked to scenes used in previous optimization steps. Dr.Jit isn't aware of the scope of a Mitsuba scene and will therefore assume every alive Mitsuba object is potentially reachable. This means that over the course of the optimization, if objects are not carefully destroyed, the kernels that are submitted to the GPU/CPU by Dr.Jit will keep growing.

[maxfrei750]

Thanks for the quick reply, especially for the pointer to the tutorial and the best practices! I'll try to share my insights, as soon as I had the time for some tests. But I might take a while.

[maxfrei750]

Based on the input from @njroussel, I created the following MWE. To run it, the Mitsuba test scenes are required.

"""Minimal working example to test an optimization, with a reloaded scene for every
iteration."""

import pathlib

import drjit as dr
import mitsuba as mi
import numpy as np
from tqdm import tqdm

mi.set_variant("cuda_ad_mono")    # llvm_ad_rgb works

BASE_SCENE_DICT = {
    "type": "scene",
    "material": {
        "type": "principled",
        "base_color": 1.0
    },
    "fill_light": {
        "type": "constant",
        "radiance": 0.5,
    },
    "integrator": {
        "type": "path"
    },
    "sensor": {
        "type":
            "orthographic",
        "to_world":
            mi.ScalarTransform4f.look_at(origin=[0, 0, 100],
                                         target=[0, 0, 0],
                                         up=[0, 1, 0]),
        "film": {
            "type": "hdrfilm",
            "pixel_format": "luminance" if "mono" in mi.variant() else "rgb",    # pylint: disable=not-callable
        }
    },
}


def get_random_scene(seed: int):
    rng = np.random.default_rng(seed=seed)

    object_type_id = rng.integers(2)

    object_types = {0: "bunny", 1: "teapot"}
    object_scales = {"bunny": 5, "teapot": 0.1}

    scene_dict = BASE_SCENE_DICT.copy()

    object_type = object_types[object_type_id]
    object_dict = {
        "particle": {
            "type":
                "ply",
            "filename":
                f"./scenes/meshes/{object_type}.ply",
            "to_world":
                mi.ScalarTransform4f.scale(object_scales[object_type]).rotate(
                    (1, 0, 0), rng.uniform(0, 359)).rotate(
                        (0, 1, 0), rng.uniform(0, 359)).rotate((0, 0, 1),
                                                               rng.uniform(0, 359)),
            "bsdf": {
                "type": "ref",
                "id": "material",
            },
        },
    }

    scene_dict.update(object_dict)

    return mi.load_dict(scene_dict)


def mse(image, image_ref):
    return dr.mean(dr.sqr(image - image_ref))


def mwe_bsdf_optimization():

    output_root = pathlib.Path("output/test")
    output_root.mkdir(exist_ok=True, parents=True)

    reference_scene = get_random_scene(10001)
    parameters = mi.traverse(reference_scene)
    key = "material.base_color.value"
    target_value = 0.0
    parameters[key] = target_value
    parameters.update()
    image_reference = mi.render(reference_scene, parameters, spp=1024)
    mi.util.write_bitmap(str(output_root / "reference.png"), image_reference)

    base_scene = get_random_scene(0)
    parameters = mi.traverse(base_scene)
    optimizer = mi.ad.Adam(lr=0.05)
    optimizer[key] = parameters[key]
    parameters.update(optimizer)

    for iteration_id in tqdm(range(100)):
        scene_new = get_random_scene(iteration_id)
        parameters_new = mi.traverse(scene_new)

        for parameter_name in list(
                parameters.keys()):    # needs to be a list, cause dict changes size
            if "material" in parameter_name:
                continue

            if parameter_name in parameters_new:
                parameters[parameter_name] = parameters_new[parameter_name]
            else:
                raise RuntimeError(
                    "Base scene is incompatible with newly loaded scene.")

        parameters.update()

        image = mi.render(base_scene, parameters)
        loss = mse(image, image_reference)
        dr.backward(loss)
        optimizer.step()
        parameters.update(optimizer)
        mi.util.write_bitmap(str(output_root / f"progress{iteration_id}.png"), image)

    image = mi.render(base_scene, parameters, spp=1024)
    mi.util.write_bitmap(str(output_root / "result.png"), image)


if __name__ == "__main__":
    mwe_bsdf_optimization()

Obviously, using MSE as optimization criterion does not work very well, but that's not the point of this demo.