GJM1555C1H1R0WB01D - Murata

"""Orthographic renderer for small component STEPs (face colours kept).

render(step, view, ppmm) -> PIL RGBA image, model origin at the image centre.
views: top (look down -Z, +X right, +Y up), bottom (look up +Z, +X right, +Y down,
i.e. the part flipped about X), front (look +Y, +X right, +Z up),
end (look -X, +Y right... from +X end: -Y right, +Z up).
"""
import sys
import numpy as np
from PIL import Image, ImageDraw
from OCP.STEPCAFControl import STEPCAFControl_Reader
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_LabelSequence
from OCP.TopExp import TopExp_Explorer
from OCP.TopAbs import TopAbs_FACE, TopAbs_REVERSED
from OCP.TopoDS import TopoDS
from OCP.Quantity import Quantity_Color
from OCP.BRepMesh import BRepMesh_IncrementalMesh
from OCP.BRep import BRep_Tool
from OCP.TopLoc import TopLoc_Location


def load_tris(path):
    doc = TDocStd_Document(TCollection_ExtendedString("x"))
    r = STEPCAFControl_Reader(); r.SetColorMode(True); r.SetNameMode(True)
    r.ReadFile(path); r.Transfer(doc)
    st = XCAFDoc_DocumentTool.ShapeTool_s(doc.Main())
    ct = XCAFDoc_DocumentTool.ColorTool_s(doc.Main())
    labs = TDF_LabelSequence(); st.GetFreeShapes(labs)
    out = []
    leaves = []
    def walk(lab, loc, pcol):
        c = Quantity_Color(); col = pcol
        for t in (XCAFDoc_ColorType.XCAFDoc_ColorSurf, XCAFDoc_ColorType.XCAFDoc_ColorGen):
            if ct.GetColor_s(lab, t, c):
                col = (c.Red(), c.Green(), c.Blue()); break
        if st.IsReference_s(lab):
            ref = type(lab)(); st.GetReferredShape_s(lab, ref)
            walk(ref, loc.Multiplied(st.GetLocation_s(lab)), col); return
        kids = TDF_LabelSequence(); st.GetComponents_s(lab, kids)
        if kids.Length() > 0:
            for i in range(1, kids.Length() + 1): walk(kids.Value(i), loc, col)
            return
        leaves.append((st.GetShape_s(lab).Moved(loc), col))
    for li in range(1, labs.Length() + 1):
        walk(labs.Value(li), TopLoc_Location(), None)
    for shp, pcol in leaves:
        BRepMesh_IncrementalMesh(shp, 0.002, False, 0.2, True)
        e = TopExp_Explorer(shp, TopAbs_FACE)
        while e.More():
            f = TopoDS.Face_s(e.Current())
            c = Quantity_Color(); col = None
            for t in (XCAFDoc_ColorType.XCAFDoc_ColorSurf, XCAFDoc_ColorType.XCAFDoc_ColorGen):
                if ct.GetColor(f, t, c):
                    col = (c.Red(), c.Green(), c.Blue()); break
            if col is None:
                col = pcol or (0.6, 0.6, 0.6)
            loc = TopLoc_Location()
            tri = BRep_Tool.Triangulation_s(f, loc)
            if tri is not None:
                trsf = loc.Transformation()
                pts = np.array([(lambda p: (p.X(), p.Y(), p.Z()))(tri.Node(i).Transformed(trsf)) for i in range(1, tri.NbNodes() + 1)])
                rev = f.Orientation() == TopAbs_REVERSED
                for i in range(1, tri.NbTriangles() + 1):
                    a, b, cc = tri.Triangle(i).Get()
                    if rev:
                        b, cc = cc, b
                    out.append((pts[[a - 1, b - 1, cc - 1]], col))
            e.Next()
    return out


VIEWS = {
    # (right axis vector, up axis vector, toward-viewer vector)
    "top": ((1, 0, 0), (0, 1, 0), (0, 0, 1)),
    "bottom": ((1, 0, 0), (0, -1, 0), (0, 0, -1)),
    "front": ((1, 0, 0), (0, 0, 1), (0, -1, 0)),
    "end": ((0, -1, 0), (0, 0, 1), (1, 0, 0)),
}


def render(tris, view, ppmm, size=None, shade=True, ss=3):
    R, U, V = [np.array(v, float) for v in VIEWS[view]]
    allp = np.concatenate([t[0] for t in tris])
    ext = max(np.abs(allp @ R).max(), np.abs(allp @ U).max())
    if size is None:
        size = (int(2 * ext * ppmm + 40), int(2 * ext * ppmm + 40))
    W, H = size[0] * ss, size[1] * ss
    img = Image.new("RGBA", (W, H), (0, 0, 0, 0))
    d = ImageDraw.Draw(img)
    items = []
    light = np.array([0.3, 0.4, 0.85]); light /= np.linalg.norm(light)
    for p, col in tris:
        n = np.cross(p[1] - p[0], p[2] - p[0])
        nn = np.linalg.norm(n)
        if nn == 0:
            continue
        n /= nn
        if n @ V <= 1e-6:
            continue  # back face
        depth = (p @ V).mean()
        lv = 0.55 + 0.45 * abs(n @ (R * light[0] + U * light[1] + V * light[2])) if shade else 1.0
        rgb = tuple(int(255 * min(1, c * lv)) for c in col)
        xy = [(W / 2 + (q @ R) * ppmm * ss, H / 2 - (q @ U) * ppmm * ss) for q in p]
        items.append((depth, xy, rgb))
    items.sort(key=lambda x: x[0])
    for _, xy, rgb in items:
        d.polygon(xy, fill=rgb + (255,), outline=rgb + (255,))
    return img.resize(size, Image.LANCZOS)


if __name__ == "__main__":
    tris = load_tris(sys.argv[1])
    out = sys.argv[2]
    ppmm = float(sys.argv[3]) if len(sys.argv) > 3 else 300
    imgs = [render(tris, v, ppmm) for v in ("top", "bottom", "front", "end")]
    w = sum(i.width for i in imgs); h = max(i.height for i in imgs)
    sheet = Image.new("RGBA", (w, h), (255, 255, 255, 255)); x = 0
    for i in imgs:
        sheet.alpha_composite(i, (x, 0)); x += i.width
    sheet.save(out)