wip

model2.py

from dataclasses import dataclass
from pathlib import Path
import xml.etree.ElementTree as XmlET

import bpy


STEAMLIB_DIR = """C:\Program Files (x86)\Steam"""
PACKAGE_NAME = "vanilla"

RWR_INSTALL_PATH = Path(STEAMLIB_DIR) / "steamapps/common/RunningWithRifles"
PACKAGE_PATH = RWR_INSTALL_PATH / f"media/packages/{PACKAGE_NAME}"
PACKAGE_MODELS_PATH = PACKAGE_PATH / "models"
MODEL_NAME = "soldier_a1.xml"
SCALING_FACTOR = 0.2


@dataclass
class Voxel:
    index: int
    x: int
    y: int
    z: int
    r: float
    g: float
    b: float

    def __str__(self):
        return f"Voxel {self.index} [" \
               f"xyz=({self.x},{self.y},{self.z}) " \
               f"rgb=({self.r},{self.g},{self.b})]"

@dataclass
class Particle:
    index: int
    id_: int
    name: str
    x: float
    y: float
    z: float
    invMass: float
    bodyAreaHint: int

    def __str__(self):
        return f"Particle '{self.name}' [" \
               f"index={self.index}, id={self.id_} " \
               f"xyz=({self.x},{self.y},{self.z}) " \
               f"invMass={self.invMass}, bodyAreaHint={self.bodyAreaHint}]"

@dataclass
class Stick:
    index: int
    a: Particle
    b: Particle
    voxels: list[Voxel]

    def __str__(self):
        return f"Stick '{self.a.name}' -> '{self.b.name}' [" \
               f"index={self.index}, count_voxels={len(self.voxels)}]"

@dataclass
class Skeleton:
    particles: dict[int, Particle]
    sticks: list[Stick]

    def __str__(self):
        return f"Skeleton [count_bones={len(self.sticks)}]"

