.blend to .tmd (PS1) export

tmd_export.py

"""
.blend to TMD export.
Only objects in "Collection" collection are exported, they're automatically joined and triangulated.
Only flat shaded meshes are supported:
    Use Color Attribute -> Face Corner (NOT Vertex color).

Run via CLI:
    /usr/bin/blender <BLEND_FILE> \
            --quiet --python-exit-code 1 \
            --background --python tmd_export.py \
            -- <OUTPUT_FILE_PATH>
"""

import bmesh
import bpy
import math
import struct
import sys

from operator import attrgetter

from bpy_extras.io_utils import ExportHelper
from bpy.props import StringProperty
from bpy.types import Operator

bl_info = {
    "name": ".blend to .tmd export",
    "description": "A plugin for exporting meshes to .tmd",
    "author": "Elias Daler",
    "version": (0, 1),
    "blender": (4, 1, 0),
    "category": "Import-Export",
}


def triangulate_mesh(mesh):
    bm = bmesh.new()
    bm.from_mesh(mesh)
    bmesh.ops.triangulate(bm, faces=bm.faces)
    bm.to_mesh(mesh)
    mesh.update()
    bm.free()


def float_to_fixed_4_12(f):
    frac, integer = math.modf(f)
    scaled_frac = int(frac * 4096)
    fixed_point_value = (int(integer) << 12) | scaled_frac
    return fixed_point_value


# Rounds normal to 0.01 precision to find similar normals
def to_psx_normal(normal):
    precision = 0.01
    # Blender -> PSX coordinate conversion
    # X' = +X
    # Y' = -Z
    # Z' = +Y
    # You can change the normals to change the lighting direction
    return (
        round(+normal.x / precision) * precision,
        round(-normal.z / precision) * precision,
        round(+normal.y / precision) * precision,
    )


def write_tmd_from_mesh(mesh, path):
    triangulate_mesh(mesh)

    vertex_colors = None
    vertex_colors_domain = None
    if mesh.color_attributes:
        vertex_colors = mesh.color_attributes[0].data
        vertex_colors_domain = mesh.color_attributes[0].domain
    else:
        raise ValueError(f'The mesh {mesh.name} should have a color attribute')

    f = open(path, 'wb')
    f.write(struct.pack('I', 0x41))  # id
    f.write(struct.pack('I', 0x0))  # flags
    f.write(struct.pack('I', 1))  # num objects

    # vertTopAddr
    objectStartPos = f.tell()
    vertTopAddrPos = f.tell()
    f.write(struct.pack('I', 0))

    f.write(struct.pack('I', len(mesh.vertices)))  # num vertices

    # collect unique normals
    normals = []
    unique_normals_map = {}  # normal -> index in normals list
    for poly in mesh.polygons:
        psx_normal = to_psx_normal(poly.normal)
        if psx_normal not in unique_normals_map:
            normal_index = len(normals)
            unique_normals_map[psx_normal] = normal_index
            normals.append(list(psx_normal))

    # normalTopAddr
    normalTopAddrPos = f.tell()
    f.write(struct.pack('I', 0))
    f.write(struct.pack('I', len(normals)))

    # primitiveTopAddr
    primitiveTopAddPos = f.tell()
    f.write(struct.pack('I', 0))
    f.write(struct.pack('I', len(mesh.polygons)))

    f.write(struct.pack('I', 7))  # scale (idk if it matters)

    # return to primitiveTopAddr and write it
    currPos = f.tell()
    primitiveTopAddr = currPos - objectStartPos
    f.seek(primitiveTopAddPos)
    f.write(struct.pack('I', primitiveTopAddr))  # primitiveTopAddr
    f.seek(currPos)

    # prims - F3 only
    for face_idx, poly in enumerate(mesh.polygons):
        verts = []

        psx_normal = to_psx_normal(poly.normal)
        normal_index = unique_normals_map[psx_normal]

        f.write(struct.pack('B', 4))   # olen
        f.write(struct.pack('B', 3))   # ilen
        f.write(struct.pack('B', 0))   # flag
        f.write(struct.pack('B', 32))  # mode

        faceColor = [0, 0, 0]
        for loop_index in poly.loop_indices:
            vi = mesh.loops[loop_index].vertex_index
            vertex = mesh.vertices[vi]
            verts.append(vi)

            if vertex_colors_domain == "POINT":
                color = vertex_colors[vi].color_srgb
            else:
                color = vertex_colors[loop_index].color_srgb
            faceColor = [
                int(round(color[0] * 255.0)),
                int(round(color[1] * 255.0)),
                int(round(color[2] * 255.0)),
            ]

        f.write(struct.pack('B', faceColor[0]))     # r
        f.write(struct.pack('B', faceColor[1]))     # g
        f.write(struct.pack('B', faceColor[2]))     # b
        f.write(struct.pack('B', 32))               # mode
        f.write(struct.pack('H', normal_index))     # normal
        f.write(struct.pack('H', verts[2]))         # v2
        f.write(struct.pack('H', verts[1]))         # v1
        f.write(struct.pack('H', verts[0]))         # v0

    # write vertTopAddr
    currPos = f.tell()
    vertTopAddr = currPos - objectStartPos
    f.seek(vertTopAddrPos)
    f.write(struct.pack('I', vertTopAddr))  # vertTopAddr
    f.seek(currPos)

    for vertex in mesh.vertices:
        # Blender -> PSX coordinate conversion
        # X' = X
        # Y' = -Z
        # Z' = Y
        f.write(struct.pack('h', float_to_fixed_4_12(+vertex.co.x)))
        f.write(struct.pack('h', float_to_fixed_4_12(-vertex.co.z)))
        f.write(struct.pack('h', float_to_fixed_4_12(+vertex.co.y)))
        f.write(struct.pack('H', 0))  # pad

    # write normalTopAddr
    currPos = f.tell()
    normalTopAddr = currPos - objectStartPos
    f.seek(normalTopAddrPos)
    f.write(struct.pack('I', normalTopAddr))  # normalTopAddr
    f.seek(currPos)

    for normal in normals:
        f.write(struct.pack('h', float_to_fixed_4_12(normal[0])))
        f.write(struct.pack('h', float_to_fixed_4_12(normal[1])))
        f.write(struct.pack('h', float_to_fixed_4_12(normal[2])))
        f.write(struct.pack('H', 0))  # pad

