phone-case-cq/main.py

#!/usr/bin/env python3
"""
toplevel file used for interactive modeling and also data exports.

use like:
- `main.py --editor` => to edit the case interactively
  - then press green play button to render
  - edit files externally, and press render again to refresh the view
  - this is shorthand for `cq-editor ./main.py`
- `main.py --export-stl build/case.stl` => to export the case as stl

every operation works over a parameterized model. you can control this parameterization with environment variables (or flags -- see extended --help):
- CASE_RENDER_PHONE=1  => render not just the case, but also what it would look like with the phone inside
- CASE_RENDER_PHONE_ONLY=1  => render ONLY the phone, no case
- CASE_BATTERY=<name>  => design the case to fit a specific battery model, or "none".
"""

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

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

from config import Config
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: cq.Workplane, file_: str, orientation: str) -> None:
    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.Yaw(180)
        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(config: Config, as_assy: bool=False) -> cq.Workplane:

    phone = pinephone.PinePhone()

    if config.render_phone_only:
        return case.orient_for_printing(phone)

    battery = None
    if config.battery == "ldtek":
        battery = ldtek_battery.LdtekBattery()
    elif config.battery:
        assert False, f"unknown battery: {battery!r}"
    return case.case(phone, battery=battery, config=config, render_phone=config.render_phone, as_assy=as_assy)

_computedModels: dict[(bool,), cq.Workplane] = {}
def model(config: Config, as_assy: bool=False) -> cq.Workplane:
    """ memoized wrapper around `_model` """
    global _computedModels
    key = (as_assy, )
    if key not in _computedModels:
        _computedModels[key] = _model(config, as_assy=as_assy)
    return _computedModels[key]

def main() -> None:
    logging.basicConfig()
    logging.getLogger().setLevel(logging.INFO)

    parser = argparse.ArgumentParser(description="toplevel cadquery interface")
    parser.add_argument("--verbose", action="store_true", help="debug-level logging")
    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("--battery", choices=["none", "ldtek"], help="name of the battery for which to create a pocket cutout, or 'none'")
    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.verbose or os.environ.get("CASE_DEBUG"):
        logging.getLogger().setLevel(logging.DEBUG)
        logger.debug("enabled debug-level logging")

    config = Config.from_env()

    if args.render_phone:
        config.render_phone = args.render_phone

    if args.render_phone_only:
        config.render_phone_only = args.render_phone_only

    if args.battery:
        config.battery = args.battery

    if args.export_stl:
        model_ = model(config)
        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(config, as_assy=True)
        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(config)
        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(Config.from_env())