Implementation of "Large Steps in Inverse Rendering of Geometry" in Mitsuba3

[DoeringChristian]

Hi,

As I did not find any examples of how to optimize the geometry of a mesh, I decided to integrate the reparametrization from
I integrated the reparametrization from the paper Large Steps in Inverse
Rendering of Geometry in
Mitsuba3 and I thought this might be interesting to some.

Setup

As in the tutorials we need to import mitsuba and drjit and set the
variant.
In addition, we also need trimesh for creating an icosphere mesh and NumPy for converting between trimesh and Mitsuba.
The largesteps package is also a necessary requirement though we are going to
import it in the helper functions to avoid name conflicts.
Finally, PyTorch is necessary to perform the matrix multiplications in the largesteps package.

import torch
import numpy as np
import trimesh

import mitsuba as mi
import drjit as dr

mi.set_variant("cuda_ad_rgb", "llvm_ad_rgb")

Helper Functions for Largesteps

To make the use of the largesteps package easier we can implement some helper functions.
First, the construction of the parameterization matrix (which is calculated
from the Laplacian matrix of the mesh) does not require gradients, and so we
can use a PyTorch tensor directly.
The .torch() can be used to convert Mitsuba tensors to PyTorch tensors,
however the faces and positions have to be converted to tensors.
Since the positions and faces are stored in a flat array we can construct the
Mitsuba tensors in the following way.