def apply_modifiers(obj):
    ctx = bpy.context.copy()
    ctx['object'] = obj
    for _, m in enumerate(obj.modifiers):
        try:
            ctx['modifier'] = m
            with bpy.context.temp_override(**ctx):
                bpy.ops.object.modifier_apply(modifier=m.name)
        except RuntimeError:
            print(f"Error applying {m.name} to {obj.name}, removing it instead.")
            obj.modifiers.remove(m)

    for m in obj.modifiers:
        obj.modifiers.remove(m)

def collect_objects(collection_objects):
    obj_set = set(o for o in collection_objects if o.type == 'MESH')
    obj_list = list(obj_set)
    obj_list.sort(key=attrgetter("name"))
    return obj_list

def collect_meshes(scene):
    # find objects with meshes
    mesh_objects = []
    for obj in scene.objects:
        if obj.type == 'MESH':
            mesh_objects.append(obj)
            apply_modifiers(obj)
            if obj.parent: # possibly armature?
                # apply all transforms
                with bpy.context.temp_override(
                    active_object=obj.parent,
                    selected_editable_objects=[obj.parent]
                ):
                    bpy.ops.object.transform_apply(
                        location=True,
                        rotation=True,
                        scale=True)

        # apply all transforms
        with bpy.context.temp_override(
            active_object=obj,
            selected_editable_objects=[obj]
        ):
            bpy.ops.object.transform_apply(
                location=True,
                rotation=True,
                scale=True)

    meshes_set = set(o.data for o in mesh_objects)

    mesh_list = list(meshes_set)
    mesh_list.sort(key=attrgetter("name"))

    return mesh_list


def write_tmd(context, filepath):
    bpy.ops.object.mode_set(mode='OBJECT')

    scene = context.scene
    meshes = collect_meshes(scene)
    if len(meshes) != 1:
        default_collection = bpy.data.collections.get("Collection")
        obj_list = collect_objects(default_collection.all_objects)
        for obj in obj_list:
            obj.select_set(True)
        bpy.context.view_layer.objects.active = obj_list[0]
        bpy.ops.object.join()
        meshes = collect_meshes(scene)

    write_tmd_from_mesh(meshes[0], filepath)


class ExportTMD(Operator, ExportHelper):
    """Save a PSX TMD file"""
    bl_idname = "psx_tmd.save"
    bl_label = "Export PSX TMD file"

    filename_ext = ".tmd"

    filter_glob = StringProperty(
        default="*.tmd",
        options={'HIDDEN'},
        maxlen=255,
    )

    def execute(self, context):
        return write_tmd(context, self.filepath)


def menu_func_export(self, context):
    self.layout.operator(ExportTMD.bl_idname, text="TMD file")


classes = {
    ExportTMD
}


def register():
    for cls in classes:
        bpy.utils.register_class(cls)
    bpy.types.TOPBAR_MT_file_export.append(menu_func_export)


def unregister():
    bpy.types.TOPBAR_MT_file_export.remove(menu_func_export)

    try:
        for cls in classes:
            bpy.utils.unregister_class(cls)
    except RuntimeError:
        pass


if __name__ == "__main__":
    cli_mode = False
    argv = sys.argv[sys.argv.index("--") + 1:]  # get all args after "--"
    if len(argv) > 0:
        cli_mode = True
        export_filename = argv[0]
        print(f"CLI mode, export to: {export_filename}")
        write_tmd(bpy.context, export_filename)
        sys.exit(0)

    if not cli_mode:
        register()

    # show export menu
    bpy.ops.psxtools_json.save('INVOKE_DEFAULT')