rpgpt/art/Exporter.py

import bpy
import bmesh
import os
import shutil
import struct
from mathutils import *; from math import *
from bpy_extras.io_utils import (axis_conversion)

C = bpy.context
D = bpy.data

LEVEL_EXPORT_NAME = "level"
EXPORT_DIRECTORY = "exported"

print("\n\nLet's get it started")

if os.path.exists(bpy.path.abspath(f"//{EXPORT_DIRECTORY}")):
    shutil.rmtree(bpy.path.abspath(f"//{EXPORT_DIRECTORY}"))
os.makedirs(bpy.path.abspath(f"//{EXPORT_DIRECTORY}"))

def write_b8(f, boolean: bool):
    f.write(bytes(struct.pack("?", boolean)))

def write_f32(f, number: float):
    f.write(bytes(struct.pack("f", number)))

def write_u64(f, number: int):
    f.write(bytes(struct.pack("Q", number)))

def write_i32(f, number: int):
    f.write(bytes(struct.pack("i", number)))

def write_u16(f, number: int): # unsigned short, used in shaders to figure out which bone index is current
    f.write(bytes(struct.pack("H", number)))

def write_v3(f, vector):
    write_f32(f, vector.x)
    write_f32(f, vector.y)
    write_f32(f, vector.z)

def write_quat(f, quat):
    write_f32(f, quat.x)
    write_f32(f, quat.y)
    write_f32(f, quat.z)
    write_f32(f, quat.w)

def write_string(f, s: str):
    encoded = s.encode("utf8")
    write_u64(f, len(encoded))
    f.write(encoded)

def write_4x4matrix(f, m):
    # writes each row, sequentially, row major
    for row in range(4):
        for col in range(4):
            write_f32(f, m[row][col])

def normalize_joint_weights(weights):
    total_weights = sum(weights)
    result = [0,0,0,0]
    if total_weights != 0:
        for i, weight in enumerate(weights): result[i] = weight/total_weights
    return result

# for the level.bin 
level_object_data = []
collision_cubes = []
placed_entities = []
saved_meshes = set()

mapping = axis_conversion(
    from_forward = "Y",
    from_up = "Z",
    to_forward = "-Z",
    to_up = "Y",
)
mapping.resize_4x4()

with open(bpy.path.abspath(f"//{EXPORT_DIRECTORY}/shorttest.bin"), "wb") as f:
    for i in range(4):
        write_u16(f, i)

# meshes can either be Meshes, or Armatures. Armatures contain all mesh data to draw it, and any anims it has

