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"))

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

DEFAULT_BATTERY = "ldtek"


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.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(as_assy: bool=False):
    logger.info("computing model ...")
    battery_name = os.environ.get("CASE_BATTERY", DEFAULT_BATTERY)
    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 = None
    if battery_name == "ldtek":
        battery = ldtek_battery.LdtekBattery()
    return case.case(phone, battery=battery, render_phone=render_phone, as_assy=as_assy)

_computedModels = {}
def model(as_assy: bool=False):
    """ memoized wrapper around `_model` """
    global _computedModels
    key = (as_assy, )
    if key not in _computedModels:
        _computedModels[key] = _model(as_assy=as_assy)
    return _computedModels[key]

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("--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.render_phone:
        os.environ["CASE_RENDER_PHONE"] = "1"

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

    if args.battery:
        os.environ["CASE_BATTERY"] = args.battery

    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(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()
        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()