osmWorkflow/blender/osmassets/tree.py

"""Import Poly Haven island_tree_01 as instanced-tree asset.

This is the analog of load_tuft_variants() for trees.  It appends only the
LOD1 meshes from the vendored .blend, which automatically pulls in their
material definitions and texture images.  The top-level object parenting
(branch and leaf parts share one origin) is handled by the caller who
instances them as a single logical tree.
"""

import os

import bpy


TREE_BLEND = os.path.join(
    os.path.dirname(os.path.abspath(__file__)), "..", "..",
    "assets", "models", "polyhaven", "island_tree_01",
    "island_tree_01_1k.blend",
)

# Branch / leaf LOD1 pairings for the four variants.
# Each pairing is a single logical tree that the caller instances as one.
VARIANT_PAIRS = [("a", "b"), ("b", "a"), ("c", "a"), ("d", "b")]


def load_tree_variants():
    """Import the vendored Poly Haven tree and return meshes ready to instance.

    Returns a list of (branch_mesh, leaf_mesh) tuples, one per variant.  The
    meshes share the material datablocks that were appended alongside them, so
    the textures and shader nodes are wired up automatically.

    Returns [] when the source .blend is absent so a clean checkout still
    builds.  The caller falls back to procedural trees.
    """
    if not os.path.exists(TREE_BLEND):
        return []

    before_objects = set(bpy.data.objects)
    before_meshes = set(bpy.data.meshes)

    # Gather the object names we need from the library.
    wanted = []
    for branch_id, leaf_id in VARIANT_PAIRS:
        wanted.append(f"island_tree_01_branches_{branch_id}_LOD1")
        wanted.append(f"island_tree_01_leaves_{leaf_id}_LOD1")

    # bpy.ops.wm.append cannot handle a list of files in a single call
    # reliably (it may crash when objects in the same library share data
    # that was already linked by an earlier append in the same batch), so
    # we append one at a time and skip the "already linked" warning.
    for name in wanted:
        try:
            bpy.ops.wm.append(
                filepath=TREE_BLEND + "/Object/" + name,
                directory=TREE_BLEND + "/Object/",
                files=[{"name": name}],
                link=False,
            )
        except RuntimeError:
            # Object already linked by an earlier append of a sibling mesh
            # that shared materials; this is expected and harmless.
            pass

    imported = [obj for obj in set(bpy.data.objects) - before_objects
                if obj.type == "MESH"]

    # Pair branch + leaf into variants.  We append in the order above so
    # a simple zip works.
    variant_meshes = []
    for i in range(0, len(imported), 2):
        if i + 1 >= len(imported):
            break
        # The import order is interleaved (branch, leaf, branch, leaf, …)
        # because `wanted` alternated.  Sort by name to be safe.
        pass

    # Rebuild the pairs from the imported object names.
    variants = []
    for branch_id, leaf_id in VARIANT_PAIRS:
        b_obj = bpy.data.objects.get(
            f"island_tree_01_branches_{branch_id}_LOD1")
        l_obj = bpy.data.objects.get(
            f"island_tree_01_leaves_{leaf_id}_LOD1")
        if b_obj and l_obj:
            # Copy the meshes so they survive object removal.
            b_mesh = b_obj.data.copy()
            b_mesh.name = "PolyHavenTree_Branch_" + branch_id
            b_mesh.use_fake_user = True
            l_mesh = l_obj.data.copy()
            l_mesh.name = "PolyHavenTree_Leaf_" + leaf_id
            l_mesh.use_fake_user = True
            variants.append((b_mesh, l_mesh))

    for obj in imported:
        bpy.data.objects.remove(obj, do_unlink=True)

    tris = sum(len(bm.loop_triangles) + len(lm.loop_triangles)
               for bm, lm in variants)
    print(f"Poly Haven tree variants loaded: {len(variants)} "
          f"({tris} tris per instance)")
    return variants


def assemble(positions, collection):
    """Place island_tree_01 instances at `positions`.

    positions is a list of (x, y, height) tuples.  Returns the number of
    trees placed (0 when the model is unavailable).
    """
    import math  # noqa — tree.py has no top-level math import by design

    variants = load_tree_variants()
    if not variants:
        return 0

    for index, (x, y, height) in enumerate(positions):
        branch_mesh, leaf_mesh = variants[index % len(variants)]
        target = max(4.0, height)
        factor = target / 3.2

        b_obj = bpy.data.objects.new(f"Tree_PH_{index:03d}_Branch", branch_mesh)
        b_obj.location = (x, y, 0.0)
        b_obj.scale = (factor, factor, factor)
        b_obj.rotation_euler = (0, 0, (index * 1.61803398875) % 1.0 * math.tau)
        collection.objects.link(b_obj)

        l_obj = bpy.data.objects.new(f"Tree_PH_{index:03d}_Leaf", leaf_mesh)
        l_obj.location = (x, y, 0.0)
        l_obj.scale = (factor, factor, factor)
        l_obj.rotation_euler = (0, 0, (index * 1.61803398875) % 1.0 * math.tau)
        collection.objects.link(l_obj)

    return len(positions)