for o in D.objects:
    if o.type == "MESH":
        if o.parent and o.parent.type == "ARMATURE":
            mesh_object = o
            o = o.parent
            object_transform_info = (mesh_name, mapping @ o.location, o.rotation_euler, o.scale)
            if o.users_collection[0].name == 'Level':
                assert(False, "Cannot put armatures in the level. The level is for static placed meshes. For dynamic entities, you put them outside of the level collection, their entity kind is encoded, and the game code decides how to draw them")
            else:
                placed_entities.append((mesh_object.name,) + object_transform_info)
            armature_name = o.data.name
            output_filepath = bpy.path.abspath(f"//{EXPORT_DIRECTORY}/{armature_name}.bin")
            print(f"Exporting armature to {output_filepath}")
            with open(output_filepath, "wb") as f:
                write_b8(f, True)
                bones_in_armature = []
                for b in o.data.bones:
                    bones_in_armature.append(b)
                    
                # the inverse model space pos of the bones
                write_u64(f, len(bones_in_armature))
                for b in bones_in_armature:
                    model_space_pose = mapping @ b.matrix_local
                    inverse_model_space_pose = (mapping @ b.matrix_local).inverted()

                    write_4x4matrix(f, model_space_pose)
                    write_4x4matrix(f, inverse_model_space_pose)
                    write_f32(f, b.length)
                
                # write the pose information
                write_u64(f, len(o.pose.bones))
                for pose_bone in o.pose.bones:
                    parent_index = -1
                    if pose_bone.parent:
                        for i in range(len(bones_in_armature)):
                            if bones_in_armature[i] == pose_bone.parent.bone:
                                parent_index = i
                                break
                        if parent_index == -1:
                            assert(false, f"Couldn't find parent of bone {pose_bone}")
                        print(f"Parent of bone {pose_bone.name} is index {parent_index} in list {bones_in_armature}")
                    parent_space_pose = None
                    
                    if pose_bone.parent:
                        parent_space_pose = pose_bone.parent.matrix.inverted() @ pose_bone.matrix
                    else:
                        parent_space_pose = mapping @ pose_bone.matrix
                        #parent_space_pose = pose_bone.matrix
                        print("parent_space_pose of the bone with no parent:")
                        print(parent_space_pose)
                    
                    write_string(f, pose_bone.bone.name)
                    write_i32(f, parent_index)
                    #parent_space_pose = mapping @ pose_bone.matrix
                    translation = parent_space_pose.to_translation()
                    rotation = parent_space_pose.to_quaternion()
                    scale = parent_space_pose.to_scale()
                    write_v3(f, translation)
                    write_quat(f, rotation)
                    write_v3(f, scale)

                # write the mesh data for the armature
                bm = bmesh.new()
                mesh = mesh_object.to_mesh()
                bm.from_mesh(mesh)
                bmesh.ops.triangulate(bm, faces=bm.faces)
                bm.transform(mapping)
                bm.to_mesh(mesh)
                
                
                vertices = []
                armature = o
                for polygon in mesh.polygons:
                    if len(polygon.loop_indices) == 3:
                        for loopIndex in polygon.loop_indices:
                            loop = mesh.loops[loopIndex]
                            position = mesh.vertices[loop.vertex_index].undeformed_co
                            uv = mesh.uv_layers.active.data[loop.index].uv
                            normal = loop.normal
                            
                            jointIndices = [0,0,0,0]
                            jointWeights = [0,0,0,0]
                            for jointBindingIndex, group in enumerate(mesh.vertices[loop.vertex_index].groups):
                                if jointBindingIndex < 4:
                                    groupIndex = group.group
                                    boneName = mesh_object.vertex_groups[groupIndex].name
                                    jointIndices[jointBindingIndex] = armature.data.bones.find(boneName)
                                    if jointIndices[jointBindingIndex] == -1:
                                        # it's fine that this references a real bone, the bone at index 0,
                                        # because the weight of its influence is 0
                                        jointIndices[jointBindingIndex] = 0
                                        jointWeights[jointBindingIndex] = 0.0
                                    else:
                                        jointWeights[jointBindingIndex] = group.weight
                            
                            
                            vertices.append((position, uv, jointIndices, normalize_joint_weights(jointWeights)))
                
                write_u64(f, len(vertices))
                vertex_i = 0
                for v_and_uv in vertices:
                    v, uv, jointIndices, jointWeights = v_and_uv
                    write_f32(f, v.x)
                    write_f32(f, v.y)
                    write_f32(f, v.z)
                    write_f32(f, uv.x)
                    write_f32(f, uv.y)
                    for i in range(4):
                        write_u16(f, jointIndices[i])
                    for i in range(4):
                        write_f32(f, jointWeights[i])
                    vertex_i += 1
                print(f"Wrote {len(vertices)} vertices")

        else: # if the parent type isn't an armature, i.e just a bog standard mesh
            mesh_name = o.to_mesh().name # use this over o.name so instanced objects which refer to the same mesh, both use the same serialized mesh.
            
            object_transform_info = (mesh_name, mapping @ o.location, o.rotation_euler, o.scale)
            
            if o.users_collection[0].name == 'Level' and mesh_name == "CollisionCube":
                collision_cubes.append((o.location, o.dimensions))
            else:
                if o.users_collection[0].name == 'Level':
                    print(f"Object {o.name} has mesh name {o.to_mesh().name}")
                    assert(o.rotation_euler.order == 'XYZ')
                    level_object_data.append(object_transform_info)
                else:
                    placed_entities.append((o.name,) + object_transform_info)
                
                if mesh_name in saved_meshes:
                    continue
                saved_meshes.add(mesh_name)
                
                assert(mesh_name != LEVEL_EXPORT_NAME)
                output_filepath = bpy.path.abspath(f"//{EXPORT_DIRECTORY}/{mesh_name}.bin")
                print(f"Exporting mesh to {output_filepath}")
                with open(output_filepath, "wb") as f:
                    write_b8(f, False)
                    bm = bmesh.new()
                    mesh = o.to_mesh()
                    bm.from_mesh(mesh)
                    bmesh.ops.triangulate(bm, faces=bm.faces)
                    bm.transform(mapping)
                    bm.to_mesh(mesh)
                    
                    
                    vertices = []

                    for polygon in mesh.polygons:
                        if len(polygon.loop_indices) == 3:
                            for loopIndex in polygon.loop_indices:
                                loop = mesh.loops[loopIndex]
                                position = mesh.vertices[loop.vertex_index].undeformed_co
                                uv = mesh.uv_layers.active.data[loop.index].uv
                                normal = loop.normal
                                
                                vertices.append((position, uv))
                    
                    write_u64(f, len(vertices))
                    for v_and_uv in vertices:
                        v, uv = v_and_uv
                        write_f32(f, v.x)
                        write_f32(f, v.y)
                        write_f32(f, v.z)
                        write_f32(f, uv.x)
                        write_f32(f, uv.y)
                    print(f"Wrote {len(vertices)} vertices")

with open(bpy.path.abspath(f"//{EXPORT_DIRECTORY}/{LEVEL_EXPORT_NAME}.bin"), "wb") as f:
    write_u64(f, len(level_object_data))
    for o in level_object_data:
        mesh_name, blender_pos, blender_rotation, blender_scale = o
        print(f"Writing instanced object of mesh {mesh_name}")
        write_string(f, mesh_name)
        write_f32(f, blender_pos.x)
        write_f32(f, blender_pos.y)
        write_f32(f, blender_pos.z)
        
        write_f32(f, blender_rotation.x)
        write_f32(f, blender_rotation.y)
        write_f32(f, blender_rotation.z)
        
        write_f32(f, blender_scale.x)
        write_f32(f, blender_scale.y)
        write_f32(f, blender_scale.z)
    
    write_u64(f, len(collision_cubes))
    for c in collision_cubes:
        blender_pos, blender_dims = c
        write_f32(f,  blender_pos.x)
        write_f32(f, -blender_pos.y)
        write_f32(f, blender_dims.x)
        write_f32(f, blender_dims.y)
        assert(blender_dims.x > 0.0)
        assert(blender_dims.y > 0.0)
    
    write_u64(f, len(placed_entities))
    for p in placed_entities:
        # underscore is mesh name, prefer object name for name of npc. More obvious
        object_name, _, blender_pos, blender_rotation, blender_scale = p
        print(f"Writing placed entity '{object_name}'")
        write_string(f, object_name)
        write_f32(f, blender_pos.x)
        write_f32(f, blender_pos.y)
        write_f32(f, blender_pos.z)
        
        write_f32(f, blender_rotation.x)
        write_f32(f, blender_rotation.y)
        write_f32(f, blender_rotation.z)
        
        write_f32(f, blender_scale.x)
        write_f32(f, blender_scale.y)
        write_f32(f, blender_scale.z)