osmWorkflow/blender/osmassets/materials.py

"""Blender material construction.

Requires `bpy`; only runs inside Blender.  The catalog (`osmassets.catalog`)
declares *what* a material is, this module builds it — that split is what keeps
the catalog importable by plain Python, and by anything else that wants to read
the scene's material definitions without launching Blender.
"""

import os

import bpy


TEXTURE_ROOT = os.path.abspath(os.path.join(
    os.path.dirname(os.path.abspath(__file__)), "..", "..",
    "assets", "textures", "polyhaven"
))


def principled_bsdf(material):
    if not material.use_nodes:
        return None
    for node in material.node_tree.nodes:
        if node.type == "BSDF_PRINCIPLED":
            return node
    return None


def make_material(name, color, roughness=0.8, metallic=0.0):
    material = bpy.data.materials.get(name) or bpy.data.materials.new(name)
    material.diffuse_color = (*color, 1.0)
    material.use_nodes = True
    bsdf = principled_bsdf(material)
    if bsdf:
        bsdf.inputs["Base Color"].default_value = (*color, 1.0)
        bsdf.inputs["Roughness"].default_value = roughness
        bsdf.inputs["Metallic"].default_value = metallic
    return material


def add_procedural_surface(material, colors, scale=2.0, detail=2.0, bump_strength=0.08,
                           object_space=False):
    nodes = material.node_tree.nodes
    links = material.node_tree.links
    bsdf = principled_bsdf(material)
    if not bsdf:
        return
    noise = nodes.new("ShaderNodeTexNoise")
    noise.inputs["Scale"].default_value = scale
    noise.inputs["Detail"].default_value = detail
    noise.inputs["Roughness"].default_value = 0.65
    texcoord = nodes.new("ShaderNodeTexCoord")
    ramp = nodes.new("ShaderNodeValToRGB")
    ramp.color_ramp.elements[0].color = (*colors[0], 1.0)
    ramp.color_ramp.elements[1].color = (*colors[1], 1.0)
    bump = nodes.new("ShaderNodeBump")
    bump.inputs["Strength"].default_value = bump_strength
    bump.inputs["Distance"].default_value = 0.12
    # "Generated" normalises across the object bounding box, so on a mesh that
    # spans the whole scene the noise stretches to tens of metres and vanishes.
    # Object space keeps the scale in metres, which is what foliage needs.
    source = "Object" if object_space else "Generated"
    links.new(texcoord.outputs[source], noise.inputs["Vector"])
    links.new(noise.outputs["Fac"], ramp.inputs["Fac"])
    links.new(ramp.outputs["Color"], bsdf.inputs["Base Color"])
    links.new(noise.outputs["Fac"], bump.inputs["Height"])
    links.new(bump.outputs["Normal"], bsdf.inputs["Normal"])


def tint_base_color(material, tint, factor):
    """Mix an existing material's base colour toward `tint`.

    Imported assets arrive with their own diffuse texture wired up. Rather than
    replacing it — which throws away the leaf detail — this splices a mix node
    in front of the Base Color input so the texture survives at (1 - factor).
    """
    if factor <= 0.0 or not material.use_nodes:
        return
    bsdf = principled_bsdf(material)
    if not bsdf:
        return
    nodes = material.node_tree.nodes
    links = material.node_tree.links
    base = bsdf.inputs["Base Color"]
    tint_node = nodes.new("ShaderNodeRGB")
    tint_node.outputs["Color"].default_value = (*tint, 1.0)
    mix = nodes.new("ShaderNodeMixRGB")
    mix.blend_type = "MIX"
    mix.inputs["Fac"].default_value = factor
    if base.is_linked:
        # Capture the upstream socket before relinking; Blender drops the old
        # link as soon as the input takes a new one.
        links.new(base.links[0].from_socket, mix.inputs[1])
    else:
        mix.inputs[1].default_value = base.default_value
    links.new(tint_node.outputs["Color"], mix.inputs[2])
    links.new(mix.outputs["Color"], base)


