osmWorkflow/blender/osmassets/tree.py

"""Instanced tree assets — the model side of `--tree-style`.

Two vendored models, reduced to one runtime shape: import once, bake the source
object's orientation into a mesh copy, measure it, then link one lightweight
object per tree that reuses that datablock.  Nothing is duplicated per tree, so
the .blend and the exported GLB carry each mesh and each texture exactly once
no matter how many trees are planted.

    apple     SpeedTree Red Delicious, 4.5k tris, alpha-cut leaf cards
    fattree   low-poly cartoon tree, 2.2k tris, opaque geometry

Materials are rebuilt here rather than taken from the source files, because
neither arrives usable.  57% of the apple's colour texture is transparent —
those are leaf cards, and without an alpha-clipped setup the crown renders as a
solid ball of intersecting quads.  fattree ships a bare Diffuse BSDF and a
texture path that only resolves next to the original .blend.

A third style, `polyhaven`, used to live here.  It appended what the Poly Haven
island_tree_01 file calls its LOD1 objects, but those are not whole trees: the
branch parts are 0.4-unit twigs and the leaf parts are flat clusters hanging
below their own origin, both meant to be scattered by the geometry-nodes setup
in that file.  Planting them directly gave twigs, which is what sent us looking
for these two models.  Removed along with the 78MB asset.
"""

import math
import os
from collections import namedtuple

import bpy

from osmassets.materials import link_alpha_clip


MODEL_ROOT = os.path.abspath(os.path.join(
    os.path.dirname(os.path.abspath(__file__)), "..", "..",
    "assets", "models",
))

APPLE_DIR = os.path.join(MODEL_ROOT, "speedtree", "apple_low")
APPLE_OBJ = os.path.join(APPLE_DIR, "RedDeliciousApple.obj")
APPLE_COLOR = os.path.join(APPLE_DIR, "textures", "apple_color_2k.png")
APPLE_NORMAL = os.path.join(APPLE_DIR, "textures", "apple_normal_2k.png")

FATTREE_DIR = os.path.join(MODEL_ROOT, "lyrog", "fattree")
FATTREE_BLEND = os.path.join(FATTREE_DIR, "fattree.blend")
FATTREE_COLOR = os.path.join(FATTREE_DIR, "textures", "fat_tree.png")

MIN_TREE_HEIGHT = 4.0
# Leaf cards are cut at half opacity: the apple atlas's alpha is near-binary
# already, so a different threshold only changes edge thickness.
ALPHA_CUTOFF = 0.5
# Golden angle.  Successive trees face directions that never repeat and never
# settle into a pattern, so a tree_row reads as planted rather than stamped.
GOLDEN_TURN = 0.61803398875
SCALE_JITTER = 0.14
TILT_JITTER = math.radians(3.0)


# meshes  — mesh datablocks instanced together at one transform
# height  — the variant's own height, what a target height is divided by
# base_z  — the variant's own ground line, what drops the trunk onto z=0
TreeVariant = namedtuple("TreeVariant", "meshes height base_z")


# --------------------------------------------------------------------------
# import plumbing


def _bake(obj, name):
    """Copy obj's mesh with its rotation and scale applied, but not its
    position.

    Orientation has to be baked: the OBJ importer leaves the apple's Y-up to
    Z-up conversion sitting on the object, so a raw mesh copy would plant the
    tree on its side.  Position must *not* be, because a source file's
    translation is where the artist parked the model in their own scene —
    fattree sits 2.9m up in the air — and baking that in would offset every
    instance by it.  Dropping it costs nothing: `base_z` measures whatever
    ground line the mesh ends up with, and assemble() corrects for it.
    """
    mesh = obj.data.copy()
    mesh.transform(obj.matrix_world.to_3x3().to_4x4())
    mesh.name = name
    mesh.use_fake_user = True
    return mesh


def _measure(meshes):
    """Return (height, base_z) for a variant's meshes in their shared space."""
    zs = [vertex.co.z for mesh in meshes for vertex in mesh.vertices]
    if not zs:
        return 1.0, 0.0
    low, high = min(zs), max(zs)
    return max(high - low, 1e-6), low