@dataclass
class Model:
    name: str
    skeleton: Skeleton

    @classmethod
    def load(cls, model_path: Path, scaling_factor: float = SCALING_FACTOR) -> "Model":
        with model_path.open(mode="r", encoding="utf-8") as model_file:
            model_xml = XmlET.fromstring(model_file.read())

            # load the voxels into a temporary list
            _voxels = []
            for i, voxel_elem in enumerate(model_xml.findall("./voxels/voxel")):
                x = int(voxel_elem.attrib["x"]) * scaling_factor
                # swap y and z here (in OGRE, y is height)
                y, z = (int(voxel_elem.attrib["z"]) * scaling_factor,
                        int(voxel_elem.attrib["y"]) * scaling_factor)
                r, g, b = float(voxel_elem.attrib["r"]), float(voxel_elem.attrib["g"]), float(voxel_elem.attrib["b"])
                _voxels.append(Voxel(i, x, y, z, r, g, b))

            # load the particles into a dict(id: Particle)
            _particles = dict()
            for i, particle_elem in enumerate(model_xml.findall("./skeleton/particle")):
                id_ = int(particle_elem.attrib["id"])
                name = particle_elem.attrib["name"]
                x = float(particle_elem.attrib["x"])  * scaling_factor
                # swap y and z here
                y, z = (float(particle_elem.attrib["z"]) * scaling_factor,
                        float(particle_elem.attrib["y"]) * scaling_factor)
                invMass = float(particle_elem.attrib["invMass"])
                bodyAreaHint = int(particle_elem.attrib["bodyAreaHint"])
                _particles[id_] = Particle(i, id_, name, x, y, z, invMass, bodyAreaHint)

            # load the sticks into a temporary list of tuple(index: int, a: Particle, b: Particle)
            _sticks = []
            for i, stick_elem in enumerate(model_xml.findall("./skeleton/stick")):
                a, b = int(stick_elem.attrib["a"]), int(stick_elem.attrib["b"])
                # find particles a and b using their int id
                particle_a, particle_b = _particles[a], _particles[b]
                _sticks.append((i, particle_a, particle_b))

            # load skeleton voxel binding groups
            sticks = []
            _indices = set()
            for group in model_xml.findall("./skeletonVoxelBindings/group"):
                constraint_index = int(group.attrib["constraintIndex"])
                # the constraint index is the index of the stick in the _sticks list
                _stick = _sticks[constraint_index]
                voxels = []
                for v in group.findall("./voxel"):
                    _index = int(v.attrib["index"])
                    # sanity check to make sure that a voxel is only bound to one stick
                    if _index in _indices:
                        raise Exception(f"Voxel {_index} bound to more than 1 stick")
                    else:
                        _indices.add(_index)
                        # add the Voxel from _voxels by index to the voxels list
                        voxels.append(_voxels[_index])
                # make a proper Stick from the stick tuple and the voxels we now have here
                s = Stick(_stick[0], _stick[1], _stick[2], voxels)
                sticks.append(s)

            return cls(model_path.stem, Skeleton(_particles, sticks))

    def make(self, voxel_size: float = SCALING_FACTOR):
        voxel_mat = bpy.data.materials.get("VoxelMaterial")
        # create a collection to hold all the model things
        soldier_collection = bpy.data.collections.new(self.name)
        bpy.context.scene.collection.children.link(soldier_collection)
        # create empty skeleton root object
        bpy.ops.object.empty_add(type="ARROWS")
        skeleton_root_obj = bpy.context.object
        # set skeleton root obj name and location
        skeleton_root_obj.name = f"{self.name}_controller"
        skeleton_root_obj.location = (0, 0, 0)
        skeleton_root_obj.show_in_front = True
        # remove the user collection links
        skeleton_root_obj_uc = skeleton_root_obj.users_collection
        for o in skeleton_root_obj_uc:
            o.objects.unlink(skeleton_root_obj)
        # link skeleton root object to the soldier collection
        soldier_collection.objects.link(skeleton_root_obj)

        # make some verts and edges from our particles and sticks
        vertices = [(p.x, p.y, p.z)
                    for i, p in self.skeleton.particles.items()]
        edges = [(s.a.index, s.b.index) for s in self.skeleton.sticks]

        # make skeleton mesh from vertices and edges
        skeleton_mesh = bpy.data.meshes.new(f"{self.name}_skeleton_mesh")
        skeleton_mesh.from_pydata(vertices, edges, [])
        skeleton_mesh.validate()
        skeleton_obj = bpy.data.objects.new(f"{self.name}_skeleton", skeleton_mesh)
        # show the skeleton in front in viewport
        skeleton_obj.show_in_front = True
        # make the controller empty the parent of the skeleton
        skeleton_obj.parent = skeleton_root_obj
        soldier_collection.objects.link(skeleton_obj)
        # add vertex groups to mark the bones
        for stick in self.skeleton.sticks:
            vg = skeleton_obj.vertex_groups.new(name=f"{stick.a.name}->{stick.b.name}")
            vg.add([stick.a.index, stick.b.index], 1.0, "ADD")

        # add some voxels :D
        # create a singular primitive cube to use as a template
        bpy.ops.mesh.primitive_cube_add(size=voxel_size)
        # set a reference to this template voxel
        template_voxel = bpy.context.object
        voxels_obj = bpy.data.objects.new(f"{self.name}_voxels", None)
        voxels_obj.hide_viewport = True
        voxels_obj.parent = skeleton_root_obj
        soldier_collection.objects.link(voxels_obj)

        # draw the voxels - it is much faster to copy and adjust the template voxel than
        # use `bpy.ops.mesh.primitive_cube_add(size=1)` for each and every voxel
        for stick in self.skeleton.sticks:
            # make an object for the voxels attached to this stick
            stick_obj = bpy.data.objects.new(f"{self.name}_stick{stick.index}_" \
                                             f"{stick.a.name}->{stick.b.name}", None)
            stick_obj.empty_display_size = voxel_size
            stick_obj.show_in_front = True
            # set stick_obj position to stick (bone) midpoint
            mx, my, mz = ((stick.a.x + stick.b.x) /2,
                          (stick.a.y + stick.b.y) /2,
                          (stick.a.z + stick.b.z) /2)
            stick_obj.location = (mx, my, mz)
            stick_obj.parent = voxels_obj
            soldier_collection.objects.link(stick_obj)
            linkable = set()
            for voxel in stick.voxels:
                # copy the template, set location, etc
                copy = template_voxel.copy()
                copy.data = template_voxel.data.copy()
                copy.name = f"{self.name}_voxel_{voxel.index}"
                # set the voxel location relative to the bone midpoint empty
                copy.location = (voxel.x - mx, voxel.y - my, voxel.z - mz)
                copy.color = (voxel.r, voxel.g, voxel.b, 1.)
                copy.active_material = voxel_mat
                linkable.add(copy)
            # set voxel parent and link to soldier collection
            for obj in linkable:
                obj.parent = stick_obj
                soldier_collection.objects.link(obj)

        # destroy the template voxel
        # bpy.ops.object.delete({"selected_objects": [template_voxel]})
        with bpy.context.temp_override(selected_objects=[template_voxel]):
            bpy.ops.object.delete()



if __name__ == '__main__':
    print("RwR Voxel Model Renderer!")
    print("Current parameters:")
    print(f"{RWR_INSTALL_PATH=}")
    print(f"{PACKAGE_PATH=}")
    print(f"{PACKAGE_MODELS_PATH=}")
    print(f"{MODEL_NAME=}")

    model_path = PACKAGE_MODELS_PATH / MODEL_NAME

    if not model_path.exists():
        raise Exception(f"Model not found at '{model_path}'")

    print(f"Loading '{model_path}'...")
    voxel_model = Model.load(model_path)
    #print(f"{voxel_model.skeleton.particles[50]=}")
    #print(f"{voxel_model.skeleton.sticks[0]=}")
    #print(f"{voxel_model.skeleton.particles[50]}\n")
    #print(f"{voxel_model.skeleton.sticks[0].voxels[0]}\n")
    #print(f"{voxel_model.skeleton.sticks[0]}\n")
    #print(f"{voxel_model.skeleton}\n")
    print(f"Making model...")
    voxel_model.make()