main
Caleb Publish 0.1.1 83c9715 1d ago
#!/usr/bin/env python3
"""Give a step2glb GLB physical materials (step2glb writes base colours only, so glTF's defaults, metallic 1 and
roughness 1, render every part as rough metal). Adapted from adom/stm32g431cbt6 provenance/tools/kicad-generator/
materials.py, keyed by base colour instead of node name because KiCad-generator STEPs are one fused solid.

usage: glbmat.py in.glb out.glb RULE [RULE ...]
  RULE = "r,g,b=metallic,roughness,Name[,alpha]"  (r,g,b = the GLB's linear baseColorFactor, matched within 0.02;
         alpha < 1 makes the material alphaMode BLEND, e.g. a water-clear LED lens)
  The Adom.LaserEtch node's material always becomes metallic 0, roughness 0.7 (colour kept).
Geometry, node transforms, quantization and colours are untouched; materials that match no rule are reported."""
import sys
from pygltflib import GLTF2, PbrMetallicRoughness

def main(src, out, rules):
    R = []
    for r in rules:
        col, spec = r.split('=')
        m, ro, rest = spec.split(',', 2)
        name, _, alpha = rest.partition(',')
        R.append(([float(v) for v in col.split(',')], float(m), float(ro), name, float(alpha) if alpha else 1.0))
    g = GLTF2().load(src)
    gen = g.asset.generator
    etch = set()
    for n in g.nodes:
        if n.mesh is not None and (n.name or '').startswith('Adom.LaserEtch'):
            etch |= {p.material for p in g.meshes[n.mesh].primitives}
    unmatched = []
    for i, m in enumerate(g.materials):
        pbr = m.pbrMetallicRoughness or PbrMetallicRoughness()
        c = list(pbr.baseColorFactor or [1, 1, 1, 1])
        if i in etch:
            pbr.metallicFactor, pbr.roughnessFactor = 0.0, 0.7; m.name = 'Adom.LaserEtch'
        else:
            hit = next((x for x in R if all(abs(a - b) <= 0.02 for a, b in zip(x[0], c[:3]))), None)
            if hit is None:
                unmatched.append((m.name, [round(v, 3) for v in c])); continue
            pbr.metallicFactor, pbr.roughnessFactor, m.name = hit[1], hit[2], hit[3]
            if hit[4] < 1.0:
                c[3] = hit[4]; pbr.baseColorFactor = c; m.alphaMode = 'BLEND'; m.doubleSided = True
        m.pbrMetallicRoughness = pbr
    g.asset.generator = gen  # keep the glTF-Transform stamp; the material pass is recorded in extras
    ex = g.asset.extras if isinstance(g.asset.extras, dict) else {}
    ex['adomMaterials'] = 'PBR metallic/roughness set by base colour (Adom glbmat.py); geometry untouched'
    g.asset.extras = ex
    g.save(out)
    # pygltflib stamps its own generator on save: put the glTF-Transform stamp back in the JSON chunk
    import json, struct
    b = open(out, 'rb').read(); cl = struct.unpack('<I', b[12:16])[0]
    j = json.loads(b[20:20 + cl]); j['asset']['generator'] = gen
    js = json.dumps(j, separators=(',', ':')).encode(); js += b' ' * ((4 - len(js) % 4) % 4)
    rest = b[20 + cl:]
    body = struct.pack('<I', len(js)) + b'JSON' + js + rest
    open(out, 'wb').write(b[:8] + struct.pack('<I', 12 + len(body)) + body)
    print(f'OK {out}: ' + '; '.join(f'{m.name} m{m.pbrMetallicRoughness.metallicFactor} r{m.pbrMetallicRoughness.roughnessFactor}' for m in g.materials))
    if unmatched:
        sys.exit(f'ERROR: materials matched no rule: {unmatched}')

if __name__ == '__main__':
    main(sys.argv[1], sys.argv[2], sys.argv[3:])