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 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