---
language: en
license: apache-2.0
tags:
- vision
- image-to-code
- cad
- cadquery
- vision-encoder-decoder
- vit
- gpt2
datasets:
- CADCODER/GenCAD-Code
metrics:
- rouge
widget:
- src: https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/tasks/car.jpg
  example_title: Example CAD Image
---

# VIT-CodeGPT CAD Code Generator

This model generates CADQuery Python code from images of 3D CAD objects. It uses a Vision Transformer (ViT) encoder and CodeGPT decoder in a vision-encoder-decoder architecture.

## Model Details

- **Architecture**: Vision Encoder-Decoder (ViT + CodeGPT)
- **Encoder**: google/vit-base-patch16-224
- **Decoder**: microsoft/CodeGPT-small-py
- **Task**: Image-to-Code Generation (CAD)
- **Dataset**: CADCODER/GenCAD-Code
- **Training Samples**: 10,000 (8,500 train / 1,500 val)
- **Training Time**: ~4 hours 12 minutes

## Training Configuration

- **Batch Size**: 4 (effective: 16 with gradient accumulation)
- **Learning Rate**: 3e-5
- **Epochs**: 3
- **Max Length**: 256 tokens
- **Optimizer**: AdamW with warmup
- **Mixed Precision**: FP16

## Performance

Final training metrics:
- **ROUGE-1**: 0.0944
- **ROUGE-2**: 0.0040
- **ROUGE-L**: 0.0863

## Usage

```python
from transformers import VisionEncoderDecoderModel, ViTFeatureExtractor, AutoTokenizer
from PIL import Image
import torch

# Load the model
model = VisionEncoderDecoderModel.from_pretrained("Thehunter99/vit-codegpt-cadcoder")
feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
tokenizer = AutoTokenizer.from_pretrained("microsoft/CodeGPT-small-py")

# Load and process image
image = Image.open("path/to/your/cad_image.png")
pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values

# Generate CAD code
with torch.no_grad():
    generated_ids = model.generate(
        pixel_values,
        max_length=256,
        num_beams=4,
        early_stopping=True,
        pad_token_id=tokenizer.eos_token_id
    )

generated_code = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
print(generated_code)
```

## Example Output

Input: Image of a 3D cube
Output:
```python
import cadquery as cq

# Create a simple cube
result = cq.Workplane("XY").box(10, 10, 10)
```

## Training Data

The model was trained on the CADCODER/GenCAD-Code dataset, which contains pairs of 3D CAD images and their corresponding CADQuery Python code.

## Limitations

- Limited to CADQuery syntax
- Best performance on geometric shapes similar to training data
- May struggle with very complex or unusual CAD designs
- Maximum output length: 256 tokens

## Citation

If you use this model, please cite:

```bibtex
@misc{vit-codegpt-cadcoder,
  title={VIT-CodeGPT CAD Code Generator},
  author={Your Name},
  year={2024},
  publisher={Hugging Face},
  url={https://huggingface.co/Thehunter99/vit-codegpt-cadcoder}
}
```