Green SMD LED 0603 - Inolux IN-S63BTG

#!/usr/bin/env python3
"""Colour an LED STEP's lens (and give body/terminals colours where they have none) with XCAF.
Geometry, topology, product names and the file's leading comment block are kept; only face colours
change. Every coloured face gets BOTH XCAFDoc_ColorSurf and XCAFDoc_ColorGen. Colours are given in
sRGB (as written to COLOUR_RGB in the file) and passed to OCCT as Quantity_TOC_sRGB.

Families (classification is by the source model's own colours + geometry, never by guesswork):
  kicad-chip  KiCad LED_0402_1005Metric / LED_0603_1608Metric: one solid; ivory faces above the
              substrate top form the moulded lens block (the clear lens); ivory faces at/below it are
              the PCB substrate; gold = terminals; green = underside cathode mark.
  template    a colourless copy of a kicad-chip model with the same face order (IN-S42ATR is the KiCad
              0402 model scaled to 0.45 mm in Z and re-exported as B-splines): classes are copied from
              --template by face index after checking every face's X/Y extent matches.
  ws5050      KiCad LED_WS2812B_PLCC4_5.0x5.0mm: ivory = body, blue disc = cavity/die window, grey = leads.

Usage: colour_led.py SRC OUT --family F --lens R,G,B [--body R,G,B] [--term R,G,B] [--mark R,G,B]
                     [--template TEMPLATE.step] [--report report.json]"""
import sys, json, argparse, re
from OCP.STEPCAFControl import STEPCAFControl_Reader, STEPCAFControl_Writer
from OCP.STEPControl import STEPControl_AsIs
from OCP.TDocStd import TDocStd_Document
from OCP.TCollection import TCollection_ExtendedString
from OCP.XCAFDoc import XCAFDoc_DocumentTool, XCAFDoc_ColorType
from OCP.TDF import TDF_ChildIterator, TDF_Label
from OCP.Quantity import Quantity_Color, Quantity_TOC_sRGB
from OCP.Bnd import Bnd_Box
from OCP.BRepBndLib import BRepBndLib
from OCP.TopExp import TopExp_Explorer
from OCP.TopAbs import TopAbs_FACE

SURF, GEN = XCAFDoc_ColorType.XCAFDoc_ColorSurf, XCAFDoc_ColorType.XCAFDoc_ColorGen
KICAD_IVORY, KICAD_GOLD, KICAD_GREEN = (0.895, 0.891, 0.813), (0.859, 0.738, 0.496), (0.34, 0.68, 0.445)

def load(path):
    doc = TDocStd_Document(TCollection_ExtendedString("d"))
    r = STEPCAFControl_Reader(); r.SetColorMode(True); r.SetNameMode(True); r.SetLayerMode(True)
    assert r.ReadFile(path) == 1, path
    r.Transfer(doc)
    st = XCAFDoc_DocumentTool.ShapeTool_s(doc.Main()); ct = XCAFDoc_DocumentTool.ColorTool_s(doc.Main())
    free = []
    it = TDF_ChildIterator(st.Label(), False)
    while it.More():
        l = it.Value(); it.Next()
        if st.IsFree_s(l) and st.IsTopLevel(l): free.append(l)
    assert len(free) == 1 and not st.IsAssembly_s(free[0]), 'expected one simple part'
    return doc, st, ct, free[0]

def faces(st, ct, lab):
    out = []; ex = TopExp_Explorer(st.GetShape_s(lab), TopAbs_FACE)
    while ex.More():
        f = ex.Current(); ex.Next()
        b = Bnd_Box(); BRepBndLib.Add_s(f, b); mn, mx = b.CornerMin(), b.CornerMax()
        c = Quantity_Color(); col = None
        for t in (SURF, GEN):
            if ct.GetColor(f, t, c):
                r, g, bb = c.Values(Quantity_TOC_sRGB); col = (round(r, 3), round(g, 3), round(bb, 3)); break
        out.append(dict(face=f, col=col, mn=(mn.X(), mn.Y(), mn.Z()), mx=(mx.X(), mx.Y(), mx.Z())))
    return out

def near(a, b, tol=0.02): return a is not None and all(abs(x - y) < tol for x, y in zip(a, b))