positions = mi.TensorXf(positions, shape=(len(positions) // 3, 3))
faces = mi.TensorXi(faces, shape=(len(faces) // 3, 3))

The helper function for computing the parameterization matrix then looks like
this.

def compute_matrix(positions: mi.Float, faces: mi.UInt, lambda_: float) -> torch.Tensor:
    positions = mi.TensorXf(positions, shape=(len(positions) // 3, 3))
    faces = mi.TensorXi(faces, shape=(len(faces) // 3, 3))

    positions = positions.torch()
    faces = faces.torch()

    from largesteps.geometry import compute_matrix

    M = compute_matrix(positions, faces, lambda_)

    return M

The "from_differential" and "to_differential" functions on the other hand
require gradient backpropagation to be enabled, since the gradients are
backpropagated to the reparameterized coordinates.
Using the wrap_ad method we can integrate PyTorch operations into Dr.Jit.
Since we want to use PyTorch tensors in the wrapped functions we use "torch" as
the target argument and "drjit" as the source argument.

def to_differential(M: torch.Tensor, v: mi.Float) -> mi.Float:
    v = mi.TensorXf(v, shape=(len(v) // 3, 3))

    @dr.wrap_ad(source="drjit", target="torch")
    def to_differential_internal(v: torch.Tensor):
        from largesteps.parameterize import to_differential

        return to_differential(M, v)

    return to_differential_internal(v).array


def from_differential(M: torch.Tensor, u: mi.Float, method="Cholesky") -> mi.Float:
    u = mi.TensorXf(u, shape=(len(u) // 3, 3))

    @dr.wrap_ad(source="drjit", target="torch")
    def to_differential_internal(u: torch.Tensor):
        from largesteps.parameterize import from_differential

        return from_differential(M, u, method)

    return to_differential_internal(u).array

Rendering Reference Images

We can construct a scene from which we can render reference images.
In this case one reference image is sufficient to get achieve reconstruction of
the front side of the mesh.
I used a model of suzanne (the Blender monkey).

suzanne = mi.load_dict(
    {
        "type": "ply",
        "filename": "data/suzanne.ply",
        "to_world": mi.ScalarTransform4f.scale(0.5),
        "face_normals": True,
    }
)

scene = {
    "type": "scene",
    "integrator": {
        "type": "prb_reparam",
    },
    "sensor": {
        "type": "perspective",
        "to_world": mi.ScalarTransform4f.look_at(
            origin=(0, 2, 2), target=(0, 0, 0), up=(0, 0, 1)
        ),
        "film": {
            "type": "hdrfilm",
            "width": 1024,
            "height": 1024,
            "sample_border": True,
        },
    },
    "floor": {
        "type": "rectangle",
        "to_world": mi.ScalarTransform4f.translate([0, 0, -1]).scale(10),
    },
    "light": {
        "type": "point",
        "position": [1, 2, 2],
        "intensity": {
            "type": "spectrum",
            "value": 10.0,
        },
    },
}

scene_ref = mi.load_dict(
    {**scene, **{"mesh": suzanne}}
)  # Constructing a reference scene by setting "mesh" to suzanne

ref = mi.render(scene_ref, spp=128)
mi.util.write_bitmap("out/ref.jpg", ref)

Initialization

In order to perform reconstruction, initial mesh is required.
I decided to use an icosphere which has vertices distributed more uniformly than a UV sphere.
One limitation at the moment is that texture coordinates are stored at the
vertices, therefore loading OBJ or PLY file with texture coordinates would result in multiple positions per vertex.
When optimizing, the faces can then disconnect at the vertices and result in non-manifold geometry.

To construct the sphere without texture coordinates I used the
trimesh package.
An alternative would be to use Multi View Geometry (MVG) to generate the initial geometry.
In order to construct a Mitsuba mesh another helper function can be used that
converts a trimesh mesh to a Mitsuba mesh.

def trimesh2mitsuba(mesh: trimesh.Trimesh) -> mi.Mesh:
    positions: np.ndarray = np.array(mesh.vertices)
    indices: np.ndarray = np.array(mesh.faces)
    mesh = mi.Mesh(
        "trimesh",
        vertex_count=positions.shape[0],
        face_count=indices.shape[0],
        has_vertex_normals=False,
        has_vertex_texcoords=False,
    )

    params = mi.traverse(mesh)
    params["vertex_positions"] = dr.ravel(mi.Point3f(positions))
    params["faces"] = dr.ravel(mi.Vector3u(indices))
    params.update()
    return mesh

The mesh construction is derived from the example in the Mesh I/O and
manipulation
Tutorial.

Now we can construct a scene for optimizing the mesh.
This can be done by setting the mesh to the initial mesh, in this case an
icosphere with 4 subdivisions.

scene = mi.load_dict(
    {
        **scene,
        **{
            "mesh": trimesh2mitsuba(
                trimesh.creation.icosphere(subdivisions=4, radius=0.5)
            )
        },
    }
)

Parameter Optimization

We can extract the parameters of the scene using mi.traverse and print them.

params = mi.traverse(scene)
print(params)

Then, the parameterization matrix "M" and can be constructed by utilizing the
helper function defined above.

positions = params["mesh.vertex_positions"]
faces = params["mesh.faces"]

M = compute_matrix(positions, faces, lambda_=10)
u = to_differential(M, positions)

As these helper functions handle the conversion to tensors of shape [N, 3] we
can pass the flattened positions and faces directly to them.
The reparameterized coordinates "u" are then constructed from the
to_differential function.

Then we can construct the optimizer using the latent variable "u" for the
reparameterized coordinates.
In this case we use the Adam optimizer with a learning rate of 0.01.

opt = mi.ad.Adam(lr=0.01)
opt["u"] = u

Finally, in the optimization loop we need to calculate the "vertex_positions"
from the reparameterized coordinates "u" before updating the parameters.
The rest of the optimization loop is the similar to other optimization tasks.
We run the optimization for 200 iterations which took about 2 minutes on a
RTX3070 laptop graphics card.

losses = []
for i in range(200):
    print(f"Iteration {i}:")
    u = opt["u"]
    v = from_differential(M, u)
    params["mesh.vertex_positions"] = v
    params.update()

    img = mi.render(scene, params, spp=1)

    mi.util.write_bitmap(f"out/{i}.jpg", img)

    loss = dr.mean(dr.sqr(img - ref))
    print(f"loss={loss[0]}")
    losses.append(loss[0])

    dr.backward(loss)

    opt.step()

Visualizing the Results

Reference image:

Using ffmpeg -i %d.jpg -vcodec libx264 -acodec mp2 -pix_fmt yuv420p reparam.mp4 in the out directory it is possible
to generate a video from the image sequence.

reparam.mp4

With matplotlib it is also possible to plot the loss per iteration.

import matplotlib.pyplot as plt

plt.plot(losses)
plt.xlabel("Optimization Step")
plt.ylabel("Loss")
plt.show()

[gerwang]

Thanks for the tutorial! I find it works and helpful!

Seems that suzanne.ply is not uploaded. I created a version using Blender, and want to put it here suzanne.ply.zip in case anyone is interested in.

The import matplotlib as plt should be from matplotlib import pyplot as plt?

I can run the code on my machine and find it really helpful.

[DoeringChristian]

Oh, yay sorry forgot about 😅.
Here is the file I used suzanne.zip.
I also fixed the matplotlib import issue (it probably got lost when writing it in markdown) thanks for the hint.

Btw. I'm not sure weather I should open an issue/feature request about the UV faces what do you think?

[gerwang]

Now I understand what "issue about the UV faces" means. If I directly load the initialized mesh from mi.load_file, some vertices duplicate if they have multiple texture coordinates, making faces disconnected even if they are connected in the ply file. Here is what my ply mesh was before optimization:

Here is how it looks after optimization:

Seems like we need a feature similar to maintain_order=True in trimesh, so that the mesh topology is guaranteed to be the same.

[DoeringChristian]

I probably didn't explain it that well but thanks for the example showing the problem. However I don't think that can be fixed by changing anything in trimesh. The cause of the problem is that a vertex, which is part of two faces can have multiple different UV coordinates. In the image you can see that the top vertex in this case has 5 UV coordinates, depending on the face it is part of. This is necessary represent "seams".

OBJ handles this by representing each face with three sets of three indices, one for the vertex position, one for the UV coordinates and one for the normals. Mitsuba however only has one set of indices for the vertex positions and it is assumed that one vertex corresponds to only one UV coordinate. When loading OBJ files with UV coordinates Mitsuba therefore generates vertices for each of the UV coordinates.
Example, coordinates of an icosahedron:
These are the vertex coordinates, there are 12 of them.

v 0.000000 -1.000000 0.000000
v 0.723600 -0.447215 0.525720
v -0.276385 -0.447215 0.850640
v -0.894425 -0.447215 0.000000
v -0.276385 -0.447215 -0.850640
v 0.723600 -0.447215 -0.525720
v 0.276385 0.447215 0.850640
v -0.723600 0.447215 0.525720
v -0.723600 0.447215 -0.525720
v 0.276385 0.447215 -0.850640
v 0.894425 0.447215 0.000000
v 0.000000 1.000000 0.000000

These are the Texture coordinates, 22 of them:

vt 0.181819 0.000000
vt 0.363637 0.000000
vt 0.909091 0.000000
vt 0.727273 0.000000
vt 0.545455 0.000000
vt 0.272728 0.157461
vt 1.000000 0.157461
vt 0.090910 0.157461
vt 0.818182 0.157461
vt 0.636364 0.157461
vt 0.454546 0.157461
vt 0.181819 0.314921
vt 0.000000 0.314921
vt 0.909091 0.314921
vt 0.727273 0.314921
vt 0.545455 0.314921
vt 0.363637 0.314921
vt 0.272728 0.472382
vt 0.090910 0.472382
vt 0.818182 0.472382
vt 0.636364 0.472382
vt 0.454546 0.472382

And finally, the face indices of 20 faces:

f 1/1 2/6 3/8
f 2/6 1/2 6/11
f 1/3 3/7 4/9
f 1/4 4/9 5/10
f 1/5 5/10 6/11
f 2/6 6/11 11/17
f 3/8 2/6 7/12
f 4/9 3/7 8/14
f 5/10 4/9 9/15
f 6/11 5/10 10/16
f 2/6 11/17 7/12
f 3/8 7/12 8/13
f 4/9 8/14 9/15
f 5/10 9/15 10/16
f 6/11 10/16 11/17
f 7/12 11/17 12/18
f 8/13 7/12 12/19
f 9/15 8/14 12/20
f 10/16 9/15 12/21
f 11/17 10/16 12/22

As you can see every face is defined as:

f <position index>/<UV index> ...

See wiki.

When this file gets loaded, Mitsuba generates the following mesh:

OBJMesh[
  name = "untitled.obj",
  bbox = BoundingBox3f[
    min = [-0.894425, -1, -0.85064],
    max = [0.894425, 1, 0.85064]
  ],
  vertex_count = 22,
  vertices = [704 B of vertex data],
  face_count = 20,
  faces = [240 B of face data],
  face_normals = 0
]

Which has 22 vertices one for every UV coordinate.
Unfortunately when manually loading a mesh for example from trimesh we similarly cannot have multiple UV coordinates for the same vertex and so we would get overlapping UV coordinates. One way of solving this, which I haven't yet tested is to use another set of latent variables for the positions only and from which the vertex coordinates of the mesh would be gathered (using dr.gather).

[DoeringChristian]

Ok, so I have tested this approach and it seems to work. However, simultaneous shape and material reconstruction has some other difficulties as well (I don't think it is a solved problem yet). Here is my implementation of a "LatentMesh", which has latent positions, faces and a mapping that maps a unwrapped mesh onto the latent positions.
Usage is the same as with the normal mesh just that you have to use params["latent_positions"] for the positions and params["latent_faces"] for the indices.

from pathlib import Path
import numpy as np
import mitsuba as mi
import drjit as dr

if __name__ == "__main__":
    mi.set_variant("cuda_ad_rgb")


def collect(
    indices: np.ndarray,
    attributes: np.ndarray,
    mapping=False,
) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
    i0 = indices[:, 0]
    i1 = indices[:, 1]
    i2 = indices[:, 2]
    a0 = attributes[i0, :]
    a1 = attributes[i1, :]
    a2 = attributes[i2, :]
    if mapping:
        return np.concatenate([a0, a1, a2], axis=0), np.concatenate(
            [i0, i1, i2], axis=0
        )
    else:
        return np.concatenate([a0, a1, a2], axis=0)


class LatentMesh(mi.Mesh):
    def __init__(self, props: mi.Properties):
        filename = props.get("filename")

        # Parse OBJ file:

        positions = []
        uvs = []
        normals = []

        i_positions = []
        i_uvs = []
        i_normals = []

        with Path(filename).open() as f:
            for line in f:
                if line.startswith("v "):
                    s = line.split()[1:]
                    positions.append([float(s[0]), float(s[1]), float(s[2])])
                elif line.startswith("vt "):
                    s = line.split()[1:]
                    uvs.append([float(s[0]), float(s[1])])
                elif line.startswith("vn "):
                    s = line.split()[1:]
                    normals.append([float(s[0]), float(s[1]), float(s[2])])
                elif line.startswith("f "):
                    indices = line.split()[1:]
                    i0 = indices[0].split("/")
                    i1 = indices[1].split("/")
                    i2 = indices[2].split("/")

                    i_positions.append([int(i0[0]) - 1, int(i1[0]) - 1, int(i2[0]) - 1])
                    i_uvs.append([int(i0[1]) - 1, int(i1[1]) - 1, int(i2[1]) - 1])
                    i_normals.append([int(i0[2]) - 1, int(i1[2]) - 1, int(i2[2]) - 1])

        # Construct numpy arrays from lists
        positions = np.array(positions)
        uvs = np.array(uvs)
        normals = np.array(normals)
        i_positions = np.array(i_positions, dtype=int)
        i_uvs = np.array(i_uvs, dtype=int)
        i_normals = np.array(i_normals, dtype=int)

        latent_positions = positions

        # Cut faces:
        positions, mapping = collect(i_positions, positions, mapping=True)
        uvs = collect(i_uvs, uvs)
        normals = collect(i_normals, normals)

        # Calculate the corresponding indices
        n_indices = positions.shape[0] / 3
        i0 = np.arange(n_indices, dtype=int)
        i1 = np.arange(n_indices, dtype=int) + n_indices
        i2 = np.arange(n_indices, dtype=int) + n_indices * 2
        i0 = i0[..., np.newaxis]
        i1 = i1[..., np.newaxis]
        i2 = i2[..., np.newaxis]

        indices = np.concatenate([i0, i1, i2], axis=1)

        # Construct the Mitsuba mesh
        super().__init__(
            "latent mesh",
            props=props,
            vertex_count=positions.shape[0],
            face_count=indices.shape[0],
            has_vertex_normals=True,
            has_vertex_texcoords=True,
        )
        params = mi.traverse(super())
        params["faces"] = dr.ravel(mi.Vector3u(indices))
        params["vertex_positions"] = dr.ravel(mi.Point3f(positions))
        params["vertex_normals"] = dr.ravel(mi.Vector3f(normals))
        params["vertex_texcoords"] = dr.ravel(mi.Point2f(uvs))
        params.update()

        # convert latent positions, faces and mapping to mitsuba
        self.latent_positions = dr.ravel(mi.Point3f(latent_positions))
        self.latent_faces = dr.ravel(mi.Vector3u(i_positions))
        self.mapping = mi.UInt(mapping)

    def traverse(self, cb: mi.TraversalCallback):
        super(mi.Shape, self).traverse(cb)
        cb.put_parameter(
            "latent_positions",
            self.latent_positions,
            mi.ParamFlags.Differentiable | mi.ParamFlags.Discontinuous,
        )
        cb.put_parameter(
            "latent_faces", self.latent_faces, mi.ParamFlags.NonDifferentiable
        )
        cb.put_parameter("mapping", self.mapping, mi.ParamFlags.NonDifferentiable)

    def parameters_changed(self, keys):
        params = mi.traverse(super())
        params["vertex_positions"] = dr.ravel(
            dr.gather(mi.Point3f, self.latent_positions, self.mapping)
        )
        params.update()
        super().parameters_changed(keys)


mi.register_mesh("latentmesh", lambda props: LatentMesh(props))

It is not the most efficient implementation by a long shot because I'm loading the obj file in python but I couldn't find a python obj loader that exposes the original faces/positions. Furthermore, I expand the mesh into totally disjoint triangles which results in far more vertices than necessary. However, I tried to use NumPy where possible to reduce computation time. This approach also relies on secondary set of vertices, faces and the mapping which is memory inefficient. But I think this could be a jumping of point If one wants to test different hyper parameters for simultaneous shape and geometry reconstruction.
Here is an example video where I start with a black sphere and used 10 images from different angles (which is necessary for material reconstruction):

sim.mp4

As you can see the reconstruction is not very good but the texture mapping works and the mesh is not split into parts.

[huohuohuohuohuohuohuohuo]

Thanks for your amazing work！May I ask how to debug in the Class LatentMesh (with Pycharm) ? When I set breakpoints (before super().init()) in the Class LatentMesh and debug, the program hit the breakpoint and immediately exit. But the breakpoints after super().init() can be hit and I can step through debugging. Grateful for any replies that may help.

[h-OUS-e]

This is awesome thank you! I was wondering what torch version you were using though, since mine crash in the optimization loop when executing v = from_differential(M, u), and am using torch 2.0.0

[h-OUS-e]

This is the last bit of the log before it crashes. It seems to crash at v = from_differential(M, u). Not sure what is going on, but I'll try it on a different PC to see if the issue is my system.

2023-03-29 14:45:51 INFO  main  [Scene] OptiX ready. (took 121ms)
2023-03-29 14:45:53 INFO  main  [Scene] Building scene in OptiX ..
2023-03-29 14:45:53 INFO  main  [Scene] OptiX ready. (took 3ms)
SceneParameters[
  -----------------------------------------------------------------------------------
  Name                            Flags    Type  Parent
  -----------------------------------------------------------------------------------
  floor.bsdf.reflectance.value    ∂        Float UniformSpectrum
  floor.to_world                  ∂, D     Transform4f Rectangle
  light.position                           Point3f PointLight
  light.intensity.value           ∂        Color3f SRGBReflectanceSpectrum
  mesh.bsdf.reflectance.value     ∂        Float UniformSpectrum
  mesh.vertex_count                        int   Mesh
  mesh.face_count                          int   Mesh
  mesh.faces                               UInt  Mesh
  mesh.vertex_positions           ∂, D     Float Mesh
  mesh.vertex_normals             ∂, D     Float Mesh
  mesh.vertex_texcoords           ∂        Float Mesh
  sensor.near_clip                         float PerspectiveCamera
  sensor.far_clip                          float PerspectiveCamera
  sensor.shutter_open                      float PerspectiveCamera
  sensor.shutter_open_time                 float PerspectiveCamera
  sensor.film.size                         ScalarVector2u HDRFilm
  sensor.film.crop_size                    ScalarVector2u HDRFilm
  sensor.film.crop_offset                  ScalarPoint2u HDRFilm
  sensor.x_fov                    ∂, D     Float PerspectiveCamera
  sensor.to_world                 ∂, D     Transform4f PerspectiveCamera
]
Iteration 0:

[h-OUS-e]

It still didn't work on my other windows computer, but I can confirm it works on google Colab if I remove mi.util.write_bitmap(f"out/{i}.jpg", img)

[DoeringChristian]

Interesting, it could have something to do with the dr.wrap_ad decorator as I also get some memory leak messages when using it.

[sapo17]

Unfortunately, I am also having a similar problem. It seems like whenever u =to_differential_internal(u).array is called something goes wrong. I set a breakpoint there and called the following torch.tensor(u). And similarly kernel crash happens there as well. Weirdly enough, whenever I debug the value of u, the kernel simply crashes.

Does anyone have any clue what might be wrong?

[sapo17]

I found a workaround---rather than the proposed to_differential(M, positions) function which is called at the start as u = to_differential(M, positions), using this (not pretty)

def to_diff(M, positions):
    """
    Actually can be renamed as matMul.
    see also largesteps/parameterize.py "to_differential(...)" at 
    https://github.com/rgl-epfl/large-steps-pytorch 
    """
    positions = mi.TensorXf(positions, shape=(len(positions) // 3, 3))
    positions = positions.torch().cuda()
    u = (M @ positions).cpu().numpy()
    u = mi.TensorXf(u.flatten(), shape=u.shape)
    return u.array

helped me to avoid the kernel crash.

I don't know what is happening in the background, but most likely something fishy. Hope this helps.

[bathal1]

Thanks for sharing your implementation !

We are working on integrating largesteps within Mitsuba directly, in order to avoid the unnecessary dependency on pytorch.

A few notes :

• Mitsuba's Adam already has a uniform parameter in the constructor. Setting this to True when instanciating the optimizer will allow you to use the UniformAdam variant that was described in the paper.
• In order to support meshes with UV seams, you can generate a mapping from the unique vertex representation to the duplicate one by exploiting the fact that these vertex positions will be exactly the same. This is what we did in the largesteps- pytorch public implementation here. The idea is that your latent variable should be the unique vertex position representation, and at each iteration you convert it to the duplicate one using a scatter operation. That way, gradients will be automatically gathered to the unique vertex positions during differentiation. See here for our implementation.

[DoeringChristian]

Thanks for sharing your notes.

It's good to hear that there is a Mitsuba only implementation is being worked on, which will probably be a lot faster/more efficient.

[xdobetter]

Thank you all for the great sharing.

• I have a problem when I opt geometry.In the example above,I change the iter parameter=500(I hope it will recover better).I
  use MeshLab open iter_364.ply is ok,but the geometry is disappeared for the iter_365.ply and find the warning:

• The total opt loss below:
  

• If it helps you understand my problem better,I've also provided a video of the whole optimization process

reparam.mp4

• the iter_364.ply,iter_365.ply file
  iter_geo.zip

I thought produce the problem because exist invalid vertex when optimize geometry vertex, the maybe exist hole? that may be fix by repeatedly remesh and clamp vertex position？

[DoeringChristian]

I'm sorry for taking this long to reply (my University was stealing my time, beginning of semester and all 😅 ).

I tried it for 500 iterations as well. It seems to work for me. Is there any terminal output that could help to debug this issue? If not could you try to add this mi.set_log_level(mi.LogLevel.Debug) line to your python file and look if there is any output that might indicate the cause of the problem?

About the hole on the nose of Suzanne that appears at around 5 seconds of your video, this could be a problem solved by re meshing as you said. I think in the largesteps paper they talk about re meshing which becomes necessary after enough iterations.

Also, After inspecting the mesh in blender again I think optimization from multiple angles seems necessary for better reconstruction. I was somewhat misled by the images that where produced from the correct view.

Mesh from novel view in blender:

For any one interested, writing the mesh to disk is as simple as adding the following line to the optimization loop:

scene.shapes()[1].write_ply(f"out/{i}.ply")

[huohuohuohuohuohuohuohuo]

Hi, xdobetter! Do you fix the problem? I got the same problem. The object just disappeared.

[Jorgenlf]

Hey, I am very new to Mitsuba. I am playing around with this implementation in an attempt to create a 3D mesh of a human face using one or more 2D images. Trying to use this implementation I have run into the following problems:

I wonder if it is possible to use several images to create the complete 3D mesh, and just not one face of it.

From my understanding the reference is simply a .jpg image of suzanne that is synthetically created by rendering the scene of suzanne. Therefore I tried to exchange this reference with my own .jpg image of the face of a person. When running the optimization the background of the image becomes part of the 3D mesh. I therefore tried to remove the background and turning the image into a .png this seemingly helped but upon closer inspection of the 3D mesh the optimization started creating a false background shadow. I am therefore wondering what needs to be done to a .jpg image in order to get the same behaviour as the suzanne example where the background is distinct and unaffected by the optimization.

Grateful for any replies that may help directly or pinpoint me to other resources that may help me in achieving the goal of creating a 3D mesh of a face based on images.

[DoeringChristian]

Hi,
I will try to answer your questions best I can but I also had some thoghts on my own about face reconstruction (but I think there already exist a number of papers).

I wonder if it is possible to use several images to create the complete 3D mesh, and just not one face of it.

In theory nothing should prevent you from doing this. Multiple images should result in a more accurate reconstruction.

This would be an example of how you could implement such a optimization:

import torch
import numpy as np
import trimesh
import glob

import mitsuba as mi
import drjit as dr

import matplotlib.pyplot as plt

mi.set_variant("cuda_ad_rgb", "llvm_ad_rgb")

import latentmesh


def compute_matrix(positions: mi.Float, faces: mi.UInt, lambda_: float) -> torch.Tensor:
    positions = mi.TensorXf(positions, shape=(len(positions) // 3, 3))
    faces = mi.TensorXi(faces, shape=(len(faces) // 3, 3))

    positions = positions.torch()
    faces = faces.torch()

    from largesteps.geometry import compute_matrix

    M = compute_matrix(positions, faces, lambda_)

    return M


def to_differential(M: torch.Tensor, v: mi.Float) -> mi.Float:
    v = mi.TensorXf(v, shape=(len(v) // 3, 3))

    @dr.wrap_ad(source="drjit", target="torch")
    def to_differential_internal(v: torch.Tensor):
        from largesteps.parameterize import to_differential

        return to_differential(M, v)

    return to_differential_internal(v).array


def from_differential(M: torch.Tensor, u: mi.Float, method="Cholesky") -> mi.Float:
    u = mi.TensorXf(u, shape=(len(u) // 3, 3))

    @dr.wrap_ad(source="drjit", target="torch")
    def to_differential_internal(u: torch.Tensor):
        from largesteps.parameterize import from_differential

        return from_differential(M, u, method)

    return to_differential_internal(u).array


# ref = mi.TensorXf(mi.Bitmap("out/ref.exr"))
refs = []

for i, to_world_file in enumerate(glob.glob("out/refs/*.npy")):
    ref = mi.TensorXf(mi.Bitmap(f"out/refs/{i}.exr"))
    to_world = np.load(to_world_file, allow_pickle=True)
    sensor = mi.load_dict(
        {
            "type": "perspective",
            "to_world": mi.ScalarTransform4f(to_world),
            "film": {
                "type": "hdrfilm",
                "width": 1024,
                "height": 1024,
                "sample_border": True,
            },
        },
    )
    refs.append((sensor, ref))

scene: mi.Scene = mi.load_dict(
    {
        "type": "scene",
        "integrator": {
            "type": "prb_reparam",
        },
        "floor": {
            "type": "rectangle",
            "to_world": mi.ScalarTransform4f.translate([0, 0, -1]).scale(10),
        },
        "light": {
            "type": "point",
            "position": [1, 2, 2],
            "intensity": {
                "type": "spectrum",
                "value": 10.0,
            },
        },
        "mesh": {
            "type": "latentmesh",
            "filename": "data/icosphere.obj",
            "bsdf": {
                "type": "principled",
                "base_color": {
                    "type": "bitmap",
                    "bitmap": mi.Bitmap(
                        mi.TensorXf(mi.Float(0), shape=[1024, 1024, 3])
                    ),
                },
            },
        },
    }
)

initial = mi.render(scene, sensor=refs[0][0], spp=128)

params = mi.traverse(scene)
print(params)

positions = params["mesh.latent_positions"]
faces = params["mesh.latent_faces"]

M = compute_matrix(positions, faces, lambda_=10)
u = to_differential(M, positions)

opt_shape = mi.ad.Adam(lr=0.01)
opt_shape["u"] = u

opt_color = mi.ad.Adam(lr=0.1)
opt_color["color"] = params["mesh.bsdf.base_color.data"]

losses = []
for i in range(200):
    print(f"Iteration {i}:")

    u = opt_shape["u"]
    v = from_differential(M, u)
    params["mesh.latent_positions"] = v

    opt_color["color"] = dr.clamp(opt_color["color"], 0, 1)
    params["mesh.bsdf.base_color.data"] = opt_color["color"]

    params.update()

    loss_mean = 0.0
    for j, (sensor, ref) in enumerate(refs):
        img = mi.render(scene, params, sensor=sensor, spp=1, seed=i)

        if j == 0:
            mi.util.write_bitmap(f"out/{i}.jpg", img)

        loss = dr.mean(dr.sqr(img - ref))
        dr.backward(loss)

        loss_mean += loss[0]

    loss_mean /= len(refs)
    print(f"loss={loss_mean}")
    losses.append(loss_mean)

    opt_shape.step()
    opt_color.step()

scene.shapes()[1].write_ply("out/final.ply")

final = mi.render(scene, params, sensor=refs[0][0], spp=128)

fig, axs = plt.subplots(2, 2, figsize=(10, 10))

axs[0][0].plot(losses)
axs[0][0].set_xlabel("iteration")
axs[0][0].set_ylabel("Loss")
axs[0][0].set_title("Image error plot")

axs[0][1].imshow(mi.util.convert_to_bitmap(initial))
axs[0][1].axis("off")
axs[0][1].set_title("Initial Image")

axs[1][0].imshow(mi.util.convert_to_bitmap(final))
axs[1][0].axis("off")
axs[1][0].set_title("Inferred Image")

axs[1][1].imshow(mi.util.convert_to_bitmap(ref))
axs[1][1].axis("off")
axs[1][1].set_title("Reference Image")

fig.show()
plt.show()

Note, that for similtaneous color optimization I used the 'latentmesh' implementation from this comment. However simultaneous color and shape reconstruction can be quite difficult. In the sdf paper they use a 3D texture to represent the BSDF paramsters. I assume this converges better because the pixel positions do not change with updates to the geometry. You could implement sucha BSDF by implementing your own BSDF class (here is the guide that shows how to implement a custom python plugin).

From my understanding the reference is simply a .jpg image of suzanne that is synthetically created by rendering the scene of suzanne.

Thats right but you have to be careful not to use images with overexposed regions, especially for material reconstruction, since the energy of a pixel is clamped and therefore optimization would produce too dark emitters and too bright diffuse colors (I had this issue in my Bachelors Thesis so I wanted to warn you 😅 ). Using .exr images for testing from synthetic images would probably be better. There also exist some LDR2HDR approaches using neural networks.

When running the optimization the background of the image becomes part of the 3D mesh.

This sounds to me like your process converged to a "local minima". There are several factors that can contribute to such a problem. First of all, the initial shape is important for good inference. Generating such a prior using for example traditional SFM methods could improve your results. Another problem can be that you do not have enogh cameras. The single viewpoint example above only works because the scene is very simple, the mesh does not have a locally varying texture and I used a single point light source. The geometry inside of the silluette could therefore be inferred from the brighness. More complex scenes therefore require more references. The uv dependent textures could also be a problem.

Especially for faces I could imagine that there are models for which the geometry only depends on relatively vew parameters. The shape of a face cannot be a triangle for example. Maybee restricing the degrees of freedom using latent variables could also help you acheive better results (Quick aside: It would be interesting to see if it would be possible to combine such an approach with the smooth geometry reconstruction of the large steps paper. For example, reconstruct a prior using differentiable rendering on bone sizes and then fine tune it using the largesteps package while slowing the optimization down for parts of the face where bones are more prominent).

I hope this could help you with your problems and feel free to ask more questions.

[Jorgenlf]

Thank you so much for the extensive reply! I will be checking out your suggested implementation and try to understand it.

[huohuohuohuohuohuohuohuo]

Thank you all for the amazing sharings.

• I am new to Mitsuba and have a problem when opting the geometry. In the initial example, after finishing optimization, I got some warnings:

nanobind: leaked 4 instances!
nanobind: leaked 2 types!

• leaked type "cholespy.CholeskySolverF"
• leaked type "cholespy.MatrixType"
  nanobind: leaked 2 functions!
• leaked function "init"
• leaked function "solve"
  nanobind: this is likely caused by a reference counting issue in the binding code.

reparam.mp4

There may be some memory leak and the mesh also shows big holes. Have you ever encountered similar problems, or could you give me some tips to solve the problems?

[DoeringChristian]

@njroussel
I think I have an Idea of what could be wrong.
In this line, jitc_eval is called which can change the variable that the the index is pointing to.
Afterwards the size is checked against the variable, referenced by src_v which still points to the non-evaluated variable. The resize function however uses the index to acquire the new variable. In my testing this resulted in crashes in very specific cases when was printing the gradients of v and using the prb integrator without parametrization.
Changing the condition, added in the commit to if (jitc_var(src)->size == 1){ fixes the issue (both the crash and the weird behavior). I'm not sure why exactly this is the case but I presume, that with a "dirty" variable that does not fail the check in jitc_var_resize it is replaced by the resized variable and not the evaluated one, effectively circumventing the still pending side effects. However I'm not sure if that is correct.

[njroussel]

You are entirely correct @DoeringChristian !

However, this was already fixed in this commit: mitsuba-renderer/drjit-core@e0d9657 😓 In that case, I'm surprised that we can still get this weird behaviour.

[DoeringChristian]

Interestingly, It seems to work both with the current master branch as well as the v3.3.0 branch when compiling Mitsuba3 myself. However, when using the pre-compiled wheels it does not. Could it be that they have somehow been compiled from a previous commit?

[njroussel]

I double-checked the wheels, they are/were compiled from the correct commit.

The only thing that comes to mind is compiler difference. The wheels are shipped using gcc and you might be building with clang? I hope that's not the cause...

[DoeringChristian]

So it took me a while to get Mitsuba3 to compile with gcc (I'm using Fedora38 and it doesn't really support multiple versions of gcc).
Unfortunately, I can confirm that compiling it with gcc does result in the problem. However compiling using clang-15 (on fedora) does not.

For anyone who is interested in how to compile Mitsuba3 using gcc(-11) on fedora or other platforms that don't support multiple gcc versions I have summarized the steps I have taken to compile Mitsuba3 using docker/podman.

I used podman (a docker alternative) with the following Dockerfile (the commands are for podman which I have found to work great in this case however, docker should work fine as well).

FROM docker.io/nvidia/cuda:12.1.1-devel-ubuntu22.04

# Install packages
RUN apt-get update -y
RUN apt-get install build-essential -y
RUN apt-get install cmake -y
RUN apt-get install ninja-build -y

RUN apt-get install python3-dev -y

RUN apt-get clean
RUN rm -rf /var/lib/apt/lists/*

Even though I'm using the nvidia/cuda image I don't have to setup their Container Toolkit as I was only compiling Mitsuba3 which relies on nvcc and does not require a GPU to be present.

I built the image using this command:

podman build -t cuda .

Note, that when compiling Mitsuba3 it is important that the path of the result stays the same. Therefore I used the following command to gain access to the container:

podman run --privileged -it -v$PWD/:$PWD/ -w$PWD/build cuda

This gives interactive access to the docker container. Then I used the following commands to build Mitsuba3:

cd build
cmake -GNinja ..
ninja

In order to use the compiled Mitsuba3 library in python you can simply source the setpath.sh file in the build directory.

@njroussel should I add these or similar instructions to the compilation tutorial?