def _discard(objects):
    """Remove imported objects along with the meshes they brought in.

    The _bake copies carry a fake user and survive.  Dropping only the objects
    would strand their original meshes at zero users, which in turn keeps the
    source materials and their megabytes of texture alive in the file.
    """
    for obj in objects:
        mesh = obj.data if obj.type == "MESH" else None
        bpy.data.objects.remove(obj, do_unlink=True)
        if mesh is not None and mesh.users == 0:
            bpy.data.meshes.remove(mesh)


def _purge_orphans(before_materials, before_images):
    """Drop the materials and images an import created that nothing now uses.

    Both importers build a material from the source file's own description and
    load its textures.  We replace that material, so without this the .blend
    ships a second, unreferenced copy of every 2k texture.
    """
    for material in set(bpy.data.materials) - before_materials:
        if material.users == 0:
            bpy.data.materials.remove(material)
    for image in set(bpy.data.images) - before_images:
        if image.users == 0:
            bpy.data.images.remove(image)


def _image(path, non_color=False):
    """Load a texture once, keyed by filename so repeat calls share it."""
    key = os.path.basename(path)
    image = bpy.data.images.get(key)
    if image is None:
        image = bpy.data.images.load(path)
        image.name = key
    if non_color:
        image.colorspace_settings.name = "Non-Color"
    return image


def _foliage_material(name, color_path, normal_path=None, alpha_clip=False,
                      roughness=0.72):
    """Principled setup for a textured tree, alpha-clipped when asked.

    The cut-out goes through materials.link_alpha_clip rather than straight
    into the Alpha socket — see that function for why the extra two nodes are
    what makes the crown survive the trip to Cesium.
    """
    material = bpy.data.materials.get(name)
    if material:
        return material

    material = bpy.data.materials.new(name)
    material.use_nodes = True
    nodes = material.node_tree.nodes
    links = material.node_tree.links
    bsdf = next(n for n in nodes if n.type == "BSDF_PRINCIPLED")
    bsdf.inputs["Roughness"].default_value = roughness
    bsdf.inputs["Metallic"].default_value = 0.0

    color_tex = nodes.new("ShaderNodeTexImage")
    color_tex.image = _image(color_path)
    color_tex.location = (-540, 260)
    links.new(color_tex.outputs["Color"], bsdf.inputs["Base Color"])

    if normal_path and os.path.exists(normal_path):
        normal_tex = nodes.new("ShaderNodeTexImage")
        normal_tex.image = _image(normal_path, non_color=True)
        normal_tex.location = (-540, -140)
        normal_map = nodes.new("ShaderNodeNormalMap")
        normal_map.location = (-250, -140)
        links.new(normal_tex.outputs["Color"], normal_map.inputs["Color"])
        links.new(normal_map.outputs["Normal"], bsdf.inputs["Normal"])

    if alpha_clip:
        link_alpha_clip(material, color_tex.outputs["Alpha"], bsdf,
                        cutoff=ALPHA_CUTOFF)
    return material


# --------------------------------------------------------------------------
# loaders — each returns [] when its model is absent, so a clean checkout
# still builds and the caller falls back to procedural trees


def _load_apple():
    """SpeedTree Red Delicious: one mesh, alpha-cut leaf cards, normal-mapped."""
    if not os.path.exists(APPLE_OBJ):
        return []

    # Build ours first: the OBJ importer reuses an already-loaded image when the
    # .mtl resolves to the same file, so the 2k textures land in the file once.
    material = _foliage_material("AppleTree", APPLE_COLOR, APPLE_NORMAL,
                                 alpha_clip=True, roughness=0.68)

    before_objects = set(bpy.data.objects)
    before_materials = set(bpy.data.materials)
    before_images = set(bpy.data.images)
    bpy.ops.wm.obj_import(filepath=APPLE_OBJ)
    imported = [obj for obj in set(bpy.data.objects) - before_objects
                if obj.type == "MESH"]
    if not imported:
        return []

    meshes = []
    for index, obj in enumerate(imported):
        mesh = _bake(obj, "AppleTree_%02d" % index)
        mesh.materials.clear()
        mesh.materials.append(material)
        meshes.append(mesh)
    height, base_z = _measure(meshes)

    _discard(imported)
    _purge_orphans(before_materials, before_images)
    return [TreeVariant(meshes, height, base_z)]


