Support colors in cubify()

[dav-ell]

🚀 Feature

Currently, if you want to export a volume as a mesh, you have to do something like this:

from pytorch3d.structures import Volumes
from pytorch3d.ops import cubify
from pytorch3d.io import IO

volume = # ...
voxels = volume.densities()

cubed = cubify(voxels, thresh=threshold)
IO().save_mesh(cubed, filename)

This produces a mesh, for example, like this:

where the mesh is entirely without color, and the .obj file (if exported as .obj) contains only vertices and faces, no colors.

v -1.000000 -1.000000 -1.000000
v -1.000000 -1.000000 -0.984252
...
f 1 2 1313
f 2 1314 1313
...

The cubify function doesn't allow features to be passed in (https://github.com/facebookresearch/pytorch3d/blob/9446d91fae56dd86b1f58d106360aab4491a0b2a/pytorch3d/ops/cubify.py#L53C1-L243C54) which creates visualization problems.

It is far more desirable for visualization and debugging purposes to be able to see the colors, for example, like this:

Motivation

Voxel data for models learning textured volumes often contain color information. Visualizing the cubified versions of those volumes is often very helpful to get an idea for what your model should be learning. However, since cubify only outputs meshes with no color, and since many 3D viewers (e.g. Blender, XCode) don't visualize colorless meshes well (edges can be very difficult to discern), the meshes are often unhelpful. Adding color, especially if it already exists in a volume, can significantly improve the utility of this library and has the potential to make many guides/tutorials/projects pop much more if you can see the color of e.g. the voxel grid the self-driving car model is learning of the world around it.

Pitch

Being able to pass in an optional features argument would be sufficient for me:

def cubify(voxels, thresh, feats=None, device=None, align: str = "topleft") -> Meshes:

However, since volume densities are already exported with shape (N, 1, D, H, W), and cubify only takes voxels with dims (N, D, H, W) the unused 1 dimension could be expanded to support feature data, which would work fine also.

[dav-ell]

I should note I have a simple version of this working with the below code using trimesh:

def create_colored_mesh(grid_occ, grid_col):
    # Create a unit cube mesh from trimesh primitives
    cube = trimesh.primitives.Box(extents=[1, 1, 1])

    # List to hold all the cube meshes
    cubes = []

    for x in range(grid_occ.shape[0]):
        for y in range(grid_occ.shape[1]):
            for z in range(grid_occ.shape[2]):
                if grid_occ[x, y, z]:  # if voxel is filled, create a cube
                    # Create a new cube using the original cube
                    new_cube = cube.copy()
                    # Move the cube to the correct position
                    new_cube.apply_translation((x, y, z))
                    # Add the cube to the cubes list
                    cubes.append(new_cube)

    # Combine all cube meshes into a single mesh
    mesh = trimesh.util.concatenate(cubes)

    # Create a color array for the mesh
    color = grid_col[grid_occ].reshape(-1, 3)  # Select colors for filled voxels
    color = np.asarray(color, dtype=np.uint8)  # Convert to uint8
    color = np.hstack([color, np.full((color.shape[0], 1), 255, dtype=np.uint8)])  # Add alpha channel, fully opaque

    # Assign colors to vertices
    mesh.visual.vertex_colors = np.repeat(color, 8, axis=0)  # Each cube has 8 vertices

    return mesh

[bottler]

I think this would be quite easy to add to the current cubify function, and I'd be happy if someone wants to send a PR for it.

[bottler]

https://github.com/facebookresearch/pytorch3d/commit/ae9d8787ce4be8be6ac87eeac6ed82ca02919056 implements this.