def classify_kicad_chip(fs):
    ivory = [f for f in fs if near(f['col'], KICAD_IVORY)]
    ztop = max(f['mx'][2] for f in fs)
    # substrate top = the z of the small horizontal ivory faces beside the lens block
    zsub = min(f['mn'][2] for f in ivory if f['mn'][2] > 0.05 * ztop)
    cls = []
    for f in fs:
        if near(f['col'], KICAD_IVORY):
            lens = f['mn'][2] >= zsub - 1e-4 and f['mx'][2] > zsub + 0.02 * ztop
            cls.append('lens' if lens else 'substrate')
        elif near(f['col'], KICAD_GOLD): cls.append('terminal')
        elif near(f['col'], KICAD_GREEN): cls.append('cathode_mark')
        else: raise SystemExit(f'unexpected colour {f["col"]}')
    return cls, dict(z_top_mm=round(ztop, 4), z_substrate_top_mm=round(zsub, 4))

def classify_ws5050(fs):
    cls = []
    for f in fs:
        if near(f['col'], KICAD_IVORY): cls.append('body')
        elif near(f['col'], (0.15, 0.4, 0.73), 0.08) or near(f['col'], (0.02, 0.13, 0.49), 0.05): cls.append('window')
        else: cls.append('terminal')
    return cls, {}

def main():
    ap = argparse.ArgumentParser(); ap.add_argument('src'); ap.add_argument('out')
    ap.add_argument('--family', required=True, choices=['kicad-chip', 'template', 'ws5050'])
    ap.add_argument('--template'); ap.add_argument('--report')
    for k in ('lens', 'body', 'term', 'mark'): ap.add_argument('--' + k)
    a = ap.parse_args()
    rgb = lambda s: tuple(float(v) for v in s.split(',')) if s else None
    doc, st, ct, lab = load(a.src); fs = faces(st, ct, lab)
    info = {}
    if a.family == 'kicad-chip': cls, info = classify_kicad_chip(fs)
    elif a.family == 'ws5050': cls, info = classify_ws5050(fs)
    else:
        _, st2, ct2, lab2 = load(a.template); tf = faces(st2, ct2, lab2)
        assert len(tf) == len(fs), (len(tf), len(fs))
        for i, (f, t) in enumerate(zip(fs, tf)):   # same face order: X/Y extents must agree
            assert all(abs(f['mn'][k] - t['mn'][k]) < 2e-3 and abs(f['mx'][k] - t['mx'][k]) < 2e-3 for k in (0, 1)), i
        cls, info = classify_kicad_chip(tf); info['template'] = a.template.split('/')[-1]
        info['z_scale_vs_template'] = round(max(f['mx'][2] for f in fs) / max(f['mx'][2] for f in tf), 4)
    colours = {'lens': rgb(a.lens)}
    if a.family == 'ws5050':
        colours.update(body=rgb(a.body), window=rgb(a.lens), terminal=rgb(a.term)); colours.pop('lens')
    else:
        colours.update(substrate=rgb(a.body) or KICAD_IVORY, terminal=rgb(a.term) or KICAD_GOLD,
                       cathode_mark=rgb(a.mark) or KICAD_GREEN)
    counts = {}; changed = 0
    for f, c in zip(fs, cls):
        want = colours[c]; counts[c] = counts.get(c, 0) + 1
        if want is None or (f['col'] is not None and near(f['col'], want, 1e-3)): continue
        sub = TDF_Label()
        if not st.FindSubShape(lab, f['face'], sub): sub = st.AddSubShape(lab, f['face'])
        q = Quantity_Color(want[0], want[1], want[2], Quantity_TOC_sRGB)
        ct.SetColor(sub, q, SURF); ct.SetColor(sub, q, GEN); changed += 1
    w = STEPCAFControl_Writer(); w.SetColorMode(True); w.SetNameMode(True); w.SetLayerMode(True)
    w.Transfer(doc, STEPControl_AsIs); assert w.Write(a.out) == 1
    # keep the source's leading /* ... */ comment block (KiCad licence/attribution) in the HEADER
    src_txt = open(a.src, encoding='latin-1').read(); m = re.search(r'HEADER;\s*(/\*.*?\*/)', src_txt, re.S)
    if m:
        t = open(a.out, encoding='latin-1').read().replace('HEADER;', 'HEADER;\n' + m.group(1), 1)
        open(a.out, 'w', encoding='latin-1').write(t)
    rep = dict(src=a.src.split('/')[-1], out=a.out.split('/')[-1], family=a.family, faces=len(fs),
               faces_by_class=counts, faces_recoloured=changed,
               colours_srgb={k: v for k, v in colours.items()}, classification=info, kept_comment_block=bool(m))
    if a.report: json.dump(rep, open(a.report, 'w'), indent=1)
    print(json.dumps(rep))

if __name__ == '__main__': main()