def _load_fattree():
    """Low-poly cartoon tree: opaque geometry, one diffuse texture.

    The crown is real geometry and the texture's alpha is 1.0 everywhere, so
    unlike the apple this needs no cut-out — and no normal map, which the
    source does not ship.
    """
    if not os.path.exists(FATTREE_BLEND):
        return []

    before_objects = set(bpy.data.objects)
    before_materials = set(bpy.data.materials)
    before_images = set(bpy.data.images)
    try:
        bpy.ops.wm.append(
            filepath=FATTREE_BLEND + "/Object/fattree",
            directory=FATTREE_BLEND + "/Object/",
            files=[{"name": "fattree"}],
            link=False,
        )
    except RuntimeError:
        return []
    imported = [obj for obj in set(bpy.data.objects) - before_objects
                if obj.type == "MESH"]

    obj = bpy.data.objects.get("fattree")
    if obj is None:
        _discard(imported)
        _purge_orphans(before_materials, before_images)
        return []

    material = _foliage_material("FatTree", FATTREE_COLOR, alpha_clip=False,
                                 roughness=0.85)
    mesh = _bake(obj, "FatTree")
    mesh.materials.clear()
    mesh.materials.append(material)
    height, base_z = _measure([mesh])

    _discard(imported)
    _purge_orphans(before_materials, before_images)
    return [TreeVariant([mesh], height, base_z)]


LOADERS = {
    "apple": _load_apple,
    "fattree": _load_fattree,
}

# The styles this module can serve, for the CLI to validate against.
MODEL_STYLES = tuple(LOADERS)


def assemble(positions, collection, style="apple"):
    """Place instanced trees of `style` at `positions`.

    positions is a list of (x, y, height) tuples as gathered by the two
    tree-collecting loops (point nodes + tree_row samples).  `height` is the
    OSM height where tagged and a constant default otherwise, which means every
    sample along one tree_row arrives with an identical value — the per-index
    jitter below is what stops a row of forty from reading as one tree stamped
    forty times.

    Returns the number of trees placed, or 0 when the style's model is absent
    or unknown, which is the caller's signal to fall back to procedural trees.
    """
    loader = LOADERS.get(style)
    if loader is None:
        return 0
    variants = loader()
    if not variants:
        return 0

    for index, (x, y, target_height) in enumerate(positions):
        variant = variants[index % len(variants)]
        # Irrational periods stand in for an RNG: no repeat over any realistic
        # tree count, and a pure function of the index, so rebuilding an area
        # plants the identical forest.
        scale_wobble = 1.0 + SCALE_JITTER * math.sin(index * 2.399963)
        target = max(MIN_TREE_HEIGHT, target_height) * scale_wobble
        factor = target / variant.height
        yaw = ((index * GOLDEN_TURN) % 1.0) * math.tau
        tilt_x = TILT_JITTER * math.sin(index * 1.114517)
        tilt_y = TILT_JITTER * math.cos(index * 0.927295)
        # Scaled and negated, the variant's own ground line drops the trunk
        # onto z=0 whatever the source file used as its origin.
        z = -variant.base_z * factor

        for slot, mesh in enumerate(variant.meshes):
            obj = bpy.data.objects.new(
                "Tree_%s_%04d_%d" % (style, index, slot), mesh)
            obj.location = (x, y, z)
            obj.scale = (factor, factor, factor)
            obj.rotation_euler = (tilt_x, tilt_y, yaw)
            collection.objects.link(obj)

    tris = 0
    for variant in variants:
        for mesh in variant.meshes:
            mesh.calc_loop_triangles()
            tris += len(mesh.loop_triangles)
    print("Tree style %r: %d variants, %d tris per instance, %d planted"
          % (style, len(variants), tris // max(1, len(variants)), len(positions)))
    return len(positions)