Bake exact vertex positions to Float32 image?

Question

I'm trying to get a pixel/vertex perfect vector displacement map for a real-time project. The idea is essentially to bake a shape key to a texture so I can run the transformation on the vertex shader alone. Baking the VDMs is normally pretty trivial, the typical method is to get the difference between a given point in world space position of two different meshes and draw that as an RGB float 32 value.

But that method uses the point on a surface, instead of a point of a vertex, so certain complex displacements will have seams in them due to small differences in the color data of two verts. In my use case, I can guarantee that my source and target geometry for the VDM will have the exact same number and order of verts, so I'm not interested in the interpolated data on the face at all. Because the face data is irrelevant, my hope is that the UVs can be arranged at the center of each pixel, and then on that pixel, the position value is stored to the image.

More details here and some example photos: https://www.reddit.com/r/blender/comments/p4mhme/bake_exact_vertex_positions_to_float32_image/

Edit: The dimensions of the image are not important other than it cannot exceed 8k in any in width or height. Probably best if it's roughly square (a 10,000 vert image would best be 100x100 pixels) to make the image more manageable.

Answer 1

I tried to write a script for it. Put it in the text editor, make sure you're in object mode, make sure your mesh is selected, make sure the shapekey you want baked is selected, and hit run. You should get a new image and new UV map.

# Tested with Blender 2.93
import bpy
import numpy as np
import math

assert bpy.context.mode == 'OBJECT'

ob = bpy.context.object
assert ob and ob.type == 'MESH'
mesh = ob.data
nverts = len(mesh.vertices)

# Fetch shapekey vertex coords
vcoords = np.empty(nverts*3, dtype=np.float32)
sk_idx = ob.active_shape_key_index
sk = mesh.shape_keys.key_blocks[sk_idx]
sk.data.foreach_get("co", vcoords)

# Decide width x height of image
width = nverts
height = 1
# Uncomment this to get a square-ish image instead
#width = int(nverts**0.5)
#height = math.ceil(nverts / width)

im = bpy.data.images.new(
    sk.name + ".Texture",
    width,
    height,
    alpha=False,
    float_buffer=True,
    is_data=True,
)

# Fill image pixels with vertex coords
pixels = np.zeros((width*height, 4), dtype=np.float32)
pixels[:nverts, :3] = vcoords.reshape(nverts, 3)
pixels[:, 3] = 1.0
im.pixels.foreach_set(pixels.reshape(width*height*4))

# Calculate UV for each vert
vidxs = np.arange(nverts, dtype=np.uint32)
vuvs = np.stack((vidxs % width, vidxs // width), axis=-1)
vuvs = vuvs.astype(np.float32)
vuvs[:, 0] *= 1/width
vuvs[:, 1] *= 1/height
vuvs[:, 0] += 0.5/width
vuvs[:, 1] += 0.5/height

# Set UVs
nloops = len(mesh.loops)
uv_layer = mesh.uv_layers.new(name=sk.name + ".UV")
loop_vidxs = np.empty(nloops, dtype=np.uint32)
mesh.loops.foreach_get("vertex_index", loop_vidxs)
loop_uvs = vuvs[loop_vidxs]
uv_layer.data.foreach_set("uv", loop_uvs.reshape(2*nloops))

I'm not sure you need the UVs though. You could just pass the vertex index to the vertex shader and have it calculate the UV.