phone-case-cq/cq_toplevel.py

#!/usr/bin/env python3
"""
toplevel file used for interactive modeling.

- `cq-editor ./cq_toplevel.py`
- then press green play button to render
- edit files externally, and press render again to refresh the view
"""

import cadquery as cq
import argparse
import logging
import os
import subprocess
import sys

sys.path.append(os.path.join(os.getcwd(), "src"))

import case
import pinephone
import ldtek_battery

from cadquery.occ_impl.assembly import toVTK
from cadquery.vis import _to_assy
from vtkmodules.vtkRenderingCore import vtkRenderWindow, vtkWindowToImageFilter
from vtkmodules.vtkIOImage import vtkPNGWriter


logger = logging.getLogger(__name__)

def export_png_image(obj, file_: str, orientation: str):
    assy = _to_assy(obj)
    renderer = toVTK(assy)
    win = vtkRenderWindow()
    win.AddRenderer(renderer)
    win.Render()
    camera = renderer.GetActiveCamera()
    if orientation == "front":
        camera.Roll(-28)
        camera.Elevation(-50)
    elif orientation == "back":
        camera.Roll(0)
        camera.Elevation(-35)
    elif orientation == "side":
        camera.Yaw(75)
        camera.Elevation(30)
    # adjust camera so full object is visible.
    # this also resizes the window, potentially changing its aspect ratio.
    # it's important that the window has been `Render()`'d at least once by now,
    # else it'll adjust the camera based on the wrong aspect ratio.
    renderer.ResetCamera()
    renderer.SetBackground(0.8, 0.8, 0.8)

    win.Render()

    # documented here: <https://examples.vtk.org/site/Python/IO/ImageWriter/>
    win_to_input = vtkWindowToImageFilter()
    win_to_input.SetInput(win)
    win_to_input.SetInputBufferTypeToRGB()
    win_to_input.ReadFrontBufferOff()
    win_to_input.Update()

    exporter = vtkPNGWriter()
    exporter.SetFileName(file_)
    exporter.SetInputConnection(win_to_input.GetOutputPort())
    exporter.Write()


def _model():
    logger.info("computing model ...")
    render_phone = os.environ.get("CASE_RENDER_PHONE", "") not in ("", "0")
    render_phone_only = os.environ.get("CASE_RENDER_PHONE_ONLY", "") not in ("", "0")

    phone = pinephone.PinePhone()

    if render_phone_only:
        return case.orient_for_printing(phone)

    battery = ldtek_battery.LdtekBattery()
    return case.case(phone, battery=battery, render_phone=render_phone)

_computedModel = None
def model():
    """ memoized wrapper around `_model` """
    global _computedModel
    if _computedModel is None:
        _computedModel = _model()
    return _computedModel

def main():
    logging.basicConfig()
    logging.getLogger().setLevel(logging.INFO)

    parser = argparse.ArgumentParser(description="toplevel cadquery interface")
    parser.add_argument("--render-phone", action="store_true", help="render the case and also the phone within it; useful to confirm fit visually before printing")
    parser.add_argument("--render-phone-only", action="store_true", help="render *only* the phone, not even the case")
    parser.add_argument("--export-stl")
    parser.add_argument("--export-png")
    parser.add_argument("--export-vtk")
    parser.add_argument("--editor", action="store_true", help="view in cq-editor")

    args = parser.parse_args()

    if args.render_phone:
        os.environ["CASE_RENDER_PHONE"] = "1"

    if args.render_phone_only:
        os.environ["CASE_RENDER_PHONE_ONLY"] = "1"

    if args.export_stl:
        model_ = model()
        logger.info("exporting stl to %s", args.export_stl)
        cq.exporters.export(model_, args.export_stl)

    if args.export_png:
        orientation = None
        if "side" in args.export_png:
            orientation = "side"
        if "back" in args.export_png:
            orientation = "back"
        if "front" in args.export_png:
            orientation = "front"
        model_ = model()
        logger.info("exporting png to %s", args.export_png)
        export_png_image(model_, args.export_png, orientation)

    if args.export_vtk:
        vtk_file = args.export_vtk
        js_var, _ext = os.path.splitext(os.path.basename(vtk_file))
        js_file = f'{vtk_file}.js'
        model_ = model()
        logger.info("exporting VTK (for web rendering) to %s", vtk_file)
        cq.exporters.export(model_, vtk_file, cq.exporters.ExportTypes.VTP)

        logger.info("wrapping VTK data in a javascript variable (var %s) in %s", js_var, js_file)
        vtk_data = open(vtk_file).read()
        with open(js_file, 'w') as js:
            js.write(f"var {js_var} = `\n")
            js.write(vtk_data)
            js.write("`;\n")

    if args.editor:
        logger.info("launching cq-editor")
        subprocess.check_call(["cq-editor", __file__])


if __name__ == "__main__":
    main()
else:
    # this `result` var should be picked up by cadquery, in case we were imported by it.
    # note that we don't actually get here until the user presses the `>` render button.
    result = model()