Wrist Health Tracker

#!/usr/bin/env python3
"""Offscreen OpenGL renders of the enclosure (pyrender + EGL): exploded, assembled top, skin
side, a section through the optical stack, the cell in its pocket before the board, and the dock. Run with PYOPENGL_PLATFORM=egl."""
import os
os.environ.setdefault("PYOPENGL_PLATFORM", "egl")
import numpy as np
import cadquery as cq
import trimesh
import pyrender
from PIL import Image
from enclosure import OUT, BOARD_STEP, J1_PADS, PAD_Z, TONGUE_TOP, cell as cell_solid, cell_leads

COL = {"body": (0.17, 0.18, 0.21), "back_plate": (0.07, 0.07, 0.08), "pipe_pd": (0.70, 0.88, 1.0),
       "pipe_led": (0.70, 0.88, 1.0), "thermal_post": (0.78, 0.79, 0.82), "led_rod_1": (0.85, 0.95, 1.0),
       "led_rod_2": (0.85, 0.95, 1.0), "button_cap": (0.33, 0.34, 0.38), "dock_cradle": (0.55, 0.62, 0.72),
       "board": (0.08, 0.38, 0.20), "parts": (0.22, 0.22, 0.24), "module": (0.80, 0.78, 0.72), "cell": (0.62, 0.66, 0.74),
       "lead": (0.80, 0.12, 0.10)}


def tm(shape, tol=0.03):
    v, t = shape.tessellate(tol, 0.2)
    return trimesh.Trimesh(np.array([(p.x, p.y, p.z) for p in v]), np.array(t), process=False)


def add(scene, shape, color, offset=(0, 0, 0), alpha=1.0):
    m = tm(shape)
    m.apply_translation(offset)
    mat = pyrender.MetallicRoughnessMaterial(baseColorFactor=(*color, alpha), metallicFactor=0.05,
                                             roughnessFactor=0.55, alphaMode="BLEND" if alpha < 1 else "OPAQUE")
    scene.add(pyrender.Mesh.from_trimesh(m, material=mat, smooth=False))


def look_at(eye, target, up=(0, 0, 1)):
    eye, target, up = (np.array(a, dtype=float) for a in (eye, target, up))
    f = target - eye; f /= np.linalg.norm(f)
    s = np.cross(f, up); s /= np.linalg.norm(s)
    u = np.cross(s, f)
    m = np.eye(4); m[:3, 0], m[:3, 1], m[:3, 2], m[:3, 3] = s, u, -f, eye
    return m


def shoot(scene, eye, target, path, size=(1500, 1100), fov=0.55, up=(0, 0, 1)):
    pose = look_at(eye, target, up)
    scene.add(pyrender.PerspectiveCamera(yfov=fov), pose=pose)
    scene.add(pyrender.DirectionalLight(intensity=3.2), pose=pose)
    scene.add(pyrender.DirectionalLight(intensity=1.6), pose=look_at(np.array(target) + [30, 40, 60], target))
    scene.add(pyrender.DirectionalLight(intensity=1.0), pose=look_at(np.array(target) + [-40, -20, -50], target))
    r = pyrender.OffscreenRenderer(*size)
    c, d = r.render(scene, flags=pyrender.RenderFlags.SKIP_CULL_FACES)
    r.delete()
    c = c.copy()
    c[d == 0] = 255                       # EGL ignores bg_color here: paint the background white
    Image.fromarray(c).save(path)


def new_scene():
    return pyrender.Scene(bg_color=(1, 1, 1, 1), ambient_light=(0.35, 0.35, 0.35))


def board_parts(board):
    out = []
    for s in board.Solids():
        bb = s.BoundingBox()
        if bb.xlen > 30 and bb.ylen > 28:
            out.append((s, COL["board"]))
        elif bb.zmax > 3.5:
            out.append((s, COL["module"]))
        else:
            out.append((s, COL["parts"]))
    return out


