Image-to-CadQuery
Custom Vision-to-Code Model
A neural network that turns a picture of a 3D part into code
Built with
PyTorchResNet18Transformer
Hugging Face Datasets
pytesttorchvisionHugging Face Transformerstorchrun
Ruff
Code is private (company / client work)
Overview
CadQuery is a Python library for describing 3D parts as code. I wanted to know whether a model trained purely from data could look at a rendered CAD part and write the CadQuery script that rebuilds it, with no LLM API and no prompt engineering involved.
I built an encoder-decoder model in plain PyTorch: a ResNet18 turns the image into a 7x7 grid of visual tokens, and a 6-layer Transformer decoder writes the script one token at a time while attending to those image tokens. Around it I wrote a typed configuration, a data pipeline for the GenCAD-Code dataset (147k image/script training pairs), a training loop that scales from a laptop CPU to multiple GPUs and machines, an inference CLI, and a CPU-only pytest suite.
The pipeline is verified end to end: in an overfit sanity check on real data, the model memorized 8 samples and regenerated their scripts almost exactly. A full GPU training run and evaluation of the generated code (valid-syntax rate and IoU against the true solid) are the next milestone.
The question
Can a purpose-built ResNet encoder and autoregressive Transformer decoder learn to translate CAD images directly into editable CadQuery programs?
Key features
- Translates a 224x224 render of a CAD part into an executable CadQuery Python script.
- Keeps 49 spatial image tokens so the decoder can look at different parts of the image for different lines of code.
- Trains on CPU, one GPU, several GPUs, or several machines with DDP, with mixed precision, gradient accumulation, and resumable checkpoints.
- Generates code from the command line or Python with greedy decoding or temperature and top-k sampling.
- Ships a CPU-only, offline test suite that checks causality, padding, loss masking, checkpoint reloads, and that two DDP processes end with identical weights.
How it works
Prepare data
Each GenCAD-Code render is resized to 224x224 and normalized; each script is tokenized with GPT-2 BPE and wrapped as BOS, code, EOS, with a dedicated pad token so the model still learns when to stop.
Encode the image
A ResNet18 without pooling produces a 7x7 feature grid, flattened into 49 image tokens with learned 2D positions.
Decode the code
A 6-layer pre-norm Transformer decoder uses masked self-attention over the code written so far and cross-attention to the image tokens, with the output layer tied to the token embedding.
Train
Teacher-forced next-token cross-entropy, AdamW with separate learning rates for encoder and decoder, linear warmup then cosine decay, and gradient clipping.
Generate
The image is encoded once, then the decoder appends one token at a time until every sequence emits EOS or reaches the token limit.
Examples
Input
Overfit sanity check: 8 real GenCAD-Code samples, 26.4M-parameter model, 150 steps on a laptop CPU (about 15 minutes)
Output
Validation loss fell from 10.64 to 0.06; reloading the checkpoint and generating from the 8 training images reproduced the scripts at 98-100% token accuracy (4 of 8 exactly).
Tech stack
Model, training loop, DDP, and mixed precision
ResNet18
Image encoder producing a 7x7 feature grid
Transformer
Autoregressive decoder with cross-attention to image
Loading the GenCAD-Code image/script pairs
Offline CPU tests for data, model, training
ResNet18 backbone and image transforms
GPT-2 BPE tokenizer for CadQuery code
Launching multi-GPU and multi-node training
Linting and formatting checks
What I learned
- Vision-to-code needs spatial features to remain available to every decoding step
- A tiny-overfit baseline is the right gate before expensive full-dataset training
Limitations
- Training, checkpoint generation, and VSR/IoU evaluation are the next milestone
- The repository is currently private