File size: 2,897 Bytes
82877b1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
---
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}
}
```
|