Blender script to help align 3D screenshots from eg. Duckstation-3D-Screenshot

3d-screenshot-aligner.py

# 3d-screenshot-aligner.py
#
# Blender script to help align 3D screenshots from eg.
# Duckstation-3D-Screenshot. Given meshes which are overlapping
# pieces of the same scene, you select the "same" polygon in both
# meshes, and it will move the pieces to align those polys in
# world space.
#
# Blender versions tested: 3.6
#
# INSTRUCTIONS
#
#  1. Select exactly two meshes. Rename one of them to have
#     "ROOT" in the object name; the root mesh will stay still
#     and the other will be moved to align with it.
#  2. Enter edit mode so you are editing both meshes at once.
#     Select ONE poly in the root mesh and the corresponding poly
#     in the other mesh. Two polys should be selected in total.
#  3. Paste this script into the text editor and hit the "Run"
#     button. The non-root mesh should now be aligned with the
#     root so that the two polys coincide.
#  4. Blender will be back in Object mode with the root object
#     selected, ready for you to repeat the process with another
#     mesh.

import bpy
import bmesh
from mathutils import Matrix


# Whether to join the objects after aligning them.
# Change this if you want.
JOIN_OBJECTS_AFTER = False


def find_matching_index(seq1, seq2, eq):
    """
    Given two lists which are rotations of one another, returns
    the unique i such that left rotating seq2 by i makes it the
    same as seq1.

    If no such unique rotation exists, returns None.
    """
    matches = []
    for i in range(0, len(seq1)):
        rotated = seq2[i:] + seq2[:i]
        if all(eq(x, y) for x, y in zip(seq1, rotated)):
            matches.append(i)
    return matches[0] if len(matches) == 1 else None


def find_matching_corner(bmesh1, poly1, bmesh2, poly2):
    """
    Assuming poly1 and poly2 are the "same", finds the index i
    such that poly2.loops[i] corresponds to poly1.loops[0].
    """

    # Matching corners by UVs is prefered, since it isn't
    # sensitive to error in the vertex positions.
    epsilon = 0.001
    uvs1 = bmesh1.loops.layers.uv.active
    uvs2 = bmesh2.loops.layers.uv.active
    if uvs1 and uvs2:
        m = find_matching_index(
            [l[uvs1].uv for l in poly1.loops],
            [l[uvs2].uv for l in poly2.loops],
            eq=lambda uv1, uv2: abs(uv1[0] - uv2[0]) < epsilon and abs(uv1[1] - uv2[1]) < epsilon,
        )
        if m is not None:
            print('Corners matched by UV')
            return m

    # Try corner angles next.
    m = find_matching_index(
        [l.calc_angle() for l in poly1.loops],
        [l.calc_angle() for l in poly2.loops],
        eq=lambda x, y: abs(x - y) < 0.01,
    )
    if m is not None:
        print('Corners matched by edge angle')
        return m

    # Fallback to assuming the corner indices are in the same
    # order. If the polys really did both come from the same
    # original mesh, this is probably true.
    print('No corners correspondence found. Corners *may* match up wrong. If they do, try using a different poly.')
    return 0


def move_to_align(obj1, poly1, obj2, poly2, i):
    """
    Moves (translates + rotates) obj2 so that poly2 in obj2 is
    coincident with poly1 in obj1 in world space.

    i identifies how the corners should match up; poly2.loops[i]
    correspond to poly1.loops[0] (and poly2.loops[i+1] with
    poly1.loops[1], etc.)
    """
    # Verts
    a1 = obj1.matrix_world @ poly1.verts[0].co
    b1 = obj1.matrix_world @ poly1.verts[1].co
    c1 = obj1.matrix_world @ poly1.verts[-1].co

    a2 = obj2.matrix_world @ poly2.verts[i].co
    b2 = obj2.matrix_world @ poly2.verts[(i + 1) % len(poly2.verts)].co
    c2 = obj2.matrix_world @ poly2.verts[i - 1].co

    # Edge vectors
    eb1 = (b1 - a1).normalized()
    ec1 = (c1 - a1).normalized()

    eb2 = (b2 - a2).normalized()
    ec2 = (c2 - a2).normalized()

    # Now transform
    #       A2                  A1
    #  ec2 /  \ eb2  into  ec1 /  \ eb1
    #    C2    B2            C1    B1

    # Find T that rotates e2 to e1
    m1 = Matrix([eb1, ec1, eb1.cross(ec1).normalized()]).transposed()
    m2 = Matrix([eb2, ec2, eb2.cross(ec2).normalized()]).transposed()
    T = m1 @ m2.inverted_safe()                   # T @ m2 = m1
    rot = T.to_quaternion().to_matrix().to_4x4()  # force rotation matrix (orthogonal)

    obj2.matrix_world = (
        Matrix.Translation(a1) @
        rot @
        Matrix.Translation(-a2) @
        obj2.matrix_world
    )


def main():
    assert bpy.context.mode == 'EDIT_MESH', "Must be in mesh edit mode"

    edit_objs = [obj for obj in bpy.data.objects if obj.type == 'MESH' and obj.mode == 'EDIT']
    assert len(edit_objs) == 2, "Must be editing exactly two meshes"
    obj1, obj2 = edit_objs

    # Choose the root object
    if 'ROOT' in obj2.name and 'ROOT' not in obj1.name:
        obj1, obj2 = obj2, obj1

    # Make root object active
    bpy.context.view_layer.objects.active = obj1

    bmesh1 = bmesh.from_edit_mesh(obj1.data)
    bmesh2 = bmesh.from_edit_mesh(obj2.data)

    # Find selected poly in each mesh
    selpolys1 = [face for face in bmesh1.faces if face.select]
    selpolys2 = [face for face in bmesh2.faces if face.select]
    assert len(selpolys1) == 1 and len(selpolys2) == 1, "Must select one poly in each mesh"
    poly1 = selpolys1[0]
    poly2 = selpolys2[0]

    assert len(poly1.verts) == len(poly2.verts), "Polys must have the same number of verts"

    i = find_matching_corner(bmesh1, poly1, bmesh2, poly2)
    move_to_align(obj1, poly1, obj2, poly2, i)

    bpy.ops.object.mode_set(mode='OBJECT')

    if JOIN_OBJECTS_AFTER:
        bpy.ops.object.join()
    else:
        obj2.select_set(False)


if __name__ == "__main__":
    main()