def make_textured_material(name, diffuse_file, normal_file, roughness,
                           scale, normal_is_bump=False, metallic=0.0,
                           tint=None, tint_factor=0.0):
    diffuse_path = os.path.join(TEXTURE_ROOT, diffuse_file)
    normal_path = os.path.join(TEXTURE_ROOT, normal_file)
    if not os.path.exists(diffuse_path) or not os.path.exists(normal_path):
        return make_material(name, (0.5, 0.5, 0.5), roughness, metallic)

    material = make_material(name, (0.5, 0.5, 0.5), roughness, metallic)
    nodes = material.node_tree.nodes
    links = material.node_tree.links
    bsdf = principled_bsdf(material)
    if not bsdf:
        return material
    texcoord = nodes.new("ShaderNodeTexCoord")
    mapping = nodes.new("ShaderNodeMapping")
    mapping.inputs["Scale"].default_value = (scale, scale, scale)
    diffuse = nodes.new("ShaderNodeTexImage")
    diffuse.image = bpy.data.images.load(diffuse_path, check_existing=True)
    diffuse.extension = "REPEAT"
    normal = nodes.new("ShaderNodeTexImage")
    normal.image = bpy.data.images.load(normal_path, check_existing=True)
    normal.image.colorspace_settings.name = "Non-Color"
    normal.extension = "REPEAT"
    links.new(texcoord.outputs["Generated"], mapping.inputs["Vector"])
    links.new(mapping.outputs["Vector"], diffuse.inputs["Vector"])
    links.new(mapping.outputs["Vector"], normal.inputs["Vector"])
    if tint and tint_factor > 0.0:
        tint_node = nodes.new("ShaderNodeRGB")
        tint_node.outputs["Color"].default_value = (*tint, 1.0)
        mix = nodes.new("ShaderNodeMixRGB")
        mix.blend_type = "MIX"
        mix.inputs["Fac"].default_value = tint_factor
        links.new(diffuse.outputs["Color"], mix.inputs[1])
        links.new(tint_node.outputs["Color"], mix.inputs[2])
        links.new(mix.outputs["Color"], bsdf.inputs["Base Color"])
    else:
        links.new(diffuse.outputs["Color"], bsdf.inputs["Base Color"])
    if normal_is_bump:
        bump = nodes.new("ShaderNodeBump")
        bump.inputs["Strength"].default_value = 0.22
        bump.inputs["Distance"].default_value = 0.12
        links.new(normal.outputs["Color"], bump.inputs["Height"])
        links.new(bump.outputs["Normal"], bsdf.inputs["Normal"])
    else:
        normal_map = nodes.new("ShaderNodeNormalMap")
        normal_map.inputs["Strength"].default_value = 0.52
        links.new(normal.outputs["Color"], normal_map.inputs["Color"])
        links.new(normal_map.outputs["Normal"], bsdf.inputs["Normal"])
    return material


def link_alpha_clip(material, alpha_output, bsdf, cutoff=0.5):
    """Wire a texture's alpha into `bsdf` as a hard cut-out.

    The obvious wiring — alpha straight into the Alpha socket — is wrong for
    anything destined for glTF.  Blender 4.2 stopped deriving a material's
    alpha mode from `blend_method` (still writable, now a no-op: setting 'CLIP'
    reads back 'HASHED') and made the exporter infer it from the node tree
    instead.  It recognises exactly a few shapes; a bare link is not one of
    them, and falls through to alphaMode=BLEND.  Foliage exported as BLEND
    makes Cesium depth-sort thousands of leaf quads it cannot order correctly.

    So build the shape the exporter looks for — `1 - (alpha < cutoff)` — which
    it reads back as alphaMode=MASK with this cutoff.  EEVEE gets the same
    thing for free: alpha is 0 or 1 by the time it reaches the BSDF, so the
    viewport shows the crisp cut-out Cesium will, not a dithered approximation.

    See the exporter's `detect_alpha_clip` in
    scripts/addons_core/io_scene_gltf2/blender/exp/material/search_node_tree.py.
    """
    nodes = material.node_tree.nodes
    links = material.node_tree.links

    less_than = nodes.new("ShaderNodeMath")
    less_than.operation = "LESS_THAN"
    less_than.location = (-60, 320)
    less_than.inputs[1].default_value = cutoff

    invert = nodes.new("ShaderNodeMath")
    invert.operation = "SUBTRACT"
    invert.location = (110, 320)
    invert.inputs[0].default_value = 1.0

    links.new(alpha_output, less_than.inputs[0])
    links.new(less_than.outputs["Value"], invert.inputs[1])
    links.new(invert.outputs["Value"], bsdf.inputs["Alpha"])

    # EEVEE Next takes its cut-out handling from surface_render_method, not
    # from blend_method. 'DITHERED' still casts a leaf-shaped shadow;
    # 'BLENDED' does not.
    material.surface_render_method = "DITHERED"
    material.alpha_threshold = cutoff
    # Leaf cards are single-sided quads seen from both sides; culling
    # backfaces would empty out half of every crown.
    material.use_backface_culling = False


def from_spec(spec):
    """Build a material from a `catalog.MATERIALS` entry."""
    if spec["kind"] == "textured":
        return make_textured_material(
            spec["name"], spec["diffuse"], spec["normal"],
            roughness=spec.get("roughness", 0.8), scale=spec["scale"],
            normal_is_bump=spec.get("normal_is_bump", False),
            metallic=spec.get("metallic", 0.0),
            tint=spec.get("tint"), tint_factor=spec.get("tint_factor", 0.0))

    material = make_material(spec["name"], spec["color"],
                             spec.get("roughness", 0.8),
                             spec.get("metallic", 0.0))
    procedural = spec.get("procedural")
    if procedural:
        add_procedural_surface(material, procedural["colors"],
                               scale=procedural["scale"],
                               detail=procedural["detail"],
                               bump_strength=procedural["bump_strength"],
                               object_space=procedural.get("object_space", False))
    return material