def main():
    names = ["body", "back_plate", "pipe_pd", "pipe_led", "thermal_post", "led_rod_1", "led_rod_2", "button_cap", "dock_cradle"]
    P = {n: cq.importers.importStep(os.path.join(OUT, f"{n}.step")).val() for n in names}
    board = cq.importers.importStep(BOARD_STEP).val()
    cell = cell_solid().val()
    leads = cell_leads()
    bparts = board_parts(board)
    ctr = (16, -15, 0)

    # exploded, skin side up
    sc = new_scene()
    add(sc, P["body"], COL["body"], (0, 0, -14))
    add(sc, cell, COL["cell"], (0, 0, -8))
    for s, c in bparts:
        add(sc, s, c)
    add(sc, P["back_plate"], COL["back_plate"], (0, 0, 12))
    for n in ("pipe_pd", "pipe_led", "thermal_post", "led_rod_1", "led_rod_2"):
        add(sc, P[n], COL[n], (0, 0, 22))
    add(sc, P["button_cap"], COL["button_cap"], (-7, 0, 0))
    shoot(sc, (70, -95, 75), (16, -15, 2), os.path.join(OUT, "render_exploded.png"), size=(1400, 1500), fov=0.62)

    # assembled, outside of the wrist (battery side), with the board hidden
    sc = new_scene()
    for n in ("body", "back_plate", "button_cap"):
        add(sc, P[n], COL[n])
    shoot(sc, (55, -75, -55), ctr, os.path.join(OUT, "render_top.png"))

    # skin side
    sc = new_scene()
    for n in ("body", "back_plate", "pipe_pd", "pipe_led", "thermal_post", "led_rod_1", "led_rod_2", "button_cap"):
        add(sc, P[n], COL[n])
    shoot(sc, (45, -70, 70), ctr, os.path.join(OUT, "render_skin.png"))

    # section at the PPG (cut away x > 16)
    cut = cq.Workplane("XY").box(100, 100, 100).translate((16 + 50, 0, 0)).val()
    sc = new_scene()
    for n in ("body", "back_plate", "pipe_pd", "pipe_led"):
        add(sc, P[n].cut(cut), COL[n])
    for s, c in bparts:
        try:
            add(sc, s.cut(cut), c)
        except Exception:
            pass
    add(sc, cell.cut(cut), COL["cell"])
    shoot(sc, (60, -24, 2), (16, -24, -0.5), os.path.join(OUT, "render_section.png"), fov=0.42)

    # build step: the cell taped into its pocket, leads laid in the channel to the BT1 well, no board yet
    sc = new_scene()
    add(sc, P["body"], COL["body"])
    add(sc, P["button_cap"], COL["button_cap"])
    add(sc, cell, COL["cell"])
    for w in leads:
        add(sc, w, COL["lead"])
    shoot(sc, (40, -62, 70), ctr, os.path.join(OUT, "render_cell_in.png"))

    # board in the open body (battery side down), plate removed
    sc = new_scene()
    add(sc, P["body"], COL["body"])
    add(sc, cell, COL["cell"])
    for s, c in bparts:
        add(sc, s, c)
    shoot(sc, (45, -70, 75), ctr, os.path.join(OUT, "render_open.png"))

    # dock (empty), from above the cradle: recess, strap channel, pogo tongue, LED grooves, magnet pockets
    sc = new_scene()
    add(sc, P["dock_cradle"], COL["dock_cradle"])
    shoot(sc, (55, -85, -60), (16, -15, 9), os.path.join(OUT, "render_dock.png"))

    # device seated on the dock, cut through the pogo pins (board y 27.9): pins compressed onto the J1 pads
    cut = cq.Workplane("XY").box(100, 100, 100).translate((0, -27.9 + 50, 0)).val()
    sc = new_scene()
    for n in ("dock_cradle", "body", "back_plate", "button_cap"):
        add(sc, P[n].cut(cut), COL[n])
    for s, c in bparts:
        try:
            add(sc, s.cut(cut), c)
        except Exception:
            pass
    add(sc, cell.cut(cut), COL["cell"])
    for (x, y) in J1_PADS:
        add(sc, cq.Workplane("XY").workplane(offset=PAD_Z).center(x, -y).circle(0.37).extrude(TONGUE_TOP + 0.4 - PAD_Z)
            .union(cq.Workplane("XY").workplane(offset=TONGUE_TOP + 0.4).center(x, -y).circle(0.51).extrude(6)).val(),
            (0.85, 0.70, 0.25))
    shoot(sc, (9.6, 25, 3.0), (9.6, -27.9, 3.0), os.path.join(OUT, "render_docked_section.png"), fov=0.30)
    print("renders written")


if __name__ == "__main__":
    main()