Upload 50 files

CADFusion/.gitignore

@@ -0,0 +1,171 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# PyPI configuration file
+.pypirc

CADFusion/CODE_OF_CONDUCT.md

@@ -0,0 +1,9 @@
+# Microsoft Open Source Code of Conduct
+
+This project has adopted the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/).
+
+Resources:
+
+- [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/)
+- [Microsoft Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/)
+- Contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with questions or concerns

CADFusion/LICENSE

@@ -0,0 +1,21 @@
+    MIT License
+
+    Copyright (c) Microsoft Corporation.
+
+    Permission is hereby granted, free of charge, to any person obtaining a copy
+    of this software and associated documentation files (the "Software"), to deal
+    in the Software without restriction, including without limitation the rights
+    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+    copies of the Software, and to permit persons to whom the Software is
+    furnished to do so, subject to the following conditions:
+
+    The above copyright notice and this permission notice shall be included in all
+    copies or substantial portions of the Software.
+
+    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+    SOFTWARE

CADFusion/README.md

@@ -0,0 +1,194 @@
+# CADFusion
+
+
+This repo is the official implementation of paper **[ICML 2025] Text-to-CAD Generation Through Infusing Visual Feedback in Large Language Models** by *Ruiyu Wang, Yu Yuan, Shizhao Sun, Jiang Bian*.
+
+[Paper](https://arxiv.org/abs/2501.19054) | [Video](https://www.youtube-nocookie.com/embed/LK8LAzR0v5M?si=FD1Vg9wjkROTKjDV) | [Huggingface](https://huggingface.co/microsoft/CADFusion)
+
+CADFusion is a text-to-CAD generation framework that leverages visual feedback to enhance the performance of large language models (LLMs) in generating CAD models from textual descriptions. It consists of two main components: sequential learning and visual learning. The sequential learning component fine-tunes LLMs on a text-to-CAD dataset, while the visual learning component alternates between training a visual feedback model and fine-tuning the LLM with the generated visual feedback.
+
+## Installation
+
+- Create a conda environment and install the generic dependencies.
+
+```
+name=<your-env-name>
+conda create -n $name python=3.9
+conda activate $name
+python -m pip install -e .
+```
+
+- Install the additional dependencies for training.
+
+```
+python -m pip install -e .["train"]
+```
+
+- Install the additional dependencies for evaluation and rendering.
+
+```
+python -m pip install -e .["render"]
+conda install -c conda-forge pythonocc-core=7.7.0
+python -m pip install git+https://github.com/otaheri/chamfer_distance@dc9987dcf70888d387d96893ba1fb9ba9a333992
+python -m pip install -e .["eval"]
+```
+
+## Data Preparation
+CADFusion is trained by alternating the **Sequential Learning (SL)** stage and the **Visual Feedback (VF)** stage.
+We introduce how to prepare the training data for these two stages in the below.
+
+### Data for Sequential Learning
+
+#### Approach 1: use human-annotated textual descriptions provided by us
+We provide human-annoated textual descriptions and their correspoding CAD model IDs in [Skexgen](https://github.com/samxuxiang/SkexGen) under `data/sl_data/sl_data.zip`. It should contain the following files after unzipping:
+```
+data/sl_data
+├── train.json
+├── val.json
+├── test.json
+```
+To use our annotated data, download the SkexGen data, unzip it as the reference dataset and run the convertion script to get the dataset. In detail, run the following command:
+```
+# make sure you are in the root directory of this repo and have the 'data/sl_data/sl_data.zip' unzipped
+gdown --id 1so_CCGLIhqGEDQxMoiR--A4CQk4MjuOp 
+unzip cad_data.zip
+python3 data/sl_data/convert.py
+```
+The `train.json`, `val.json` and `test.json` under `data/sl_data` are the datasets.
+
+#### Approach 2: create human-annotated textual descriptions by yourself
+We provide a script to execute all the preprocessing steps until human annotation. 
+```
+./scripts/preprocess_skexgen.sh
+```
+If you want to customize the internal steps, expand the following section for more details.
+<details>
+<summary>Start from scratch (click to expand).</summary>
+
+1. Download the [SkexGen](https://github.com/samxuxiang/SkexGen) data by: [Google Drive link](https://drive.google.com/file/d/1so_CCGLIhqGEDQxMoiR--A4CQk4MjuOp/view).
+
+```
+gdown --id 1so_CCGLIhqGEDQxMoiR--A4CQk4MjuOp
+unzip cad_data.zip
+```
+
+2. Convert the SkexGen data into sequences. Note that `train_deduplicate_s.pkl`, `val.pkl` and `test.pkl` should be converted separately.
+```
+python3 src/data_preprocessing/convert.py --in_path <skexgen_path> --out_path <sequence_path>
+```
+
+3. Render the sequences into images. *Note that running the last step on linux requires the installation of an x server (e.g. `xvfb`). See [this discussion.](https://github.com/tpaviot/pythonocc-core/issues/1302#issuecomment-2053526444)*
+```
+python3 src/rendering_utils/parser.py --in-path <sequence_path> --out-path <visual_object_folder>
+timeout 180 python3 src/rendering_utils/parser_visual.py --data_folder <visual_object_folder>
+python3 src/rendering_utils/img_renderer.py --input_dir <visual_object_folder> --output_dir <image_folder>
+```
+
+4. Annotate these data with LLM captioning.
+```
+# Generic:
+python3 src/data_preprocessing/captioning.py --image-folder-path <image_folder> --out-path <sl_data_path>
+
+```
+* We use openai and azure system for LLM calling. You are welcome to use your own LLMs and prompts by changing `line 21, 22` of `src/data_preprocessing/captioning.py` with your own client definition and function calls.
+</details>
+
+
+### Data for Visual Feedback
+
+The Visual Feedback dataset should be automatically generated from the Visual Feedback pipeline described in the Training section. 
+We provide an example under `data/vf_data/example_vf_data.json` to help people understand how it should look like.
+You can retrieve this file by unzipping `data/vf_data/example_vf_data.zip`.
+We do not recommend using this example data as the training data, as the policy update should depend on its own generations.
+
+
+## Training 
+Our training receipe contains two parts. In the first part, we conduct initial sequential learning. In the second part, we conduct alternate training between sequential learning and visual feedback.
+### Initial Sequential Learning
+We use the following script to train the model in the sequential learning stage.
+```
+./scripts/train_with_shuffling.sh <run_name>
+```
+
+You are also welcome to customize the training procedure. A normal training script on multiple GPUs is provided. Change `num_processes` in `ds_config.yaml` to specify how many GPUs will be used.
+```
+CUDA_VISIBLE_DEVICES=<gpu_ids> accelerate launch --config_file ds_config.yaml src/train/llama_finetune.py \
+    --num-epochs <num_epochs> --run-name <run_name> --data-path <train_data> --eval-data-path <eval_data> \
+    --device-map accelerate --model-name llama3 --expdir <model_saving_path>
+```
+
+In our work we shuffle the dataset per x epochs. To train model with this implementation, inspect and modify `scripts/train_with_shuffling.sh`.
+
+### Alternate Training between Sequential Learning and Visual Feedback
+We provide a script for executing our alternate training round. See `scripts/alternate_VF.sh`.
+```
+./scripts/alternate_VF.sh  # change the value of base_name in the script as instructed
+```
+We also provide a script for training on multiple gpus for saving time: `scripts/alternate_VF_quadra_gpu.sh`. In our setting, we use 4 GPUs for training. You can change the script to use more GPUs if you have them available.
+
+If you only want to conduct a single round of visual learning, run
+```
+python src/train/dpo.py --run-name <dpo_run_name> --pretrained-path <pretrained_model_path> --data-path <dpo_data_Path> --output-path <model_saving_path>
+```
+By default it runs dpo for 3 epochs, but you can change by adding flag `--num-epochs x`.
+
+
+## Model Checkpoints
+We provide two versions. 
+v1.0 has 5 rounds of alternate training and is used for evaluation in our paper.
+v1.1 has 9 rounds of alternate training and is considered to have better performance than v1.0.
+- [CADFusion v1.0](https://huggingface.co/microsoft/CADFusion/tree/main/v1_0)
+- [CADFusion v1.1](https://huggingface.co/microsoft/CADFusion/tree/main/v1_1)
+
+You should download, unzip and place them under the `exp/model_ckpt` folder for using.
+
+## Inference & Visualization
+Use `scripts/generate_samples.sh`.
+```
+./scripts/generate_samples.sh <run_name> test --full
+```
+You can find samples generated in `exp/model_generation/<run_name>.jsonl` and rendered figures under the `exp/figures/<run_name>` folder. The point clouds, .obj files, .step and .stl files are saved under `exp/visual_objects/<run_name>` directory for your own usage and evaluation.
+
+## Evaluation
+Use the functions in `src/test`. This includes the Chamfer Distance (`chamfer_dist.py`), Minimum Matching Distance, Coverage, Jensen-Shannon Divergence (`dist_eval.py`), and the VLM score (`VLM_score.py`).
+
+For VLM Score, we use Azure OpenAI API to access the GPT-4o model for scoring the CAD objects. 
+In this way, you should log in your own azure account before using this module.
+If your are using other LLM/VLM service and feel difficult to adapt to our setup, we provide the prompt in the python module that is available for you to integrate into your own testing pipeline.
+
+### 
+
+## Acknowledgements
+We would like to acknowledge that the CAD rendering and distributional metrics in this repository is partially based on and adapted from the [SkexGen](https://github.com/samxuxiang/SkexGen) project.
+
+## Citation
+If you find our work useful, please cite the following paper
+```
+@inproceedings{wang2025texttocad, 
+  title = {Text-to-CAD Generation Through Infusing Visual Feedback in Large Language Models},
+  author = {Wang, Ruiyu and Yuan, Yu and Sun, Shizhao and Bian, Jiang},
+  booktitle = {International Conference on Machine Learning},
+  year={2025}
+}
+```
+## Contributing
+
+This project welcomes contributions and suggestions.  Most contributions require you to agree to a
+Contributor License Agreement (CLA) declaring that you have the right to, and actually do, grant us
+the rights to use your contribution. For details, visit https://cla.opensource.microsoft.com.
+
+When you submit a pull request, a CLA bot will automatically determine whether you need to provide
+a CLA and decorate the PR appropriately (e.g., status check, comment). Simply follow the instructions
+provided by the bot. You will only need to do this once across all repos using our CLA.
+
+This project has adopted the [Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct/).
+For more information see the [Code of Conduct FAQ](https://opensource.microsoft.com/codeofconduct/faq/) or
+contact [opencode@microsoft.com](mailto:opencode@microsoft.com) with any additional questions or comments.
+
+## Trademarks
+
+This project may contain trademarks or logos for projects, products, or services. Authorized use of Microsoft 
+trademarks or logos is subject to and must follow 
+[Microsoft's Trademark & Brand Guidelines](https://www.microsoft.com/en-us/legal/intellectualproperty/trademarks/usage/general).
+Use of Microsoft trademarks or logos in modified versions of this project must not cause confusion or imply Microsoft sponsorship.
+Any use of third-party trademarks or logos are subject to those third-party's policies.

CADFusion/SECURITY.md

@@ -0,0 +1,41 @@
+<!-- BEGIN MICROSOFT SECURITY.MD V0.0.9 BLOCK -->
+
+## Security
+
+Microsoft takes the security of our software products and services seriously, which includes all source code repositories managed through our GitHub organizations, which include [Microsoft](https://github.com/Microsoft), [Azure](https://github.com/Azure), [DotNet](https://github.com/dotnet), [AspNet](https://github.com/aspnet) and [Xamarin](https://github.com/xamarin).
+
+If you believe you have found a security vulnerability in any Microsoft-owned repository that meets [Microsoft's definition of a security vulnerability](https://aka.ms/security.md/definition), please report it to us as described below.
+
+## Reporting Security Issues
+
+**Please do not report security vulnerabilities through public GitHub issues.**
+
+Instead, please report them to the Microsoft Security Response Center (MSRC) at [https://msrc.microsoft.com/create-report](https://aka.ms/security.md/msrc/create-report).
+
+If you prefer to submit without logging in, send email to [secure@microsoft.com](mailto:secure@microsoft.com).  If possible, encrypt your message with our PGP key; please download it from the [Microsoft Security Response Center PGP Key page](https://aka.ms/security.md/msrc/pgp).
+
+You should receive a response within 24 hours. If for some reason you do not, please follow up via email to ensure we received your original message. Additional information can be found at [microsoft.com/msrc](https://www.microsoft.com/msrc). 
+
+Please include the requested information listed below (as much as you can provide) to help us better understand the nature and scope of the possible issue:
+
+  * Type of issue (e.g. buffer overflow, SQL injection, cross-site scripting, etc.)
+  * Full paths of source file(s) related to the manifestation of the issue
+  * The location of the affected source code (tag/branch/commit or direct URL)
+  * Any special configuration required to reproduce the issue
+  * Step-by-step instructions to reproduce the issue
+  * Proof-of-concept or exploit code (if possible)
+  * Impact of the issue, including how an attacker might exploit the issue
+
+This information will help us triage your report more quickly.
+
+If you are reporting for a bug bounty, more complete reports can contribute to a higher bounty award. Please visit our [Microsoft Bug Bounty Program](https://aka.ms/security.md/msrc/bounty) page for more details about our active programs.
+
+## Preferred Languages
+
+We prefer all communications to be in English.
+
+## Policy
+
+Microsoft follows the principle of [Coordinated Vulnerability Disclosure](https://aka.ms/security.md/cvd).
+
+<!-- END MICROSOFT SECURITY.MD BLOCK -->

CADFusion/SUPPORT.md

@@ -0,0 +1,25 @@
+# TODO: The maintainer of this repo has not yet edited this file
+
+**REPO OWNER**: Do you want Customer Service & Support (CSS) support for this product/project?
+
+- **No CSS support:** Fill out this template with information about how to file issues and get help.
+- **Yes CSS support:** Fill out an intake form at [aka.ms/onboardsupport](https://aka.ms/onboardsupport). CSS will work with/help you to determine next steps.
+- **Not sure?** Fill out an intake as though the answer were "Yes". CSS will help you decide.
+
+*Then remove this first heading from this SUPPORT.MD file before publishing your repo.*
+
+# Support
+
+## How to file issues and get help  
+
+This project uses GitHub Issues to track bugs and feature requests. Please search the existing 
+issues before filing new issues to avoid duplicates.  For new issues, file your bug or 
+feature request as a new Issue.
+
+For help and questions about using this project, please **REPO MAINTAINER: INSERT INSTRUCTIONS HERE 
+FOR HOW TO ENGAGE REPO OWNERS OR COMMUNITY FOR HELP. COULD BE A STACK OVERFLOW TAG OR OTHER
+CHANNEL. WHERE WILL YOU HELP PEOPLE?**.
+
+## Microsoft Support Policy  
+
+Support for this **PROJECT or PRODUCT** is limited to the resources listed above.

CADFusion/data/sl_data/convert.py

@@ -0,0 +1,125 @@
+import json
+import pickle
+
+SKETCH_R = 1
+RADIUS_R = 1
+EXTRUDE_R = 1.0
+SCALE_R = 1.4
+OFFSET_R = 0.9
+PIX_PAD = 4
+CMD_PAD = 3
+COORD_PAD = 4
+EXT_PAD = 1
+EXTRA_PAD = 1
+R_PAD = 2
+
+def create_curve_str(se_xy, se_cmd):
+    curve_str = ""
+    xy_offset = 0
+    if se_cmd == 0:  # line
+        curve_str = " line," + ",".join(str(x) for x in se_xy[0])
+        xy_offset = 2
+    elif se_cmd == 1:  # arc
+        curve_str = " arc," + ",".join(str(x) for x in se_xy[0:2].flatten())
+        xy_offset = 3
+    elif se_cmd == 2:  # circle
+        curve_str = " circle," + ",".join(str(x) for x in se_xy[0:4].flatten())
+        xy_offset = 5
+    curve_str += " <curve_end>"
+    return curve_str, xy_offset
+
+
+def create_sketch_str(se_xy, se_cmd):
+    sketch_str = ""
+    len_xy, len_cmd = len(se_xy), len(se_cmd)
+    xy_idx = 0
+    for cmd_item in se_cmd:  # for each command
+        if 0 <= cmd_item <= 2:  # curve
+            curve_str, xy_offset = create_curve_str(se_xy[xy_idx:], cmd_item)
+            sketch_str += curve_str
+            xy_idx += xy_offset
+        elif cmd_item == -1:  # loop
+            sketch_str += " <loop_end>"
+            xy_idx += 1
+        elif cmd_item == -2:  # face
+            sketch_str += " <face_end>"
+            xy_idx += 1
+        elif cmd_item == -3:  # sketch
+            sketch_str += " <sketch_end>"
+            xy_idx += 1
+        else:
+            raise ValueError("Invalid command: " + str(cmd_item))
+    if xy_idx != len_xy:
+        raise ValueError("xy_idx != len_xy")
+    return sketch_str
+
+
+def create_extrude_str(se_ext):
+    extrude_str = ""
+    # extrude operation
+    if se_ext[14] == 1:
+        extrude_str += "add"
+    elif se_ext[14] == 2:
+        extrude_str += "cut"
+    elif se_ext[14] == 3:
+        extrude_str += "intersect"
+    else:
+        raise ValueError("Invalid extrude operation: " + str(se_ext[14]))
+    # other extrude parameters
+    extrude_str = (
+        extrude_str + "," + ",".join(str(x - EXT_PAD) for x in se_ext[0:5])
+    )  # ext_v, ext_T
+    extrude_str = (
+        extrude_str + "," + ",".join(str(x - R_PAD) for x in se_ext[5:14])
+    )  # ext_R
+    extrude_str = (
+        extrude_str + "," + ",".join(str(x - EXT_PAD) for x in se_ext[15:18])
+    )  # scale, offset
+    # extrude end
+    extrude_str += " <extrude_end>"
+    return extrude_str
+
+def create_command_sequence(item):
+    se_str = ""
+    num_se = item["num_se"]
+    for se_idx in range(num_se):  # for each sketch-extrude
+        xy, cmd, ext = (
+            item["se_xy"][se_idx] - COORD_PAD,
+            item["se_cmd"][se_idx] - CMD_PAD,
+            item["se_ext"][se_idx],
+        )
+        se_str = se_str + " " + create_sketch_str(xy, cmd).strip()
+        se_str = se_str + " " + create_extrude_str(ext).strip()
+    return se_str.strip()
+
+with open("data/sl_data/train.json", "r") as f:
+    train_data = json.load(f)
+with open("data/sl_data/test.json", "r") as f:
+    test_data = json.load(f)
+with open("data/sl_data/val.json", "r") as f:
+    val_data = json.load(f)
+
+with open("cad_data/train_deduplicate_s.pkl", "rb") as f:
+    sk_data = pickle.load(f)
+
+for item in train_data:
+    serial_num = item['serial_num']
+    description = item['description']
+    item["command_sequence"] = create_command_sequence(sk_data[serial_num])
+
+for item in test_data:
+    serial_num = item['serial_num']
+    description = item['description']
+    item["command_sequence"] = create_command_sequence(sk_data[serial_num])
+
+for item in val_data:
+    serial_num = item['serial_num']
+    description = item['description']
+    item["command_sequence"] = create_command_sequence(sk_data[serial_num])
+
+with open("data/sl_data/train.json", "w+") as f:
+    json.dump(train_data, f, indent=4)
+with open("data/sl_data/test.json", "w+") as f:
+    json.dump(test_data, f, indent=4)
+with open("data/sl_data/val.json", "w+") as f:
+    json.dump(val_data, f, indent=4)

CADFusion/data/sl_data/sl_data.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a080e00591a07420d916e82365d8602ebeab00ffd909f87bc9911b231f2f5ea0
+size 1084518

CADFusion/data/vf_data/example_vf_data.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:907df4efd2ceafd9d8c336dfbf62d1754f692c0aab72b1b212ea7b844125e702
+size 2142

CADFusion/ds_config.yaml

@@ -0,0 +1,22 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+deepspeed_config:
+  gradient_accumulation_steps: 1
+  gradient_clipping: 1.0
+  offload_optimizer_device: none
+  offload_param_device: none
+  zero3_init_flag: true
+  zero_stage: 2
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+machine_rank: 0
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 1
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

CADFusion/pyproject.toml

@@ -0,0 +1,38 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "CADFusion"
+version = "1.0.0"
+description = "Enhancing Text-to-CAD generation via sequential learning and visual feedback."
+readme = "README.md"
+requires-python = ">=3.8"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: Apache Software License",
+]
+dependencies = [
+    "torch==2.7.1",
+    "transformers==4.50.0",
+    "huggingface_hub==0.26.0",
+    "peft==0.9.0",
+    "accelerate==0.28.0",
+    "psutil==5.9.8",
+    "pillow==10.4.0",
+    "datasets==3.1.0",
+    "trl==0.11.4",
+    "gdown==5.2.0"
+]
+
+[project.optional-dependencies]
+train = ["wandb==0.16.4", "deepspeed==0.15.0"]
+render = ["trimesh==4.4.9", "plyfile==1.0.3"]
+eval = ["openai==1.75.0", "azure-identity==1.21.0", "scikit-learn==1.3.2"]
+build = ["build", "twine"]
+
+[tool.setuptools.packages.find]
+exclude = ["assets*", "benchmark*", "docs", "dist*", "playground*", "scripts*", "tests*"]
+
+[tool.wheel]
+exclude = ["assets*", "benchmark*", "docs", "dist*", "playground*", "scripts*", "tests*"]

CADFusion/scripts/alternate_VF.sh

@@ -0,0 +1,47 @@
+# set it to your data path
+data_path=data/sl_data
+# by default set it to CADFusion/exp
+exp_path=exp/model_ckpt
+# by default set it to CADFusion/data
+vf_path=data/vf_data
+train_data=$data_path/train.json
+eval_data=$data_path/val.json
+
+# This script requires your SL run named as xxxx0, because for each VF stage, the final digit increments 
+# to show the number of VF rounds finished.
+# e.g. SL name: CAD-0
+#         base_name: CAD- (remove the last digit, the script autofills it)
+#         VF run 1: CAD-1 (automatically)
+#         VF run 2: CAD-2 (automatically)
+#         ...
+base_name=model_name_you_trained_for_SL_with_last_digit_removed
+
+run_name=${base_name}0
+./scripts/generate_samples.sh $run_name test "--full --device-map auto"
+./scripts/generate_samples.sh $run_name train "--sample-len 1000 --device-map auto"
+
+./scripts/make_dpo_data.sh $run_name --score-only 
+./scripts/make_dpo_data.sh $run_name-train "--gpu 0"
+
+
+for LOOP in 1 2 3 4 5
+do
+    echo "Starting VF round $LOOP"
+    run_name=$base_name$LOOP    
+    dpo_training_path=$vf_path/$base_name$((LOOP-1))-train.json
+    dpo_run_name=$base_name$LOOP-dpo
+    dpo_save_path=$exp_path/$dpo_run_name
+    sft_run_name=$base_name$LOOP
+
+    python src/train/dpo.py --run-name $dpo_run_name --pretrained-path $exp_path/$base_name$((LOOP-1)) --data-path $dpo_training_path --output-path $dpo_save_path
+    python src/train/llama_finetune.py --num-epochs 1 --run-name $sft_run_name --data-path $train_data --eval-data-path $eval_data --eval-freq 3000 --pretrained-path $dpo_save_path --expdir $exp_path
+    
+    ./scripts/generate_samples.sh $dpo_run_name test "--full --device-map auto"
+    ./scripts/generate_samples.sh $run_name test "--full --device-map auto"
+    ./scripts/generate_samples.sh $run_name train "--sample-len 1000 --device-map auto"
+
+    ./scripts/make_dpo_data.sh $dpo_run_name --score-only
+    ./scripts/make_dpo_data.sh $run_name "--score-only --gpu 0"
+    ./scripts/make_dpo_data.sh $run_name-train "--gpu 0"
+
+done

CADFusion/scripts/alternate_VF_quadra_gpu.sh

@@ -0,0 +1,50 @@
+# set it to your data path
+data_path=data/sl_data
+# by default set it to CADFusion/exp
+exp_path=exp/model_ckpt
+# by default set it to CADFusion/data
+vf_path=data/vf_data
+train_data=$data_path/train.json
+eval_data=$data_path/val.json
+
+# This script requires your SL run named as xxxx0, because for each VF stage, the final digit increments 
+# to show the number of VF rounds finished.
+# e.g. SL name: CAD-0
+#         base_name: CAD- (remove the last digit, the script autofills it)
+#         VF run 1: CAD-1 (automatically)
+#         VF run 2: CAD-2 (automatically)
+#         ...
+base_name=model_name_you_trained_for_SL_with_last_digit_removed
+
+run_name=${base_name}0
+CUDA_VISIBLE_DEVICES=0,1 ./scripts/generate_samples.sh $run_name test "--full --device-map auto" &
+CUDA_VISIBLE_DEVICES=2,3 ./scripts/generate_samples.sh $run_name train "--sample-len 10 --device-map auto"
+wait
+
+./scripts/make_dpo_data.sh $run_name --score-only &
+./scripts/make_dpo_data.sh $run_name-train "--gpu 1"
+wait
+
+
+for LOOP in 1 2 3 4 5
+do
+    echo "Starting VF round $LOOP"
+    run_name=$base_name$LOOP    
+    dpo_training_path=$vf_path/$base_name$((LOOP-1))-train.json
+    dpo_run_name=$base_name$LOOP-dpo
+    dpo_save_path=$exp_path/$dpo_run_name
+    sft_run_name=$base_name$LOOP
+
+    python src/train/dpo.py --run-name $dpo_run_name --pretrained-path $exp_path/$base_name$((LOOP-1)) --data-path $dpo_training_path --output-path $dpo_save_path
+    python src/train/llama_finetune.py --num-epochs 1 --run-name $sft_run_name --data-path $train_data --eval-data-path $eval_data --eval-freq 3000 --pretrained-path $dpo_save_path --expdir $exp_path
+    
+    CUDA_VISIBLE_DEVICES=0 ./scripts/generate_samples.sh $dpo_run_name test "--full --device-map auto" &
+    CUDA_VISIBLE_DEVICES=1 ./scripts/generate_samples.sh $run_name test "--full --device-map auto" &
+    CUDA_VISIBLE_DEVICES=2,3 ./scripts/generate_samples.sh $run_name train "--sample-len 1000 --device-map auto"
+    wait
+
+    ./scripts/make_dpo_data.sh $dpo_run_name --score-only &
+    ./scripts/make_dpo_data.sh $run_name "--score-only --gpu 1" &
+    ./scripts/make_dpo_data.sh $run_name-train "--gpu 2"
+    wait
+done

CADFusion/scripts/generate_samples.sh

@@ -0,0 +1,44 @@
+train_data_path=data/sl_data/train.json
+test_data_path=data/sl_data/test.json
+run_name=$1
+temperature=0.9
+
+if [ -z "$2" ]
+  then
+    data_path=$test_data_path
+else
+    if [ $2 = "train" ]; then
+        data_path=$train_data_path
+        run_name=$1-train
+    else
+        data_path=$test_data_path
+        temperature=0.3
+    fi
+fi
+
+model_path=exp/model_ckpt/$1
+inference_path=exp/model_generation/$run_name.jsonl
+visual_obj_path=exp/visual_objects/$run_name
+output_figure_path=exp/figures/$run_name
+log_path=exp/logs/$run_name
+
+mkdir -p $log_path
+mkdir -p exp/model_generation
+
+echo "--------------------Inferencing--------------------" > $log_path/inference.txt
+rm $inference_path
+python3 src/test/inference.py --pretrained-path $model_path --in-path $data_path --out-path $inference_path --num-samples 5 --temperature $temperature --model-name llama3 > $log_path/inference.txt $3
+
+echo "--------------------Parsing CAD objects--------------------" > $log_path/parsing_cad.txt
+rm -rf $visual_obj_path
+python3 src/rendering_utils/parser.py --in-path $inference_path --out-path $visual_obj_path > $log_path/parsing_cad.txt
+
+echo "--------------------Parsing visual objects--------------------" > $log_path/parsing_visual.txt
+python3 src/rendering_utils/parser_visual.py --data_folder $visual_obj_path > $log_path/parsing_visual.txt
+python3 src/rendering_utils/ptl_sampler.py --in_dir $visual_obj_path --out_dir ptl > $log_path/sampling_ptl.out
+
+echo "--------------------Rendering--------------------" > $log_path/rendering.txt
+rm -rf $output_figure_path
+export DISPLAY=:99
+Xvfb :99 -screen 0 640x480x24 &
+python3 src/rendering_utils/img_renderer.py --input_dir $visual_obj_path --output_dir $output_figure_path > $log_path/rendering.txt

CADFusion/scripts/make_dpo_data.sh

@@ -0,0 +1,5 @@
+source_path=exp/model_generation/$1.jsonl
+figure_path=exp/figures/$1/
+save_path=data/vf_data/$1.json
+
+python src/dpo/make_dpo_dataset.py --source-data-path $source_path --figure-path $figure_path --save-path $save_path --num-samples 5 $2

CADFusion/scripts/preprocess_skexgen.sh

@@ -0,0 +1,28 @@
+gdown --id 1so_CCGLIhqGEDQxMoiR--A4CQk4MjuOp
+unzip cad_data.zip
+
+# convert data into sequence and save in json
+mkdir data
+mkdir data/raw
+python3 src/data_preprocessing/convert.py --in-path cad_data/train_deduplicate_s.pkl --out-path data/raw/train.json
+python3 src/data_preprocessing/convert.py --in-path cad_data/val.pkl --out-path data/raw/val.json
+python3 src/data_preprocessing/convert.py --in-path cad_data/test.pkl --out-path data/raw/test.json
+
+# render the image for each entry in order to retrieve textual information by captioning:
+mkdir exp
+mkdir exp/visual_objects
+mkdir exp/figures
+for file in test val train; do
+    python3 src/rendering_utils/parser.py --in-path data/raw/$file.json --out-path exp/visual_objects/$file
+    timeout 180 python3 src/rendering_utils/parser_visual.py --data_folder exp/visual_objects/$file
+
+    export DISPLAY=:99
+    Xvfb :99 -screen 0 640x480x24 &
+    python3 src/rendering_utils/img_renderer.py --input_dir exp/visual_objects/$file --output_dir exp/figures/$file
+done
+
+# caption the images to generate descriptions
+mkdir data/sl_data
+python3 src/data_preprocessing/captioning.py --image-folder-path exp/figures/train --out-path data/sl_data/train.json
+python3 src/data_preprocessing/captioning.py --image-folder-path exp/figures/val --out-path data/sl_data/val.json
+python3 src/data_preprocessing/captioning.py --image-folder-path exp/figures/test --out-path data/sl_data/test.json

CADFusion/scripts/train_loop.sh

@@ -0,0 +1,42 @@
+# by default set it to CADFusion/data
+data_path=/your/path/to/data/folder
+# by default set it to CADFusion/exp
+exp_path=/your/path/to/exp/folder
+# by default set it to CADFusion/data
+exp_path=/your/path/to/vf_data/folder
+train_data=$data_path/train.json
+eval_data=$data_path/eval.json
+
+base_name=model_name_you_trained_for_SL
+
+run_name=${base_name}0
+CUDA_VISIBLE_DEVICES=0,1 ./scripts/inference.sh $run_name test "--full --device-map auto" &
+CUDA_VISIBLE_DEVICES=2,3 ./scripts/inference.sh $run_name train "--sample-len 1000 --device-map auto"
+wait
+
+./scripts/make_dpo_data.sh $run_name --score-only &
+./scripts/make_dpo_data.sh $run_name-train "--gpu 1"
+wait
+
+
+for LOOP in 1 2 3 4 5
+do
+    run_name=$base_name$LOOP    
+    dpo_training_path=$vf_path/$base_name$((LOOP-1))-train.json
+    dpo_run_name=$base_name$LOOP-dpo
+    dpo_save_path=$exp_path/$dpo_run_name
+    sft_run_name=$base_name$LOOP
+
+    python src/train/dpo.py --run-name $dpo_run_name --pretrained-path $exp_path/$base_name$((LOOP-1)) --data-path $dpo_training_path --output-path $dpo_save_path
+    python src/train/llama_finetune.py --num-epochs 1 --run-name $sft_run_name --data-path $train_data --eval-data-path $eval_data --eval-freq 3000 --pretrained-path $dpo_save_path --expdir $exp_path
+    
+    CUDA_VISIBLE_DEVICES=0 ./scripts/inference.sh $dpo_run_name test "--full --device-map auto" &
+    CUDA_VISIBLE_DEVICES=1 ./scripts/inference.sh $run_name test "--full --device-map auto" &
+    CUDA_VISIBLE_DEVICES=2,3 ./scripts/inference.sh $run_name train "--sample-len 1000 --device-map auto"
+    wait
+
+    ./scripts/make_dpo_data.sh $dpo_run_name --score-only &
+    ./scripts/make_dpo_data.sh $run_name "--score-only --gpu 1" &
+    ./scripts/make_dpo_data.sh $run_name-train "--gpu 2"
+    wait
+done

CADFusion/scripts/train_with_shuffling.sh

@@ -0,0 +1,20 @@
+# set it to your data path
+data_path=data/sl_data
+# set it to your experiment path
+exp_path=exp/model_ckpt
+train_data=$data_path/train.json
+eval_data=$data_path/val.json
+shuffle_dataset_between_x_epochs=2
+mkdir -p $exp_path
+
+# round 0
+accelerate launch --config_file ds_config.yaml src/train/llama_finetune.py --lora-rank 32 --lora-alpha 32 \
+        --num-epochs $shuffle_dataset_between_x_epochs --run-name $1 --data-path $train_data --eval-data-path $eval_data \
+        --device-map accelerate --eval-freq 1000 --save-freq 50000 --model-name llama3 --expdir $exp_path
+
+for round in 1 2 3 4 5 6 7 8 9
+do
+    python src/train/llama_finetune.py --lora-rank 32 --pretrained-path $exp_path/$1 --lora-alpha 32 \
+        --num-epochs $shuffle_dataset_between_x_epochs --run-name $1 --data-path $train_data --eval-data-path $eval_data \
+        --eval-freq 4000 --save-freq 50000 --expdir $exp_path
+done

CADFusion/src/data_preprocessing/call_openai.py

@@ -0,0 +1,37 @@
+from openai import AzureOpenAI
+from azure.identity import AzureCliCredential, get_bearer_token_provider
+
+import time
+
+def setup_client():
+    scope = "api://trapi/.default"
+    credential = get_bearer_token_provider(AzureCliCredential(), scope)
+
+    api_version = '2024-12-01-preview'
+    deployment_name = 'gpt-4o_2024-08-06'
+    instance = 'gcr/shared/' # See https://aka.ms/trapi/models for the instance name, remove /openai (library adds it implicitly)
+    endpoint = f'https://trapi.research.microsoft.com/{instance}'
+
+    client = AzureOpenAI(
+        azure_endpoint=endpoint,
+        azure_ad_token_provider=credential,
+        api_version=api_version,
+    )
+    return client, deployment_name
+
+
+def call_openai(client, deployment, prompt):
+    output = None
+    while output is None:
+        try:
+            time.sleep(0.5)
+            completion = client.chat.completions.create(
+                model = deployment,
+                messages = prompt,
+            )
+            output = completion.choices[0].message.content
+        except Exception as e:  
+            print("API error:", e)  
+            time.sleep(1)
+            output = None
+    return output

CADFusion/src/data_preprocessing/captioning.py

@@ -0,0 +1,101 @@
+import os
+import requests
+import base64
+import json
+import time
+from mimetypes import guess_type
+from tqdm import tqdm
+# from parse_sequence import parse_sequence
+# from parse_visual import run_parallel
+# from parse_image import render_file
+from call_openai import setup_client, call_openai
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--image-folder-path', type=str, default='exp/figures/test', help='Path to the input folder')
+parser.add_argument('--out-path', type=str, default='data/raw', help='Path to the output file')
+args = parser.parse_args()
+file_path = args.image_folder_path
+out_path = args.out_path
+
+client, deployment_name = setup_client()
+call_client = call_openai
+
+def local_image_to_data_url(image_path):
+    # Encode a local image into data URL
+    mime_type, _ = guess_type(image_path)
+    if mime_type is None:
+        mime_type = 'application/octet-stream' 
+    with open(image_path, "rb") as image_file:
+        base64_encoded_data = base64.b64encode(image_file.read()).decode('utf-8')
+    return f"data:{mime_type};base64,{base64_encoded_data}"
+
+def call_model_1(prompt, image_path):
+    message_text = [
+        {"role":"system","content":"You are an AI assistant that helps people find information."},
+        {"role":"user","content":[
+            {
+                "type": "text",
+                "text": prompt
+            },
+            {
+            "type": "image_url",
+            "image_url": {"url": local_image_to_data_url(image_path)}
+            }
+        ]}
+    ]
+    return call_client(client, deployment_name, message_text)
+
+def call_model_2(prompt1, image_path, output1, prompt2):
+    message_text = [
+        {"role":"system","content":"You are an AI assistant that helps people find information."},
+        {"role":"user","content":[
+            {
+                "type": "text",
+                "text": prompt1
+            },
+            {
+            "type": "image_url",
+            "image_url": {"url": local_image_to_data_url(image_path)}
+            }
+        ]},
+        {"role":"assistant","content":output1},
+        {"role":"user","content":prompt2}
+    ]
+    return call_client(client, deployment_name, message_text)
+
+files = [f for f in os.listdir(args.image_folder_path) if os.path.isfile(os.path.join(args.image_folder_path, f))]  
+files.sort()  
+results = []
+for filename in tqdm(files):
+    time.sleep(0.5)
+    output1 = None
+    output2 = None
+    image_path = os.path.join(file_path, filename)  
+    # Send request
+    prompt1 = """Propose a series of questions about the 3D shape and give the answers. The first question should ask for a detailed description and others should focus on the specific geometric properties, number, size proportions and positional relationship, and other details."""
+    prompt2 = """Based on the dialogue, please give a final description of the 3D shape. No more than 70 words."""
+    while output1 is None or str(output1).startswith("I'm sorry"):
+        try:
+            output1 = call_model_1(prompt1, image_path)
+        except requests.RequestException as e:  
+            print(f"Request failed: {e}")
+            time.sleep(1)  
+            output1 = None  
+    while output2 is None or str(output2).startswith("I'm sorry"):
+        try:
+            output2 = call_model_2(prompt1, image_path, output1, prompt2)
+        except requests.RequestException as e:  
+            print(f"Request failed: {e}")
+            time.sleep(1)  
+            output2 = None  
+
+    result = {
+        "pic_name":filename,
+        "questions": output1,
+        "description":output2
+    }
+    results.append(result)
+
+with open(out_path, 'w+', encoding='utf-8') as f:
+    json.dump(results, f, ensure_ascii=False, indent=4)

CADFusion/src/data_preprocessing/convert.py

@@ -0,0 +1,120 @@
+import pickle
+import argparse
+import json
+# hyperparameters from SkexGen project
+SKETCH_R = 1
+RADIUS_R = 1
+EXTRUDE_R = 1.0
+SCALE_R = 1.4
+OFFSET_R = 0.9
+PIX_PAD = 4
+CMD_PAD = 3
+COORD_PAD = 4
+EXT_PAD = 1
+EXTRA_PAD = 1
+R_PAD = 2
+
+
+def create_curve_str(se_xy, se_cmd):
+    curve_str = ""
+    xy_offset = 0
+    if se_cmd == 0:  # line
+        curve_str = " line," + ",".join(str(x) for x in se_xy[0])
+        xy_offset = 2
+    elif se_cmd == 1:  # arc
+        curve_str = " arc," + ",".join(str(x) for x in se_xy[0:2].flatten())
+        xy_offset = 3
+    elif se_cmd == 2:  # circle
+        curve_str = " circle," + ",".join(str(x) for x in se_xy[0:4].flatten())
+        xy_offset = 5
+    curve_str += " <curve_end>"
+    return curve_str, xy_offset
+
+
+def create_sketch_str(se_xy, se_cmd):
+    sketch_str = ""
+    len_xy, len_cmd = len(se_xy), len(se_cmd)
+    xy_idx = 0
+    for cmd_item in se_cmd:  # for each command
+        if 0 <= cmd_item <= 2:  # curve
+            curve_str, xy_offset = create_curve_str(se_xy[xy_idx:], cmd_item)
+            sketch_str += curve_str
+            xy_idx += xy_offset
+        elif cmd_item == -1:  # loop
+            sketch_str += " <loop_end>"
+            xy_idx += 1
+        elif cmd_item == -2:  # face
+            sketch_str += " <face_end>"
+            xy_idx += 1
+        elif cmd_item == -3:  # sketch
+            sketch_str += " <sketch_end>"
+            xy_idx += 1
+        else:
+            raise ValueError("Invalid command: " + str(cmd_item))
+    if xy_idx != len_xy:
+        raise ValueError("xy_idx != len_xy")
+    return sketch_str
+
+
+def create_extrude_str(se_ext):
+    extrude_str = ""
+    # extrude operation
+    if se_ext[14] == 1:
+        extrude_str += "add"
+    elif se_ext[14] == 2:
+        extrude_str += "cut"
+    elif se_ext[14] == 3:
+        extrude_str += "intersect"
+    else:
+        raise ValueError("Invalid extrude operation: " + str(se_ext[14]))
+    # other extrude parameters
+    extrude_str = (
+        extrude_str + "," + ",".join(str(x - EXT_PAD) for x in se_ext[0:5])
+    )  # ext_v, ext_T
+    extrude_str = (
+        extrude_str + "," + ",".join(str(x - R_PAD) for x in se_ext[5:14])
+    )  # ext_R
+    extrude_str = (
+        extrude_str + "," + ",".join(str(x - EXT_PAD) for x in se_ext[15:18])
+    )  # scale, offset
+    # extrude end
+    extrude_str += " <extrude_end>"
+    return extrude_str
+
+
+def convert(in_path, out_path):
+    with open(in_path, "rb") as f:
+        data = pickle.load(f)
+    print("Data loaded: " + str(len(data)) + " samples")
+
+    results = []
+    for item in data:  # for each data
+        se_str = ""
+        num_se = item["num_se"]
+        for se_idx in range(num_se):  # for each sketch-extrude
+            xy, cmd, ext = (
+                item["se_xy"][se_idx] - COORD_PAD,
+                item["se_cmd"][se_idx] - CMD_PAD,
+                item["se_ext"][se_idx],
+            )
+            se_str = se_str + " " + create_sketch_str(xy, cmd).strip()
+            se_str = se_str + " " + create_extrude_str(ext).strip()
+        results.append(se_str.strip())
+
+    # with open(out_path, "wb") as f:
+    #     pickle.dump(results, f)
+    # print("Data converted: " + str(len(results)) + " samples")
+    with open(out_path, "w") as f:
+        json.dump(results, f, indent=4)
+    print("Data converted: " + str(len(results)) + " samples")
+    # with open(out_path, "w") as f:  # Open in text mode  
+    #     for result in results:  
+    #         f.write(result + "\n") 
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--in-path", type=str, required=True)
+    parser.add_argument("--out-path", type=str, required=True)
+    args = parser.parse_args()
+
+    convert(args.in_path, args.out_path)

CADFusion/src/dpo/llava_utils.py

@@ -0,0 +1,95 @@
+import os
+import base64
+import time
+import json
+import requests
+from mimetypes import guess_type
+from transformers import pipeline
+from transformers import LlavaNextProcessor
+from transformers import AutoProcessor, LlavaOnevisionForConditionalGeneration
+import torch
+from PIL import Image
+dev='cuda:0'
+
+# processor = AutoProcessor.from_pretrained("llava-hf/llava-onevision-qwen2-7b-ov-hf") 
+# model = LlavaOnevisionForConditionalGeneration.from_pretrained("llava-hf/llava-onevision-qwen2-7b-ov-hf", torch_dtype=torch.float16, low_cpu_mem_usage=True) 
+# model.to(device)
+
+def restart_model(device):  
+    global dev
+    dev = device
+    processor = AutoProcessor.from_pretrained("llava-hf/llava-onevision-qwen2-7b-ov-hf") 
+    model = LlavaOnevisionForConditionalGeneration.from_pretrained("llava-hf/llava-onevision-qwen2-7b-ov-hf", torch_dtype=torch.float16, low_cpu_mem_usage=True) 
+    model.to(device)
+    return model, processor
+
+def ask_llm_on_figure(data, model, processor):
+    """
+    The layout of a typical data item
+    {
+        "index": 1,
+        "pic_name": "000000_001_final.png",
+        "ground_truth": "line,9,9 <curve_end> line,9,53 <curve_end> line,53,53 <curve_end> line,53,9 <curve_end> <loop_end> circle,31,29,31,20,35,25,27,25 <curve_end> <loop_end> circle,31,41,31,32,35,37,27,37 <curve_end> <loop_end> <face_end> <sketch_end> add,31,32,31,31,31,0,1,0,0,0,1,1,0,0,62,31,31 <extrude_end>",
+        "description": "Create a rectangular panel with two circular through-holes centrally aligned on the vertical axis.",
+        "prompt": "Below is a description of a 3D shape:\nCreate a rectangular panel with two circular through-holes centrally aligned on the vertical axis.\nGenerate a Computer-Aided Design(CAD) command sequence of the 3D shape:\n",
+        "output": "line,se,9 <curve_end> line,ne,9 <curve_end> line,ne,53 <curve_end> line,se,53 <curve_end> <loop_end> circle,22,41,22, Twenty1 ,31,30,12,30 <curve_end> <loop_end> circle,40,21,40, Ten2 ,50,32,29,32 <curve_end> <loop_end> <face_end> <sketch_end> add,31,33,31,31,31,1,0,0,0,0,1,0,-1,0,62,31,31 <extrude_end>"
+    },
+    """
+    url = data['figure_path']
+    image = Image.open(url)
+    description = data['description']
+    # data_scale = 10
+    # measurement = 'the degree of correspondence between them'
+
+    prompt = 'You are a harsh grader for new CAD designers\' works. The following is a text description of a CAD figure that they designed and an image of a CAD instance.' +\
+    f'\nDescription: {description}\n ' + \
+    f'Comment on this work for \n '+\
+    '1. If the overall shape remains correct; \n '+\
+    '2. If the number of components are correct, especially the circular holes; \n '+\
+    '3. If the distribution of the components are natural, i.e. they are not clustered together or collide with each other.\n'+\
+    'After that, give a score out of 10. Do not comment on issues such as texture, smoothness and colors'
+
+    
+    conversation = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": prompt},
+                {"type": "image"},
+            ],
+        },
+    ]
+    prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
+    inputs = processor(images=image, text=prompt, return_tensors="pt",).to(dev, torch.float16)
+
+    # autoregressively complete prompt
+    output = model.generate(**inputs, max_new_tokens=256, pad_token_id=processor.tokenizer.eos_token_id)
+    output = processor.decode(output[0], skip_special_tokens=True)
+    idx = output.index('assistant\n')
+    response = output[idx+10:]
+    return(response)
+
+
+def ask_llm(data, model, processor):
+    description = data['gpt_label']
+
+    prompt = 'The following is an evaluation of an CAD object.' +\
+        f'\n evaluation: {description}\n' +\
+        'Extract the integer score of the evaluation. The score is between 0 to 10. Return the number only.'
+    
+    conversation = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": prompt},
+            ],
+        },
+    ]
+    prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
+    inputs = processor(text=prompt, return_tensors="pt",).to(dev, torch.float16)
+
+    output = model.generate(**inputs, max_new_tokens=16, pad_token_id=processor.tokenizer.eos_token_id)
+    output = processor.decode(output[0], skip_special_tokens=True)
+    idx = output.index('assistant\n')
+    response = output[idx+10:]
+    return(response)

CADFusion/src/dpo/make_dpo_dataset.py

@@ -0,0 +1,162 @@
+import json
+import os
+import time
+import argparse
+
+from openai_utils import ask_gpt_on_figure, ask_gpt
+from llava_utils import ask_llm, ask_llm_on_figure, restart_model
+from tqdm import tqdm
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--source-data-path", type=str, required=True)
+    parser.add_argument("--figure-path", type=str, required=True)
+    parser.add_argument("--save-path", type=str, required=True)
+    parser.add_argument("--num-samples", type=int, required=True)
+    parser.add_argument("--gpu", type=int, default=0)
+    parser.add_argument("--score-only", action="store_true", default=False)
+    parser.add_argument("--gpt", action="store_true", default=False)
+    args = parser.parse_args()
+    
+    source_path = args.source_data_path
+    folder_path = args.figure_path
+    save_path = args.save_path
+    num_samples = args.num_samples
+    device=f'cuda:{args.gpu}'
+    if args.gpt:
+        func1, func2 = ask_gpt_on_figure, ask_gpt
+        model = None
+        processor = None
+    else:
+        func1, func2 = ask_llm_on_figure, ask_llm
+        model, processor = restart_model(device)
+    
+    with open(source_path, 'r') as f:
+        test_data = json.load(f)
+    
+    ####### Stage 1 #######
+    # for model generations that are able to render pictures,
+    # ask gpt to rate the generation quality.
+    for data in tqdm(test_data):
+        file_id = str(data['index']).zfill(6)
+        file = None
+        for f in os.listdir(folder_path):
+            if f.startswith(file_id):
+                file = folder_path + f
+                data['figure_path'] = file
+                error_cnt = 0
+                while 1:
+                    try:
+                        data['gpt_label'] = func1(data, model, processor)
+                        break
+                    except Exception as e:
+                        print(e)
+                        if args.gpt:
+                            time.sleep(3)
+                        else:
+                            if error_cnt == 5:
+                                exit()
+                            model, processor = restart_model(device)
+                    error_cnt += 1
+    with open(save_path, 'w+') as f:
+        json.dump(test_data, f, indent=4)   
+    
+    with open(save_path, 'r') as f:
+        test_data = json.load(f)
+    ####### Stage 2 #######
+    # clean up the dataset to summarize the generation quality estimation to a numerical score, and
+    # remove the failed ones, i.e. the generations that cannot render
+    for data in tqdm(test_data):
+        if "gpt_label" in data.keys():
+            error_cnt = 0
+            while 1:
+                try:
+                    score = func2(data, model, processor)
+                    print(score)
+                    break
+                except Exception as e:
+                    print(e)
+                    if args.gpt:
+                        time.sleep(3)
+                    else:
+                        if error_cnt == 5:
+                            exit()
+                        model, processor = restart_model(device)
+                error_cnt += 1
+            try:
+                data['gpt_score'] = int(score)
+            except:
+                print(f'ERROR: {score}')
+                pass    
+                
+    saved_data = [data for data in test_data if 'gpt_score' in data.keys()]     
+    with open(save_path, 'w+') as f:
+        json.dump(saved_data, f, indent=4)   
+    
+    if args.score_only:
+        exit()
+    
+    ####### Stage 3 #######
+    # 1. group up the scored generations by their description: we do not compare 
+    #    generation results that come from different origin prompts
+    temp_data = []
+    max_idx = test_data[-1]['index']
+    sample_size = max_idx // num_samples + 1
+    # a. select if any above 6
+
+    # for i in range(sample_size):
+    #     next_sample = test_data[i*num_samples:(i+1)*num_samples]
+    #     next_sample = [item for item in next_sample if 'gpt_score' in item.keys()]
+    #     above_score = [item['gpt_score'] >= 6 for item in next_sample]
+    #     if any(above_score):
+    #         temp_data.extend(next_sample)
+    # temp_data = [data for data in temp_data if 'gpt_score' in data.keys()] 
+    
+    # b. select if avg above 6
+
+    # for i in range(sample_size):
+    #     next_sample = test_data[i*num_samples:(i+1)*num_samples]
+    #     next_sample = [item for item in next_sample if 'gpt_score' in item.keys()]
+    #     if len(next_sample) == 0:
+    #         continue
+    #     scores = sum(item['gpt_score'] for item in next_sample) / len(next_sample)
+    #     if scores >= 6:
+    #         temp_data.extend(next_sample)
+    # temp_data = [data for data in temp_data if 'gpt_score' in data.keys()]
+
+    # c. select if individual above 6
+    test_data = saved_data
+    for item in test_data:
+        if 'gpt_score' not in item.keys():
+            continue
+        if item['gpt_score'] >= 6:
+            temp_data.append(item)
+    print(test_data[-1]['index'], max_idx)
+    
+    grouped = [[] for _ in range(max_idx)]
+    for item in temp_data:
+        idx = item['index']
+        grouped[idx // num_samples].append(item)
+    grouped = [item for item in grouped if len(item) > 0]
+    
+    # 2. within each group, make pairs where the chosens have higher score than the rejected ones.
+    # TODO: find a way to balance the data generated from each group
+    final_data = []
+    for group in grouped:
+        for item1 in group:
+            for item2 in group:
+                if item2['gpt_score'] > item1['gpt_score']:
+                    info_dict = {
+                        "description": item1['description'],
+                        "prompt": item1['prompt'],
+                        "chosen": item2['output'],
+                        "rejected": item1['output']
+                        }
+                    final_data.append(info_dict)
+                    # uncomment this break if you do not want too many data.
+                    # break
+            
+
+    with open(save_path, 'w+') as f:
+        json.dump(final_data, f, indent=4)   

CADFusion/src/dpo/openai_utils.py

@@ -0,0 +1,88 @@
+import os
+import base64
+import time
+import json
+
+from mimetypes import guess_type
+from openai import AzureOpenAI
+from azure.identity import DefaultAzureCredential, get_bearer_token_provider
+
+END_POINT = '<endpoint>'
+MODEL_NAME = 'gpt-4o_2024-08-06'
+API_VER = '2024-02-01'
+
+def local_image_to_data_url(image_path):
+    # Encode a local image into data URL
+    mime_type, _ = guess_type(image_path)
+    if mime_type is None:
+        mime_type = 'application/octet-stream' 
+    with open(image_path, "rb") as image_file:
+        base64_encoded_data = base64.b64encode(image_file.read()).decode('utf-8')
+    return f"data:{mime_type};base64,{base64_encoded_data}"
+
+def ask_gpt_on_figure(data, _, __):
+    endpoint = END_POINT
+    token_provider = get_bearer_token_provider(
+        DefaultAzureCredential(),
+        "https://cognitiveservices.azure.com/.default"
+    )
+    deployment_name = MODEL_NAME
+
+    client = AzureOpenAI(
+        azure_ad_token_provider=token_provider,
+        azure_endpoint=endpoint,
+        api_version=API_VER
+    )
+    description = data['description']
+    data_scale = 10
+    measurement = 'if the figure corresponds to the given description'
+
+    prompt = 'The following is a text description of a 3D CAD figure and an image of a CAD instance. ' +\
+        f'Measure {measurement}, and give a score in the scale of {data_scale}. Do not comment on issues such as texture, smoothness and colors' +\
+        f'\n description: {description}\n'
+    image_path = data['figure_path']
+    response = client.chat.completions.create(
+        model=deployment_name, 
+        messages=[
+            {'role': 'system', 'content': 'You are a helpful assistant'},
+            {'role': 'user', 'content': [
+                {'type': 'text', 'text': prompt},
+                {'type': 'image_url', 'image_url': {'url': local_image_to_data_url(image_path)}},  
+            ]}
+        ]
+    )
+    time.sleep(3)
+    return(response.choices[0].message.content)
+
+
+def ask_gpt(data, _, __):
+    endpoint = END_POINT
+    token_provider = get_bearer_token_provider(
+        DefaultAzureCredential(),
+        "https://cognitiveservices.azure.com/.default"
+    )
+    deployment_name = MODEL_NAME
+
+    client = AzureOpenAI(
+        azure_ad_token_provider=token_provider,
+        azure_endpoint=endpoint,
+        api_version=API_VER
+    )
+    description = data['gpt_label']
+
+    prompt = 'The following is an evaluation of an CAD object.' +\
+        f'\n evaluation: {description}\n' +\
+        'Extract the integer score of the evaluation. The score is between 0 to 10. Return the number only.'
+    
+    response = client.chat.completions.create(
+        model=deployment_name, 
+        messages=[
+            {'role': 'system', 'content': 'You are a helpful assistant'},
+            {'role': 'user', 'content': [            
+                {'type': 'text', 'text': prompt},
+            ]}
+        ]
+    )
+    # print(response.choices[0].message.content)
+    time.sleep(3)
+    return(response.choices[0].message.content)

CADFusion/src/rendering_utils/geometry/arc.py

@@ -0,0 +1,32 @@
+import numpy as np
+import math
+from geometry.curve import Curve
+
+
+class Arc(Curve):
+    def __init__(self, point_indices, point_data, is_outer):
+        assert len(point_indices) == 4, "Arc must be defined by 3 points"
+        assert point_data is not None
+        super(Arc, self).__init__(point_indices, point_data)
+        self.type = 'arc'
+        self.is_outer = is_outer
+        self.start = self.point_geom[0, :]
+        self.mid = self.point_geom[1, :]
+        self.center = self.point_geom[2, :]
+        self.end = self.point_geom[3, :]
+
+        self.r1 = math.sqrt( (self.start[0] - self.center[0])**2 + (self.start[1] - self.center[1])**2 )
+        self.r2 = math.sqrt( (self.end[0] - self.center[0])**2 + (self.end[1] - self.center[1])**2 )
+        self.radius = (self.r1+self.r2)/2
+
+        self.start_idx = point_indices[0]
+        self.mid_idx = point_indices[1]
+        self.center_idx = point_indices[2]
+        self.end_idx = point_indices[3]
+
+        self.bbox = self.verts_to_bbox(np.vstack([self.start, self.end, self.mid]))
+        self.bottom_left = np.array([self.bbox[0], self.bbox[2]])
+        
+    
+
+  

CADFusion/src/rendering_utils/geometry/circle.py

@@ -0,0 +1,27 @@
+import numpy as np
+from geometry.curve import Curve
+import pdb
+
+class Circle(Curve):
+    def __init__(self, point_indices, point_data, is_outer):
+        assert len(point_indices) == 2, "Circle must be defined by 1 points"
+        assert point_data is not None
+        super(Circle, self).__init__(point_indices, point_data)
+        self.type = 'circle'
+        self.center = self.point_geom[0, :]
+        self.radius = self.point_geom[1, 0]
+        self.center_idx = point_indices[0]
+        self.radius_idx = point_indices[1]
+        self.is_outer = is_outer
+
+        self.pt1 = np.array([self.center[0], self.center[1]+self.radius])
+        self.pt2 = np.array([self.center[0], self.center[1]-self.radius])
+        self.pt3 = np.array([self.center[0]+self.radius, self.center[1]])
+        self.pt4 = np.array([self.center[0]-self.radius, self.center[1]])
+        self.bbox = self.verts_to_bbox(np.vstack([self.pt1, self.pt2, self.pt3, self.pt4]))
+        self.bottom_left = np.array([self.bbox[0], self.bbox[2]])
+
+
+        
+
+   

CADFusion/src/rendering_utils/geometry/curve.py

@@ -0,0 +1,13 @@
+import math
+
+
+class Curve():
+    def __init__(self, point_indices, point_data):
+        self.point_indices = point_indices
+        self.point_geom = point_data[point_indices, 0:2]
+
+    def verts_to_bbox(self, verts):
+        xs = [v[0] for v in verts]
+        ys = [v[1] for v in verts]
+        bbox = [min(xs), max(xs), min(ys), max(ys)]
+        return bbox

CADFusion/src/rendering_utils/geometry/geom_utils.py

@@ -0,0 +1,95 @@
+import math
+import numpy as np
+
+def angle_from_vector_to_x(vec):
+    assert vec.size == 2
+    # We need to find a unit vector
+    angle = 0.0
+
+    l = np.linalg.norm(vec)
+    uvec = vec/l
+
+    # 2 | 1
+    #-------
+    # 3 | 4
+    if uvec[0] >=0:
+        if uvec[1] >= 0:
+            # Qadrant 1
+            angle = math.asin(uvec[1])
+        else:
+            # Qadrant 4
+            angle = 2.0*math.pi - math.asin(-uvec[1])
+    else:
+        if vec[1] >= 0:
+            # Qadrant 2
+            angle = math.pi - math.asin(uvec[1])
+        else:
+            # Qadrant 3
+            angle = math.pi + math.asin(-uvec[1])
+    return angle
+
+
+def convert_angle_to_1to360_range(angle_rad):
+    """
+    Converts the given angle in radians into 1-360 degrees range
+    """
+    angle = math.degrees(angle_rad)
+    # Lifted from: https://stackoverflow.com/questions/12234574/calculating-if-an-angle-is-between-two-angles
+    angle=(int(angle) % 360) + (angle-math.trunc(angle)) # converts angle to range -360 + 360
+    if angle > 0.0:
+        return angle
+    else:
+        return angle + 360.0
+
+
+def angle_is_between(angle_rad, a_rad, b_rad):
+    """
+    Checks if angle is in between the range of a and b
+    (All angles must be given in radians)
+    """
+    angle = convert_angle_to_1to360_range(angle_rad)
+    a = convert_angle_to_1to360_range(a_rad)
+    b = convert_angle_to_1to360_range(b_rad)
+    if a < b:
+        return a <= angle and angle <= b
+    return a <= angle or angle <= b
+
+
+def quantize_verts(verts, n_bits=8):
+    """Convert vertices in [-1., 1.] to discrete values in [0, n_bits**2 - 1]."""
+    min_range = -0.5
+    max_range = 0.5
+    range_quantize = 2 ** n_bits - 1
+    verts_quantize = (verts - min_range) * range_quantize / (max_range - min_range)
+    return verts_quantize.astype("int32")
+
+
+def dequantize_verts(verts, n_bits=8, add_noise=False):
+    """Convert quantized vertices to floats."""
+    min_range = -0.5
+    max_range = 0.5
+    range_quantize = 2 ** n_bits - 1
+    verts = verts.astype("float32")
+    verts = verts * (max_range - min_range) / range_quantize + min_range
+    if add_noise:
+        verts += np.random.uniform(size=verts.shape) * (1 / range_quantize)
+    return verts
+
+
+def center_vertices(vertices):
+    """Translate the vertices so that bounding box is centered at zero."""
+    vert_min = vertices.min(axis=0)
+    vert_max = vertices.max(axis=0)
+    vert_center = 0.5 * (vert_min + vert_max)
+    return vertices - vert_center, vert_center
+
+
+def scale_vertices(vertices):
+    """Scale the vertices so that the long diagonal of the bounding box is one."""
+    vert_min = vertices.min(axis=0)
+    vert_max = vertices.max(axis=0)
+    extents = vert_max - vert_min
+    scale = np.sqrt(np.sum(extents ** 2))
+    return vertices / scale, scale
+
+

CADFusion/src/rendering_utils/geometry/line.py

@@ -0,0 +1,24 @@
+import numpy as np
+from geometry.curve import Curve
+
+class Line(Curve):
+    def __init__(self, point_indices, point_data, is_outer):
+        assert len(point_indices) == 2, "Line must be defined by two points"
+        assert point_data is not None
+        super(Line, self).__init__(point_indices, point_data)
+        pt0 = self.point_geom[0, :]
+        pt1 = self.point_geom[1, :]
+        self.type = 'line'
+        self.start = pt0
+        self.end = pt1
+        self.start_idx = point_indices[0]
+        self.end_idx = point_indices[1]
+        self.is_outer = is_outer
+
+        self.bbox = self.verts_to_bbox(np.vstack([pt0, pt1]))
+        self.bottom_left = np.array([self.bbox[0], self.bbox[2]])
+
+        
+
+
+    

CADFusion/src/rendering_utils/geometry/obj_parser.py

@@ -0,0 +1,276 @@
+import os
+import sys
+import numpy as np
+
+from geometry.arc import Arc
+from geometry.circle import Circle
+from geometry.line import Line
+
+from geometry import geom_utils
+import pdb
+
+
+class OBJParser:
+    """
+    A class to read an OBJ file containing the sketch data
+    and hand it back in a form which is easy to work with.
+    """
+    def __init__(self, pathname=None):
+        self.pathname = pathname
+
+
+    def convert_vertices(self, vertices):
+        """Convert all the vertices to .obj format"""
+        vertex_strings = ""
+        for pt in vertices:
+            # e.g. v 0.123 0.234 0.345 1.0
+            vertex_string = f"v {pt[0]} {pt[1]}\n"
+            vertex_strings += vertex_string
+        return vertex_strings
+
+
+    def convert_curves(self, faces):
+        curve_strings = ""
+        total_curve = 0
+
+        # Faces (multiple closed regions)
+        for group_idx, loops in enumerate(faces):
+            curve_strings += f"\nface\n"
+            # Multiple loops (inner and outer)
+            for loop in loops: 
+                if loop[0].is_outer:  
+                    curve_strings += f"out\n"
+                else:
+                    curve_strings += f"in\n"
+                # All curves in one loop
+                for curve in loop:
+                    total_curve += 1
+                    if curve.type == 'line':
+                        curve_strings += f"l {curve.start_idx} {curve.end_idx}\n"
+                    elif curve.type == 'circle':
+                        curve_strings += f"c {curve.center_idx} {curve.radius_idx}\n"
+                    elif curve.type == 'arc':
+                        curve_strings += f"a {curve.start_idx} {curve.mid_idx} {curve.center_idx} {curve.end_idx}\n"
+
+        return curve_strings, total_curve
+
+
+    def parse3d(self, point3d):
+        x = point3d[0]
+        y = point3d[1]
+        z = point3d[2]
+        return str(x)+' '+str(y)+' '+str(z)
+
+
+    def write_obj2(self, file, vertices, faces, meta_info, scale=None):
+        """ Write to .obj file """
+        vertex_strings = self.convert_vertices(vertices)
+        curve_strings, total_curve = self.convert_curves(faces)
+        
+        with open(file, "w") as fh:
+            # Write Meta info
+            fh.write("# WaveFront *.obj file\n")
+            fh.write(f"# Vertices: {len(vertices)}\n")
+            fh.write(f"# Curves: {total_curve}\n")
+            fh.write("\n")
+
+            # Write vertex and curve
+            fh.write(vertex_strings)
+            fh.write("\n")
+            fh.write(curve_strings)
+            fh.write("\n")
+
+            #Write extrude value 
+            fh.write("ExtrudeOperation: " + meta_info['set_op']+"\n")
+            extrude_string = 'Extrude '
+            for value in meta_info['extrude_value']:
+                extrude_string += str(value)+' '
+            fh.write(extrude_string)
+            fh.write("\n")
+        
+            #Write refe plane transformation 
+            p_orig = self.parse3d(meta_info['t_orig'])
+            x_axis = self.parse3d(meta_info['t_x'])
+            y_axis = self.parse3d(meta_info['t_y'])
+            z_axis = self.parse3d(meta_info['t_z'])
+            fh.write('T_origin '+p_orig)
+            fh.write("\n")
+            fh.write('T_xaxis '+x_axis)
+            fh.write("\n")
+            fh.write('T_yaxis '+y_axis)
+            fh.write("\n")
+            fh.write('T_zaxis '+z_axis)
+            fh.write("\n")
+
+            # Normalized object 
+            if scale is not None:
+                fh.write('Scale '+str(scale))
+
+
+    def write_obj(self, file, curve_strings, total_curve, vertex_strings, total_v, meta_info, scale=None):
+        """ Write to .obj file """
+        #vertex_strings = self.convert_vertices(vertices)
+        #curve_strings, total_curve = self.convert_curves(faces)
+        
+        with open(file, "w") as fh:
+            # Write Meta info
+            fh.write("# WaveFront *.obj file\n")
+            fh.write(f"# Vertices: {total_v}\n")
+            fh.write(f"# Curves: {total_curve}\n")
+            fh.write("\n")
+
+            # Write vertex and curve
+            fh.write(vertex_strings)
+            fh.write("\n")
+            fh.write(curve_strings)
+            fh.write("\n")
+
+            #Write extrude value 
+            fh.write("ExtrudeOperation: " + meta_info['set_op']+"\n")
+            extrude_string = 'Extrude '
+            for value in meta_info['extrude_value']:
+                extrude_string += str(value)+' '
+            fh.write(extrude_string)
+            fh.write("\n")
+        
+            #Write refe plane transformation 
+            p_orig = self.parse3d(meta_info['t_orig'])
+            x_axis = self.parse3d(meta_info['t_x'])
+            y_axis = self.parse3d(meta_info['t_y'])
+            z_axis = self.parse3d(meta_info['t_z'])
+            fh.write('T_origin '+p_orig)
+            fh.write("\n")
+            fh.write('T_xaxis '+x_axis)
+            fh.write("\n")
+            fh.write('T_yaxis '+y_axis)
+            fh.write("\n")
+            fh.write('T_zaxis '+z_axis)
+            fh.write("\n")
+
+            # Normalized object 
+            if scale is not None:
+                fh.write('Scale '+str(scale))
+
+
+    def parse_file(self, scale=1.0):
+        """ 
+        Parse obj file
+        Return
+            vertex 2D location numpy
+            curve list (geometry class)
+            extrude parameters
+        """ 
+       
+        assert self.pathname is not None, "File is None"
+        assert self.pathname.exists(), "No such file"
+
+        # Parse file 
+        vertex_list = []
+        loops = []
+        closed_loop = []
+
+        # Read vertice
+        with open(self.pathname) as obj_file:
+            for line in obj_file:
+                tokens = line.split()
+                if not tokens:
+                    continue
+                line_type = tokens[0]
+                # Vertex
+                if line_type == "v":
+                    vertex_list.append([float(x) for x in tokens[1:]])
+        vertices = np.array(vertex_list, dtype=np.float64) * scale
+
+        # Read curves
+        faces = []
+        loops = []
+        loop = []
+        
+        # Read in all lines
+        lines = []
+        with open(self.pathname) as obj_file:
+            for line in obj_file:
+                lines.append(line)
+
+        # Parse all lines
+        faces = []
+        for str_idx, line in enumerate(lines):
+            tokens = line.split()
+            if not tokens:
+                continue
+            line_type = tokens[0]
+
+            # Start of a new face 
+            if line_type == "face":
+                faces.append(self.read_face(lines, str_idx+1, vertices))
+
+            # Read meta data
+            meta_data = line.strip('# ').strip(' \n').split(' ')
+            meta_name = meta_data[0]
+        
+            if meta_name == 'Extrude':
+                extrude_values = [float(x) for x in meta_data[1:]]
+                extrude_values = [x*scale for x in extrude_values]
+            elif meta_name == 'T_origin':
+                t_orig = [float(x) for x in meta_data[1:]] 
+                t_orig = [x*scale for x in t_orig] 
+            elif meta_name == 'T_xaxis':
+                t_x = [float(x) for x in meta_data[1:]] 
+            elif meta_name == 'T_yaxis':
+                t_y = [float(x) for x in meta_data[1:]] 
+            elif meta_name == 'T_zaxis':
+                t_z = [float(x) for x in meta_data[1:]] 
+            elif meta_name == 'ExtrudeOperation:':
+                set_op = meta_data[1]
+
+        meta_info = {'extrude_value': extrude_values,
+                     'set_op': set_op,
+                     't_orig': t_orig,
+                     't_x': t_x,
+                     't_y': t_y,
+                     't_z': t_z,
+                    }
+
+        return vertices, faces, meta_info    
+         
+
+
+    def read_face(self, lines, str_idx, vertices):
+        loops = []
+        loop = []
+        for line in lines[str_idx:]:
+            tokens = line.split()
+            if not tokens:
+                continue
+            line_type = tokens[0]
+
+            if line_type == 'face':
+                break
+
+            # Start of a new loop 
+            if line_type == "out" or line_type == "in":
+                if len(loop) > 0:
+                    loops.append(loop)
+                loop = []
+                is_outer = (line_type == 'out')
+
+            # Line
+            if line_type == 'l':
+                c_tok = tokens[1:]
+                curve = Line([int(c_tok[0]), int(c_tok[1])], vertices, is_outer=is_outer)
+                loop.append(curve)
+
+            # Arc 
+            if line_type == 'a':
+                c_tok = tokens[1:]
+                curve = Arc([int(c_tok[0]), int(c_tok[1]), int(c_tok[2]), int(c_tok[3])], vertices, is_outer=is_outer)
+                loop.append(curve)
+
+            # Circle 
+            if line_type == 'c':
+                c_tok = tokens[1:]
+                curve = Circle([int(c_tok[0]), int(c_tok[1])], vertices, is_outer=is_outer)
+                loop.append(curve)
+
+        loops.append(loop)
+        return loops

CADFusion/src/rendering_utils/geometry/obj_utils.py

@@ -0,0 +1,93 @@
+import json
+import numpy as np
+from pathlib import Path
+import pdb
+
+
+def read_wire_obj(obj_path):
+    """Read vertices and lines from .obj file defining a wire body."""
+    vertex_list = []
+    loops = []
+
+    # Read vertice and curves
+    with open(obj_path) as obj_file:
+
+        for line in obj_file:
+            tokens = line.split()
+            if not tokens:
+                continue
+
+            line_type = tokens[0]
+
+            if line_type == "v":
+                vertex_list.append([float(x) for x in tokens[1:]])
+
+            if line_type == "g":
+                pdb.set_trace()
+
+
+            
+       
+            # Read meta data
+            meta_data = line.strip('# ').strip(' \n').split(' ')
+            meta_name = meta_data[0]
+            if meta_name == 'Extrude':
+                extrude_values= [float(x) for x in meta_data[1:]]
+            elif meta_name == 'T_origin':
+                t_orig = [float(x) for x in meta_data[1:]] 
+            elif meta_name == 'T_xaxis':
+                t_x = [float(x) for x in meta_data[1:]] 
+            elif meta_name == 'T_yaxis':
+                t_y = [float(x) for x in meta_data[1:]] 
+            elif meta_name == 'T_zaxis':
+                t_z = [float(x) for x in meta_data[1:]] 
+            elif meta_name == 'ExtrudeOperation:':
+                set_op = meta_data[1]
+
+
+        vertices = np.array(vertex_list)
+
+
+
+        meta_info = {'extrude_value': extrude_values,
+                     'set_op': set_op,
+                     't_orig': t_orig,
+                     't_x': t_x,
+                     't_y': t_y,
+                     't_z': t_z}
+
+        total_in_outs.append(in_outs)
+        
+    return np.array(flat_vertices_list, dtype=np.float32), flat_hyperedge, total_in_outs, meta_info
+
+
+def write_wire_obj(vertices, faces, file_path, transpose=True, scale=1.0):
+    """Write vertices and hyperedges to obj."""
+    vertex_dimension = vertices.shape[1]
+    assert vertex_dimension in (2, 3)
+    if transpose and vertex_dimension == 3:
+        # Permute 3D vertices where z comes first followed by x and y
+        vertices = vertices[:, [1, 2, 0]]
+    vertices *= scale
+    if faces is not None:
+        if len(faces) > 0:
+            if min(min(faces)) == 0:
+                f_add = 1
+            else:
+                f_add = 0
+    with open(file_path, "w") as f:
+        for v in vertices:
+            if vertex_dimension == 2:
+                f.write("v {} {} {}\n".format(v[0], v[1], 0.0))
+            else:
+                f.write("v {} {} {}\n".format(v[0], v[1], v[2]))
+        for face in faces:
+            line = "l"
+            for i in face:
+                # Pradeep: always adding 1 to the face index makes sense to me. Not sure why
+                # PolyGen does this conditionally (see L95 above)
+                # Something to note.
+                line += " {}".format(i + 1)
+            line += "\n"
+            f.write(line)
+

CADFusion/src/rendering_utils/img_renderer.py

@@ -0,0 +1,84 @@
+import argparse
+from OCC.Core.Graphic3d import *
+from OCC.Display.OCCViewer import Viewer3d
+from OCC.Extend.DataExchange import read_step_file
+from OCC.Extend.TopologyUtils import TopologyExplorer
+from OCC.Core.Quantity import Quantity_Color, Quantity_TOC_RGB, Quantity_NOC_WHITE
+from OCC.Core.V3d import V3d_DirectionalLight
+from OCC.Core.gp import gp_Dir
+from glob import glob
+import pathlib
+from tqdm import tqdm
+
+
+def render(shape, filename, width=1024, height=768, face_color_rgb=(0.2, 0.2, 0.2), edge_color_rgb=(0, 0, 0), show_face_boundary=True):
+    viewer = Viewer3d()
+    viewer.Create(phong_shading=True, create_default_lights=True)
+    viewer.set_bg_gradient_color([255, 255, 255], [255, 255, 255])
+    viewer.SetModeShaded()
+    viewer.hide_triedron()
+    viewer.EnableAntiAliasing()
+    dir_light = V3d_DirectionalLight(gp_Dir(0, 0.5, -1), Quantity_Color(Quantity_NOC_WHITE))
+    dir_light.SetEnabled(True)
+    dir_light.SetIntensity(500.0)
+    viewer.Viewer.AddLight(dir_light)
+    viewer.Viewer.SetLightOn()
+
+    viewer.default_drawer.EnableDrawHiddenLine()
+    viewer.default_drawer.SetFaceBoundaryDraw(show_face_boundary)
+    ais_context = viewer.GetContext()
+    dc = ais_context.DeviationCoefficient()
+    da = ais_context.DeviationAngle()
+    factor = 10
+    ais_context.SetDeviationCoefficient(dc / factor)
+    ais_context.SetDeviationAngle(da / factor)
+    topexp = TopologyExplorer(shape)
+    for face in topexp.faces():
+        if face is not None:
+            viewer.DisplayShape(face, color=Quantity_Color(*face_color_rgb, Quantity_TOC_RGB))
+    for edge in topexp.edges():
+        if edge is not None:
+            viewer.DisplayShape(edge, color=Quantity_Color(*edge_color_rgb, Quantity_TOC_RGB))
+    viewer.FitAll()
+    viewer.SetSize(width, height)
+    viewer.View.Dump(str(filename))
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--input_dir", type=str, required=True, help="Input folder of STP/STEP files")
+    p.add_argument("--output_dir", type=str, required=True, help="Output folder of PNG files")
+    p.add_argument("--width", type=int, default=1024, help="Width of image")
+    p.add_argument("--height", type=int, default=768, help="Height of image")
+
+    args = p.parse_args()
+    
+    files = []
+    cad_folders = sorted(glob(args.input_dir+'/*/'))
+    for folder in cad_folders:
+        input_path = pathlib.Path(folder)
+        files += list(input_path.glob("*.st*p"))
+    print(len(files))
+    # files = files[36000:] # debug only (* remove *)
+    output_path = pathlib.Path(args.output_dir)
+    if not output_path.exists():
+        output_path.mkdir(parents=True, exist_ok=True)
+
+    i = 0
+    j = 0
+    for fn in tqdm(files):
+        j += 1
+        try:
+            shape = read_step_file(str(fn))
+            # render(shape, output_path.joinpath(f'{j:06d}' + ".png"), args.width, args.height)
+            render(shape, output_path.joinpath(fn.stem[:6] + ".png"), args.width, args.height)
+        except Exception as e:
+            i += 1
+            # raise e
+            print(e)
+            continue
+    print("error number: ", i)
+    print("total number: ", j)
+
+if __name__ == "__main__":
+    main()

CADFusion/src/rendering_utils/parser.py

@@ -0,0 +1,478 @@
+import numpy as np
+from collections import OrderedDict
+import re
+from pathlib import Path
+import argparse
+import os
+import json
+import math
+
+# hyperparameters from SkexGen project
+SKETCH_R = 1
+RADIUS_R = 1
+EXTRUDE_R = 1.0
+SCALE_R = 1.4
+OFFSET_R = 0.9
+PIX_PAD = 4
+CMD_PAD = 3
+COORD_PAD = 4
+EXT_PAD = 1
+EXTRA_PAD = 1
+R_PAD = 2
+
+
+class CADparser:
+    """Parse CAD sequence to CAD object."""
+
+    def __init__(self, bit):
+        self.vertex_dict = OrderedDict()
+        self.bit = bit
+
+    def perform(self, cad_seq):
+        # divide into sketch and extrude
+        sketches, extrudes = self.get_SE(cad_seq)
+        if sketches is None or extrudes is None:
+            return None
+        # sequentially parse each pair of SE into obj
+        se_datas = []
+        for sketch, extrude in zip(sketches, extrudes):
+            extrude_param, scale, offset = self.parse_extrude(extrude)
+            if extrude_param is None or scale is None or offset is None:
+                return None
+            vertex_str, se_str = self.parse_sketch(sketch, scale, offset)
+            if vertex_str is None or se_str is None:
+                return None
+            se_datas.append(
+                {"vertex": vertex_str, "curve": se_str, "extrude": extrude_param}
+            )
+            self.vertex_dict.clear()
+
+        return se_datas
+
+    def parse_sketch(self, sketch, scale, offset):
+        faces = self.get_faces(sketch)
+        if len(faces) == 0:
+            return None, None
+        se_str = ""
+        for face_idx, face in enumerate(faces):  # each face
+            face_str = "face\n"
+            loops = self.get_loops(face)
+            if len(loops) == 0:
+                return None, None
+            for loop_idx, loop in enumerate(loops):  # each loop
+                curves = self.get_curves(loop)
+                if len(curves) == 0:
+                    return None, None
+                next_curves = curves[1:]
+                next_curves += curves[:1]
+                cur_str = []
+                for curve, next_curve in zip(curves, next_curves):  # each curve
+                    if not self.obj_curve(curve, next_curve, cur_str, scale, offset):
+                        return None, None
+                loop_str = ""
+                for c in cur_str:
+                    loop_str += f"{c}\n"
+                if loop_idx == 0:
+                    face_str += f"out\n{loop_str}\n"
+                else:
+                    face_str += f"in\n{loop_str}\n"
+            se_str += face_str
+        vertex_str = self.convert_vertices()
+        return vertex_str, se_str
+
+    def parse_extrude(self, extrude):
+        ext = extrude.split(",")
+        if len(ext) != 18:
+            return None, None, None
+
+        # operation str to int
+        ext_op = {"add": 1, "cut": 2, "intersect": 3}.get(ext[0], None)
+        if ext_op is None:
+            return None, None, None
+        # dequantize ext_v, ext_T, scale and offset
+        ext_v, ext_T, scale, offset = self.dequantize_extrude_params(ext)
+        # get ext_R
+        ext_R = np.array(ext[6:15], dtype=int)
+
+        extrude_param = {"value": ext_v, "T": ext_T, "R": ext_R, "op": ext_op}
+        return extrude_param, scale, offset
+
+    def obj_curve(self, curve, next_curve, cur_str, scale, offset):
+        cur = curve.split(",")
+        next_cur = next_curve.split(",")
+        if cur[0] == "circle":
+            if len(cur) != 9:
+                return False
+            p1, p2, p3, p4 = self.dequantize_circle_points(
+                cur, next_cur, scale, offset)
+            center = np.asarray([0.5 * (p1[0] + p2[0]), 0.5 * (p3[1] + p4[1])])
+            radius = (np.linalg.norm(p1 - p2) + np.linalg.norm(p3 - p4)) / 4.0
+
+            center = center * scale + offset
+            radius = radius * scale
+
+            center_idx = self.save_vertex(center[0], center[1], "p")
+            radius_idx = self.save_vertex(radius, 0.0, "r")
+            cur_str.append(f"c {center_idx} {radius_idx}")
+        elif cur[0] == "arc":
+            if len(cur) != 5:
+                return False
+            if (
+                cur[1:3] == cur[3:5]
+                or cur[1:3] == next_cur[1:3]
+                or cur[3:5] == next_cur[3:5]
+            ):  # invalid arc
+                return False
+            start_v, mid_v, end_v = self.dequantize_arc_points(
+                cur, next_cur, scale, offset
+            )
+            try:
+                center, _, _, _ = find_arc_geometry(start_v, mid_v, end_v)
+            except Exception:
+                return False
+            start_v = start_v * scale + offset
+            mid_v = mid_v * scale + offset
+            end_v = end_v * scale + offset
+            center = center * scale + offset
+
+            center_idx = self.save_vertex(center[0], center[1], "p")
+            start_idx = self.save_vertex(start_v[0], start_v[1], "p")
+            mid_idx = self.save_vertex(mid_v[0], mid_v[1], "p")
+            end_idx = self.save_vertex(end_v[0], end_v[1], "p")
+            cur_str.append(f"a {start_idx} {mid_idx} {center_idx} {end_idx}")
+        elif cur[0] == "line":
+            if len(cur) != 3:
+                return False
+            if cur[1:3] == next_cur[1:3]:
+                return False
+            start_v, end_v = self.dequantize_line_points(
+                cur, next_cur, scale, offset)
+            start_v = start_v * scale + offset
+            end_v = end_v * scale + offset
+
+            start_idx = self.save_vertex(start_v[0], start_v[1], "p")
+            end_idx = self.save_vertex(end_v[0], end_v[1], "p")
+            cur_str.append(f"l {start_idx} {end_idx}")
+        else:
+            return False
+        return True
+
+    def get_SE(self, cad_seq):
+        # sketches: 1) between sequence start and sketch_end,
+        sketches_from_start = re.findall(r"^(.+?)(?=<sketch_end>)", cad_seq)
+        # sketches: 2) between extrude_end and sketch_end
+        sketches_after_extrude = re.findall(
+            r"(?<=<extrude_end>)(.+?)(?=<sketch_end>)", cad_seq
+        )
+        sketches = [x.strip() for x in sketches_from_start] + [
+            x.strip() for x in sketches_after_extrude
+        ]
+        # extrudes: between sketch_end and extrude_end
+        extrudes = [
+            x.strip() for x in re.findall(r"<sketch_end>(.+?)<extrude_end>", cad_seq)
+        ]
+        if len(sketches) != len(extrudes):
+            return None, None
+        return sketches, extrudes
+
+    def get_faces(self, sketch):
+        faces = sketch.split("<face_end>")
+        return [x.strip() for x in faces if x.strip() != ""]
+
+    def get_loops(self, face):
+        loops = face.split("<loop_end>")
+        return [x.strip() for x in loops if x.strip() != ""]
+
+    def get_curves(self, loop):
+        curves = loop.split("<curve_end>")
+        return [x.strip() for x in curves if x.strip() != ""]
+
+    def dequantize_circle_points(self, curve, next_curve, scale, offset):
+        p1 = dequantize_verts(
+            np.array(curve[1:3], dtype=int),
+            n_bits=self.bit,
+            min_range=-SKETCH_R,
+            max_range=SKETCH_R,
+            add_noise=False,
+        )
+        p2 = dequantize_verts(
+            np.array(curve[3:5], dtype=int),
+            n_bits=self.bit,
+            min_range=-SKETCH_R,
+            max_range=SKETCH_R,
+            add_noise=False,
+        )
+        p3 = dequantize_verts(
+            np.array(curve[5:7], dtype=int),
+            n_bits=self.bit,
+            min_range=-SKETCH_R,
+            max_range=SKETCH_R,
+            add_noise=False,
+        )
+        p4 = dequantize_verts(
+            np.array(curve[7:9], dtype=int),
+            n_bits=self.bit,
+            min_range=-SKETCH_R,
+            max_range=SKETCH_R,
+            add_noise=False,
+        )
+        return p1, p2, p3, p4
+
+    def dequantize_arc_points(self, curve, next_curve, scale, offset):
+        start_v = dequantize_verts(
+            np.array(curve[1:3], dtype=int),
+            n_bits=self.bit,
+            min_range=-SKETCH_R,
+            max_range=SKETCH_R,
+            add_noise=False,
+        )
+        mid_v = dequantize_verts(
+            np.array(curve[3:5], dtype=int),
+            n_bits=self.bit,
+            min_range=-SKETCH_R,
+            max_range=SKETCH_R,
+            add_noise=False,
+        )
+        end_v = dequantize_verts(
+            np.array(next_curve[1:3], dtype=int),
+            n_bits=self.bit,
+            min_range=-SKETCH_R,
+            max_range=SKETCH_R,
+            add_noise=False,
+        )
+        return start_v, mid_v, end_v
+
+    def dequantize_line_points(self, curve, next_curve, scale, offset):
+        start_v = dequantize_verts(
+            np.array(curve[1:3], dtype=int),
+            n_bits=self.bit,
+            min_range=-SKETCH_R,
+            max_range=SKETCH_R,
+            add_noise=False,
+        )
+        end_v = dequantize_verts(
+            np.array(next_curve[1:3], dtype=int),
+            n_bits=self.bit,
+            min_range=-SKETCH_R,
+            max_range=SKETCH_R,
+            add_noise=False,
+        )
+        return start_v, end_v
+
+    def dequantize_extrude_params(self, extrude):
+        ext_v = dequantize_verts(
+            np.array(extrude[1:3], dtype=int),
+            n_bits=self.bit,
+            min_range=-EXTRUDE_R,
+            max_range=EXTRUDE_R,
+            add_noise=False,
+        )
+        ext_T = dequantize_verts(
+            np.array(extrude[3:6], dtype=int),
+            n_bits=self.bit,
+            min_range=-EXTRUDE_R,
+            max_range=EXTRUDE_R,
+            add_noise=False,
+        )
+        scale = dequantize_verts(
+            np.array(extrude[15], dtype=int),
+            n_bits=self.bit,
+            min_range=0.0,
+            max_range=SCALE_R,
+            add_noise=False,
+        )
+        offset = dequantize_verts(
+            np.array(extrude[16:18], dtype=int),
+            n_bits=self.bit,
+            min_range=-OFFSET_R,
+            max_range=OFFSET_R,
+            add_noise=False,
+        )
+        return ext_v, ext_T, scale, offset
+
+    def save_vertex(self, h_x, h_y, text):
+        unique_key = f"{text}:x{h_x}y{h_y}"
+        index = 0
+        for key in self.vertex_dict.keys():
+            # Vertex location already exist in dict
+            if unique_key == key:
+                return index
+            index += 1
+        # Vertex location does not exist in dict
+        self.vertex_dict[unique_key] = [h_x, h_y]
+        return index
+
+    def convert_vertices(self):
+        """Convert all the vertices to .obj format"""
+        vertex_strings = ""
+        for pt in self.vertex_dict.values():
+            # e.g. v 0.123 0.234 0.345 1.0
+            vertex_string = f"v {pt[0]} {pt[1]}\n"
+            vertex_strings += vertex_string
+        return vertex_strings
+
+
+def find_arc_geometry(a, b, c):
+    A = b[0] - a[0]
+    B = b[1] - a[1]
+    C = c[0] - a[0]
+    D = c[1] - a[1]
+
+    E = A*(a[0] + b[0]) + B*(a[1] + b[1])
+    F = C*(a[0] + c[0]) + D*(a[1] + c[1])
+
+    G = 2.0*(A*(c[1] - b[1])-B*(c[0] - b[0]))
+
+    if G == 0:
+        raise Exception("zero G")
+
+    p_0 = (D*E - B*F) / G
+    p_1 = (A*F - C*E) / G
+
+    center = np.array([p_0, p_1])
+    radius = np.linalg.norm(center - a)
+
+    angles = []
+    for xx in [a, b, c]:
+        angle = angle_from_vector_to_x(xx - center)
+        angles.append(angle)
+
+    ab = b-a
+    ac = c-a
+    cp = np.cross(ab, ac)
+    if cp >= 0:
+        start_angle_rads = angles[0]
+        end_angle_rads = angles[2]
+    else:
+        start_angle_rads = angles[2]
+        end_angle_rads = angles[0]
+
+    return center, radius, start_angle_rads, end_angle_rads
+
+
+def angle_from_vector_to_x(vec):
+    assert vec.size == 2
+    # We need to find a unit vector
+    angle = 0.0
+
+    l = np.linalg.norm(vec)
+    uvec = vec/l
+
+    # 2 | 1
+    # -------
+    # 3 | 4
+    if uvec[0] >= 0:
+        if uvec[1] >= 0:
+            # Qadrant 1
+            angle = math.asin(uvec[1])
+        else:
+            # Qadrant 4
+            angle = 2.0*math.pi - math.asin(-uvec[1])
+    else:
+        if vec[1] >= 0:
+            # Qadrant 2
+            angle = math.pi - math.asin(uvec[1])
+        else:
+            # Qadrant 3
+            angle = math.pi + math.asin(-uvec[1])
+    return angle
+
+
+def dequantize_verts(verts, n_bits=8, min_range=-0.5, max_range=0.5, add_noise=False):
+    """Convert quantized vertices to floats."""
+    range_quantize = 2**n_bits - 1
+    verts = verts.astype("float32")
+    verts = verts * (max_range - min_range) / range_quantize + min_range
+    return verts
+
+
+def write_obj_sample(save_folder, data):
+    for idx, write_data in enumerate(data):
+        obj_name = Path(save_folder).stem + "_" + \
+            str(idx).zfill(3) + "_param.obj"
+        obj_file = Path(save_folder) / obj_name
+        extrude_param = write_data["extrude"]
+        vertex_strings = write_data["vertex"]
+        curve_strings = write_data["curve"]
+
+        """Write an .obj file with the curves and verts"""
+        if extrude_param["op"] == 1:  # 'add'
+            set_op = "NewBodyFeatureOperation"
+        elif extrude_param["op"] == 2:  # 'cut'
+            set_op = "CutFeatureOperation"
+        elif extrude_param["op"] == 3:  # 'cut'
+            set_op = "IntersectFeatureOperation"
+
+        with open(obj_file, "w") as fh:
+            # Write Meta info
+            fh.write("# WaveFront *.obj file\n")
+            fh.write("# ExtrudeOperation: " + set_op + "\n")
+            fh.write("\n")
+
+            # Write vertex and curve
+            fh.write(vertex_strings)
+            fh.write("\n")
+            fh.write(curve_strings)
+            fh.write("\n")
+
+            # Write extrude value
+            extrude_string = "Extrude "
+            for value in extrude_param["value"]:
+                extrude_string += str(value) + " "
+            fh.write(extrude_string)
+            fh.write("\n")
+
+            # Write refe plane value
+            p_orig = parse3d_sample(extrude_param["T"])
+            x_axis = parse3d_sample(extrude_param["R"][0:3])
+            y_axis = parse3d_sample(extrude_param["R"][3:6])
+            z_axis = parse3d_sample(extrude_param["R"][6:9])
+            fh.write("T_origin " + p_orig)
+            fh.write("\n")
+            fh.write("T_xaxis " + x_axis)
+            fh.write("\n")
+            fh.write("T_yaxis " + y_axis)
+            fh.write("\n")
+            fh.write("T_zaxis " + z_axis)
+
+
+def parse3d_sample(point3d):
+    x = point3d[0]
+    y = point3d[1]
+    z = point3d[2]
+    return str(x) + " " + str(y) + " " + str(z)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--in-path", type=str, required=True)
+    parser.add_argument("--out-path", type=str, required=True)
+    args = parser.parse_args()
+
+    # with open(args.in_path, "r") as f:
+        # data = f.readlines()
+    with open(args.in_path, 'r') as file:  
+        data = file.read()
+ 
+    data = json.loads(data)
+
+    num_valid_str = 0
+    for idx, item in enumerate(data):
+        try:
+            cad_parser = CADparser(bit=6)
+            # print(idx)
+            if type(item) == str:
+                parsed_data = cad_parser.perform(item)
+            elif type(item) == dict:
+                parsed_data = cad_parser.perform(item['output'])
+            else:
+                raise ValueError("Invalid data type")
+            out_path = os.path.join(args.out_path, str(idx).zfill(6))
+            os.makedirs(out_path, exist_ok=True)
+            if parsed_data is not None:
+                num_valid_str += 1
+                write_obj_sample(out_path, parsed_data)
+        except Exception as e:
+            print(e)
+            pass
+    print(f"Number of valid CAD strings: {num_valid_str}/{len(data)}")

CADFusion/src/rendering_utils/parser_visual.py

@@ -0,0 +1,110 @@
+import os 
+import argparse
+from pathlib import Path
+from tqdm import tqdm
+from multiprocessing import Pool
+from glob import glob 
+from utils.obj_reconverter import OBJReconverter
+from OCC.Core.BRepCheck import BRepCheck_Analyzer
+from geometry.obj_parser import OBJParser
+from utils.util import write_stl_file
+from OCC.Extend.DataExchange import write_step_file 
+
+import signal
+from contextlib import contextmanager
+@contextmanager
+def timeout(time):
+    # Register a function to raise a TimeoutError on the signal.
+    signal.signal(signal.SIGALRM, raise_timeout)
+    # Schedule the signal to be sent after ``time``.
+    signal.alarm(time)
+    try:
+        yield
+    except TimeoutError:
+        raise Exception("time out")
+    finally:
+        # Unregister the signal so it won't be triggered
+        # if the timeout is not reached.
+        signal.signal(signal.SIGALRM, signal.SIG_IGN)
+def raise_timeout(signum, frame):
+    raise TimeoutError
+    
+NUM_TRHEADS = 36 
+
+def find_files(folder, extension):
+    return sorted([Path(os.path.join(folder, f)) for f in os.listdir(folder) if f.endswith(extension)])
+
+
+def run_parallel(project_folder):
+    output_folder = project_folder
+
+    param_objs = find_files(project_folder, 'param.obj')
+
+    cur_solid = None
+    extrude_idx = 0
+    for obj in param_objs:
+        try:
+          with timeout(30):
+            parser = OBJParser(obj)
+            _, faces, meta_info = parser.parse_file(1.0)
+            converter = OBJReconverter()
+            ext_solid, _, _ = converter.parse_obj(faces, meta_info)
+            set_op = meta_info["set_op"]
+            if set_op == "NewBodyFeatureOperation" or set_op == "JoinFeatureOperation":
+                if cur_solid is None:
+                    cur_solid = ext_solid
+                else:
+                    cur_solid = converter.my_op(cur_solid, ext_solid, 'fuse')
+            elif set_op == "CutFeatureOperation":
+                cur_solid = converter.my_op(cur_solid, ext_solid, 'cut')
+            elif set_op == "IntersectFeatureOperation":
+                cur_solid = converter.my_op(cur_solid, ext_solid, 'common')
+            else:
+                raise Exception("Unknown operation type")
+
+            analyzer = BRepCheck_Analyzer(cur_solid)
+            if not analyzer.IsValid():
+                raise Exception("brep check failed")
+            
+            extrude_idx += 1
+        
+        except Exception as ex:
+            print(ex)
+            msg = [project_folder, str(ex)[:100]]
+            return None 
+    try:
+      with timeout(30):
+        stl_name = Path(output_folder).stem + '_'+ str(extrude_idx).zfill(3) + "_final.stl"
+        output_path =  os.path.join(output_folder, stl_name)
+        write_stl_file(cur_solid, output_path, linear_deflection=0.001, angular_deflection=0.5)
+
+        step_name = Path(output_folder).stem + '_'+ str(extrude_idx).zfill(3) + "_final.step"
+        output_path =  os.path.join(output_folder, step_name)
+        write_step_file(cur_solid, output_path)
+
+    except Exception as ex:
+        print(ex)
+        msg = [project_folder, str(ex)[:500]]
+        return None 
+
+    return cur_solid 
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data_folder", type=str, required=True)
+    parser.add_argument("--single-file", action='store_true', default=False)
+    args = parser.parse_args()
+
+    if args.single_file:
+    # If single file, just run the function on that file
+        run_parallel(args.data_folder)
+        exit(0)
+    else:
+        solids = []
+        # cad_folders = sorted(glob(args.data_folder+'/*'))[50000:] # why after 50000?
+        cad_folders = sorted(glob(args.data_folder+'/*'))
+        # print("len of cad_folder:", len(cad_folders))
+        convert_iter = Pool(NUM_TRHEADS).imap(run_parallel, cad_folders) 
+        for solid in tqdm(convert_iter, total=len(cad_folders)):
+            pass

CADFusion/src/rendering_utils/ptl_sampler.py

@@ -0,0 +1,88 @@
+import os
+import argparse
+import ntpath
+from tqdm import tqdm
+import multiprocessing
+from pathlib import Path
+from glob import glob
+import trimesh
+from trimesh.sample import sample_surface
+from plyfile import PlyData, PlyElement
+import numpy as np
+
+
+def write_ply(points, filename, text=False):
+    """ input: Nx3, write points to filename as PLY format. """
+    points = [(points[i,0], points[i,1], points[i,2]) for i in range(points.shape[0])]
+    vertex = np.array(points, dtype=[('x', 'f4'), ('y', 'f4'),('z', 'f4')])
+    el = PlyElement.describe(vertex, 'vertex', comments=['vertices'])
+    with open(filename, mode='wb') as f:
+        PlyData([el], text=text).write(f)
+
+
+def find_files(folder, extension):
+    return sorted([Path(os.path.join(folder, f)) for f in os.listdir(folder) if f.endswith(extension)])
+
+class SamplePoints:
+    """
+    Perform sampleing of points.
+    """
+
+    def __init__(self):
+        """
+        Constructor.
+        """
+        parser = self.get_parser()
+        self.options = parser.parse_args()
+
+
+    def get_parser(self):
+        """
+        Get parser of tool.
+
+        :return: parser
+        """
+        parser = argparse.ArgumentParser(description='Scale a set of meshes stored as OFF files.')
+        parser.add_argument('--in_dir', type=str, help='Path to input directory.')
+        parser.add_argument('--out_dir', type=str, help='Path to output directory; files within are overwritten!')
+        parser.add_argument("--single-file", action='store_true', default=False)
+        return parser
+
+
+    def run_parallel(self, project_folder):
+        out_folder = os.path.join(project_folder, self.options.out_dir)
+        if not os.path.exists(out_folder):
+            os.makedirs(out_folder)
+
+        files = find_files(project_folder, 'final.stl')
+        
+        for filepath in files:
+            N_POINTS = 2000
+            try:
+                out_mesh = trimesh.load(str(filepath))
+                out_pc, _ = sample_surface(out_mesh, N_POINTS)
+                save_path = os.path.join(out_folder, ntpath.basename(filepath)[:-4]+'_pcd.ply')
+                write_ply(out_pc, save_path)
+
+            except Exception as ex:        
+                return project_folder
+        return
+
+
+    def run(self):
+        """
+        Run simplification.
+        """
+        if self.options.single_file:
+            self.run_parallel(self.options.in_dir)
+        else:
+            project_folders = sorted(glob(self.options.in_dir+'/*/'))
+            num_cpus = multiprocessing.cpu_count()
+            convert_iter =  multiprocessing.Pool(num_cpus).imap(self.run_parallel, project_folders) 
+            for _ in tqdm(convert_iter, total=len(project_folders)):
+                pass
+       
+
+if __name__ == '__main__':
+    app = SamplePoints()
+    app.run()

CADFusion/src/rendering_utils/utils/obj_reconverter.py

@@ -0,0 +1,437 @@
+import numpy as np
+from collections import OrderedDict
+from .util import create_point, create_unit_vec, get_transform, create_sketch_plane
+
+# OCC
+from OCC.Core.BRepCheck import BRepCheck_Analyzer
+from OCC.Core.GC import GC_MakeArcOfCircle
+from OCC.Core.BRepBuilderAPI import (
+    BRepBuilderAPI_MakeFace,
+    BRepBuilderAPI_MakeWire,
+    BRepBuilderAPI_MakeEdge,
+)
+from OCC.Core.BRepAlgoAPI import BRepAlgoAPI_Fuse, BRepAlgoAPI_Cut, BRepAlgoAPI_Common
+from OCC.Core.BRepPrimAPI import BRepPrimAPI_MakePrism
+from OCC.Core.BRepAdaptor import BRepAdaptor_Surface
+from OCC.Core.BRepGProp import brepgprop_VolumeProperties, brepgprop_SurfaceProperties
+from OCC.Core.GProp import GProp_GProps
+from OCC.Core.ShapeFix import ShapeFix_Face, ShapeFix_Wire
+from OCC.Core.gp import gp_Vec, gp_Ax2, gp_Dir, gp_Circ
+from OCC.Extend.DataExchange import write_stl_file
+
+
+class OBJReconverter:
+    """OBJ Data Reconverter"""
+
+    def __init__(self):
+        self.vertex_dict = OrderedDict()
+        self.PRECISION = 1e-5
+        self.eps = 1e-7
+        self.x_axis = gp_Dir(1.0, 0.0, 0.0)
+
+    def convert_curve(self, curve):
+        """
+        convert to json dict format
+        """
+        json_curve = {}
+
+        if curve.type == "circle":
+            json_curve["type"] = "Circle3D"
+            json_curve["center_point"] = {
+                "x": curve.center[0],
+                "y": curve.center[1],
+                "z": 0,
+            }
+            json_curve["radius"] = curve.radius
+
+        if curve.type == "line":
+            json_curve["type"] = "Line3D"
+            json_curve["start_point"] = {
+                "x": curve.start[0],
+                "y": curve.start[1],
+                "z": 0,
+            }
+            json_curve["end_point"] = {"x": curve.end[0], "y": curve.end[1], "z": 0}
+
+        if curve.type == "arc":
+            json_curve["type"] = "Arc3D"
+            json_curve["start_point"] = {
+                "x": curve.start[0],
+                "y": curve.start[1],
+                "z": 0,
+            }
+            json_curve["end_point"] = {"x": curve.end[0], "y": curve.end[1], "z": 0}
+            json_curve["mid_point"] = {"x": curve.mid[0], "y": curve.mid[1], "z": 0}
+            json_curve["center_point"] = {
+                "x": curve.center[0],
+                "y": curve.center[1],
+                "z": 0,
+            }
+
+        json_curve["is_outer"] = curve.is_outer
+        return json_curve
+
+    def convert_vertices(self):
+        """Convert all the vertices to .obj format"""
+        vertex_strings = ""
+        for pt in self.vertex_dict.values():
+            # e.g. v 0.123 0.234 0.345 1.0
+            vertex_string = f"v {pt[0]} {pt[1]}\n"
+            vertex_strings += vertex_string
+        return vertex_strings
+
+    def parse_obj(self, faces, meta_info):
+        """
+        reconstruct brep from obj file
+        """
+        # At least one needs to match
+        for face in faces:
+            for loop in face:
+                if len(loop) > 1:
+                    for idx, curve in enumerate(loop[:-1]):
+                        next_curve = np.vstack([loop[idx + 1].start, loop[idx + 1].end])
+                        diff1 = np.sum(np.abs(curve.start - next_curve), 1)
+                        diff2 = np.sum(np.abs(curve.end - next_curve), 1)
+
+                        if min(diff2) == 0 or min(diff1) == 0:
+                            continue  # edge connected
+
+                        assert (
+                            min(diff1) < 1e-3 or min(diff2) < 1e-3
+                        )  # difference should be small
+
+                        if min(diff1) > min(diff2):
+                            min_idx = np.argmin(diff2)
+                            if min_idx == 0:
+                                loop[idx + 1].start_idx = curve.end_idx
+                                loop[idx + 1].start = curve.end
+                            else:
+                                loop[idx + 1].end_idx = curve.end_idx
+                                loop[idx + 1].end = curve.end
+                        else:
+                            min_idx = np.argmin(diff1)
+                            if min_idx == 0:
+                                loop[idx + 1].start_idx = curve.start_idx
+                                loop[idx + 1].start = curve.start
+                            else:
+                                loop[idx + 1].end_idx = curve.start_idx
+                                loop[idx + 1].end = curve.start
+
+                    # Solve start / end connection
+                    shared_idx = list(
+                        set([loop[-2].start_idx, loop[-2].end_idx]).intersection(
+                            set([loop[-1].start_idx, loop[-1].end_idx])
+                        )
+                    )
+
+                    assert len(shared_idx) >= 1
+
+                    if len(shared_idx) == 2:
+                        assert len(loop) == 2  # do nothing
+                    else:
+                        if shared_idx[0] == loop[-1].start_idx:
+                            do_start = False
+                        else:
+                            do_start = True
+                        start_curve = np.vstack([loop[0].start, loop[0].end])
+
+                        if do_start:
+                            diff = np.sum(np.abs(loop[-1].start - start_curve), 1)
+                        else:
+                            diff = np.sum(np.abs(loop[-1].end - start_curve), 1)
+                        assert min(diff) < 1e-3
+
+                        min_idx = np.argmin(diff)
+                        if min_idx == 0:
+                            if do_start:
+                                loop[-1].start_idx = loop[0].start_idx
+                                loop[-1].start = loop[0].start
+                            else:
+                                loop[-1].end_idx = loop[0].start_idx
+                                loop[-1].end = loop[0].start
+                        else:
+                            if do_start:
+                                loop[-1].start_idx = loop[0].end_idx
+                                loop[-1].start = loop[0].end
+                            else:
+                                loop[-1].end_idx = loop[0].end_idx
+                                loop[-1].end = loop[0].end
+
+        # Parse groups to json loop/curve profile
+        extrusion = {}
+        extrusion["profiles"] = []
+        for face in faces:
+            profile = {}
+            profile["loops"] = []
+            for loop in face:
+                pl = {}
+                pl["profile_curves"] = []
+                for curve in loop:
+                    # convert to json format
+                    pl["profile_curves"].append(self.convert_curve(curve))
+                profile["loops"].append(pl)
+            extrusion["profiles"].append(profile)
+
+        # Parse transform
+        sketch = {}
+        transform = {}
+        transform["origin"] = {
+            "x": meta_info["t_orig"][0],
+            "y": meta_info["t_orig"][1],
+            "z": meta_info["t_orig"][2],
+        }
+        transform["x_axis"] = {
+            "x": meta_info["t_x"][0],
+            "y": meta_info["t_x"][1],
+            "z": meta_info["t_x"][2],
+        }
+        transform["y_axis"] = {
+            "x": meta_info["t_y"][0],
+            "y": meta_info["t_y"][1],
+            "z": meta_info["t_y"][2],
+        }
+        transform["z_axis"] = {
+            "x": meta_info["t_z"][0],
+            "y": meta_info["t_z"][1],
+            "z": meta_info["t_z"][2],
+        }
+        sketch["transform"] = transform
+
+        # Parse extrude
+        extrude_params = {}
+        extrude_params["extrude_type"] = meta_info["set_op"]
+        extrude_params["extrude_values"] = meta_info["extrude_value"]
+
+        # Create sketch
+        all_faces = []
+        curve_strings = ""
+        curve_count = 0
+        for profile in extrusion["profiles"]:
+            ref_face, face, curve_string, c_count = self.parse_sketch(sketch, profile)
+            curve_strings += curve_string
+            curve_count += c_count
+            all_faces.append(face)
+
+        # Merge all faces in the same plane
+        plane_face = all_faces[0]
+        for face in all_faces[1:]:
+            plane_face = self.my_op(plane_face, face, "fuse")
+        solid = self.extrude_face(ref_face, plane_face, extrude_params)
+        return solid, curve_strings, curve_count
+
+    def my_op(self, big, small, op_name):
+        if op_name == "cut":
+            op = BRepAlgoAPI_Cut(big, small)
+        elif op_name == "fuse":
+            op = BRepAlgoAPI_Fuse(big, small)
+        elif op_name == "common":
+            op = BRepAlgoAPI_Common(big, small)
+        op.SetFuzzyValue(self.PRECISION)
+        op.Build()
+        return op.Shape()
+
+    def build_body(self, face, normal, value):
+        extrusion_vec = gp_Vec(normal).Multiplied(value)
+        make_prism = BRepPrimAPI_MakePrism(face, extrusion_vec)
+        make_prism.Build()
+        prism = make_prism.Prism()
+        return prism.Shape()
+
+    def extrudeBasedOnType(self, face, normal, distance):
+        # Extrude based on the two bound values
+        if not (distance[0] < distance[1]):
+            raise Exception("incorrect distance")
+        large_value = distance[1]
+        small_value = distance[0]
+
+        if large_value == 0:
+            return self.build_body(face, -normal, -small_value)
+        elif small_value == 0:
+            return self.build_body(face, normal, large_value)
+        elif np.sign(large_value) == np.sign(small_value):
+            if large_value < 0:
+                body1 = self.build_body(face, -normal, -small_value)
+                body2 = self.build_body(face, -normal, -large_value)
+                return self.my_op(body1, body2, "cut")
+            else:
+                assert large_value > 0
+                body1 = self.build_body(face, normal, small_value)
+                body2 = self.build_body(face, normal, large_value)
+                return self.my_op(body2, body1, "cut")
+        else:
+            assert np.sign(large_value) != np.sign(small_value)
+            body1 = self.build_body(face, normal, large_value)
+            body2 = self.build_body(face, -normal, -small_value)
+            return self.my_op(body1, body2, "fuse")
+
+    def extrude_face(self, ref_face, face, extrude_params):
+        distance = extrude_params["extrude_values"]
+        surf = BRepAdaptor_Surface(ref_face).Plane()
+        normal = surf.Axis().Direction()
+        extruded_shape = self.extrudeBasedOnType(face, normal, distance)
+        return extruded_shape
+
+    def parse_sketch(self, sketch, profile):
+        """
+        Sketch in one closed loop (one out, multiple ins)
+        """
+        # Transformation from local to global xyz coord
+        transform = get_transform(sketch["transform"])
+
+        # Create face region (automatically infer from all wires)
+        outer_facelist = []
+        inner_facelist = []
+        curve_count = 0
+        outer_string = []
+        inner_string = []
+        plane = create_sketch_plane(sketch["transform"])
+
+        for idx, pl in enumerate(profile["loops"]):
+            # Create loop
+            loop, curve_string, num_curve = self.parse_loop(
+                pl["profile_curves"], transform
+            )
+            # Create face
+            face_builder = BRepBuilderAPI_MakeFace(plane, loop)
+            if not face_builder.IsDone():
+                raise Exception("face builder not done")
+            face = face_builder.Face()
+            # Fix face
+            fixer = ShapeFix_Face(face)
+            fixer.SetPrecision(self.PRECISION)
+            fixer.FixOrientation()
+
+            analyzer = BRepCheck_Analyzer(fixer.Face())
+            if not analyzer.IsValid():
+                raise Exception("face check failed")
+
+            curve_count += num_curve
+
+            if pl["profile_curves"][0]["is_outer"]:
+                outer_facelist.append(fixer.Face())
+                outer_string.append(curve_string)
+            else:
+                inner_facelist.append(fixer.Face())
+                inner_string.append(curve_string)
+
+        # Create final closed loop face
+        assert len(outer_facelist) > 0
+        final_face = outer_facelist[0]
+        for face in outer_facelist[1:]:
+            final_face = self.my_op(final_face, face, "fuse")
+        for face in inner_facelist:
+            final_face = self.my_op(final_face, face, "cut")
+
+        # Append inner outer information to string
+        assert len(outer_string) == 1
+        out_str = ""
+        in_str = ""
+        for c_str in outer_string:
+            out_str += "out\n" + c_str + "\n"
+        for c_str in inner_string:
+            in_str += "in\n" + c_str + "\n"
+        final_str = "face\n" + out_str + in_str
+
+        return outer_facelist[0], final_face, final_str, curve_count
+
+    def parse_loop(self, profile_loop, transform):
+        """Create face in one closed loop"""
+        topo_wire = BRepBuilderAPI_MakeWire()
+        curve_strings = ""
+        curve_count = 0
+
+        # Loop through all the curves in one loop
+        for profile_curve in profile_loop:
+            curve_edge, curve_string = self.parse_curve(profile_curve, transform)
+            topo_wire.Add(curve_edge)
+            if not topo_wire.IsDone():
+                raise Exception("wire builder not done")
+
+            curve_string += "\n"
+            curve_count += 1
+            curve_strings += curve_string
+
+        fixer = ShapeFix_Wire()
+        fixer.Load(topo_wire.Wire())
+        fixer.SetPrecision(self.PRECISION)
+        fixer.FixClosed()
+        fixer.Perform()
+        return fixer.Wire(), curve_strings, curve_count
+
+    def parse_curve(self, curve, transform):
+        if curve["type"] == "Line3D":
+            return self.create_line(curve, transform)
+        elif curve["type"] == "Circle3D":
+            return self.create_circle(curve, transform)
+        elif curve["type"] == "Arc3D":
+            return self.create_arc(curve, transform)
+        else:
+            raise Exception("unknown curve type")
+
+    def create_line(self, line, transform):
+        start = create_point(line["start_point"], transform)
+        end = create_point(line["end_point"], transform)
+        if start.Distance(end) == 0:
+            raise Exception("start/end point same location")
+        topo_edge = BRepBuilderAPI_MakeEdge(start, end)
+
+        # Save pre-transform
+        star_idx = self.save_vertex(
+            line["start_point"]["x"] + 0.0, line["start_point"]["y"] + 0.0, "p"
+        )
+        end_idx = self.save_vertex(
+            line["end_point"]["x"] + 0.0, line["end_point"]["y"] + 0.0, "p"
+        )
+        curve_string = f"l {star_idx} {end_idx}"
+        return topo_edge.Edge(), curve_string
+
+    def create_arc(self, arc, transform):
+        start = create_point(arc["start_point"], transform)
+        mid = create_point(arc["mid_point"], transform)
+        end = create_point(arc["end_point"], transform)
+        arc_occ = GC_MakeArcOfCircle(start, mid, end).Value()
+        topo_edge = BRepBuilderAPI_MakeEdge(arc_occ)
+
+        # Save pre-transform
+        start_idx = self.save_vertex(
+            arc["start_point"]["x"] + 0.0, arc["start_point"]["y"] + 0.0, "p"
+        )
+        end_idx = self.save_vertex(
+            arc["end_point"]["x"] + 0.0, arc["end_point"]["y"] + 0.0, "p"
+        )
+        center_idx = self.save_vertex(
+            arc["center_point"]["x"] + 0.0, arc["center_point"]["y"] + 0.0, "p"
+        )
+        mid_idx = self.save_vertex(
+            arc["mid_point"]["x"] + 0.0, arc["mid_point"]["y"] + 0.0, "p"
+        )
+        curve_string = f"a {start_idx} {mid_idx} {center_idx} {end_idx}"
+        return topo_edge.Edge(), curve_string
+
+    def create_circle(self, circle, transform):
+        center = create_point(circle["center_point"], transform)
+        radius = circle["radius"]
+        normal = create_unit_vec({"x": 0.0, "y": 0.0, "z": 1.0}, transform)
+        ref_vector3d = self.x_axis.Transformed(transform)
+        axis = gp_Ax2(center, normal, ref_vector3d)
+        gp_circle = gp_Circ(axis, abs(float(radius)))
+        topo_edge = BRepBuilderAPI_MakeEdge(gp_circle)
+
+        center_idx = self.save_vertex(
+            circle["center_point"]["x"] + 0.0, circle["center_point"]["y"] + 0.0, "p"
+        )
+        radius_idx = self.save_vertex(abs(float(radius)) + 0.0, 0, "r")
+        curve_string = f"c {center_idx} {radius_idx}"
+        return topo_edge.Edge(), curve_string
+
+    def save_vertex(self, h_x, h_y, text):
+        unique_key = f"{text}:x{h_x}y{h_y}"
+        index = 0
+        for key in self.vertex_dict.keys():
+            # Vertex location already exist in dict
+            if unique_key == key:
+                return index
+            index += 1
+        # Vertex location does not exist in dict
+        self.vertex_dict[unique_key] = [h_x, h_y]
+        return index

CADFusion/src/rendering_utils/utils/util.py

@@ -0,0 +1,72 @@
+import os
+from OCC.Core.gp import gp_Pnt, gp_Vec, gp_Dir, gp_XYZ, gp_Ax3, gp_Trsf, gp_Pln
+from OCC.Core.BRepMesh import BRepMesh_IncrementalMesh
+from OCC.Core.StlAPI import StlAPI_Writer
+
+def create_xyz(xyz):
+    return gp_XYZ(xyz["x"], xyz["y"], xyz["z"])
+
+
+def get_ax3(transform_dict):
+    origin = create_xyz(transform_dict["origin"])
+    x_axis = create_xyz(transform_dict["x_axis"])
+    y_axis = create_xyz(transform_dict["y_axis"])
+    z_axis = create_xyz(transform_dict["z_axis"])
+    # Create new coord (orig, Norm, x-axis)
+    axis3 = gp_Ax3(gp_Pnt(origin), gp_Dir(z_axis), gp_Dir(x_axis))
+    return axis3
+
+
+def get_transform(transform_dict):
+    axis3 = get_ax3(transform_dict)
+    transform_to_local = gp_Trsf()
+    transform_to_local.SetTransformation(axis3)
+    return transform_to_local.Inverted()
+
+
+def create_sketch_plane(transform_dict):
+    axis3 = get_ax3(transform_dict)
+    return gp_Pln(axis3)
+
+
+def create_point(point_dict, transform):
+    pt2d = gp_Pnt(point_dict["x"], point_dict["y"], point_dict["z"])
+    return pt2d.Transformed(transform)
+
+
+def create_unit_vec(vec_dict, transform):
+    vec2d = gp_Dir(vec_dict["x"], vec_dict["y"], vec_dict["z"])
+    return vec2d.Transformed(transform)
+
+
+def write_stl_file(a_shape, filename, mode="ascii", linear_deflection=0.001, angular_deflection=0.5):
+    """ export the shape to a STL file
+    Be careful, the shape first need to be explicitely meshed using BRepMesh_IncrementalMesh
+    a_shape: the topods_shape to export
+    filename: the filename
+    mode: optional, "ascii" by default. Can either be "binary"
+    linear_deflection: optional, default to 0.001. Lower, more occurate mesh
+    angular_deflection: optional, default to 0.5. Lower, more accurate_mesh
+    """
+    if a_shape.IsNull():
+        raise AssertionError("Shape is null.")
+    if mode not in ["ascii", "binary"]:
+        raise AssertionError("mode should be either ascii or binary")
+    if os.path.isfile(filename):
+        print("Warning: %s file already exists and will be replaced" % filename)
+    # first mesh the shape
+    mesh = BRepMesh_IncrementalMesh(a_shape, linear_deflection, False, angular_deflection, True)
+    #mesh.SetDeflection(0.05)
+    mesh.Perform()
+    if not mesh.IsDone():
+        raise AssertionError("Mesh is not done.")
+
+    stl_exporter = StlAPI_Writer()
+    if mode == "ascii":
+        stl_exporter.SetASCIIMode(True)
+    else:  # binary, just set the ASCII flag to False
+        stl_exporter.SetASCIIMode(False)
+    stl_exporter.Write(a_shape, filename)
+
+    if not os.path.isfile(filename):
+        raise IOError("File not written to disk.")

CADFusion/src/test/VLM_score.py

@@ -0,0 +1,95 @@
+import os
+import requests
+import base64
+import json
+import time
+import argparse
+from mimetypes import guess_type
+from tqdm import tqdm
+import re
+
+from openai import AzureOpenAI
+from azure.identity import AzureCliCredential, get_bearer_token_provider
+
+scope = "api://trapi/.default"
+credential = get_bearer_token_provider(AzureCliCredential(),scope)
+
+api_version = '2024-12-01-preview'
+# deployment_name = 'gpt-4.1-mini_2025-04-14'
+deployment_name = 'gpt-4o_2024-08-06'
+instance = '<trapi/path>' # See https://aka.ms/trapi/models for the instance name, remove /openai (library adds it implicitly)
+endpoint = f'https://trapi.research.microsoft.com/{instance}'
+
+client = AzureOpenAI(
+    azure_endpoint=endpoint,
+    azure_ad_token_provider=credential,
+    api_version=api_version,
+)
+
+def local_image_to_data_url(image_path):
+    mime_type, _ = guess_type(image_path)
+    if mime_type is None:
+        mime_type = 'application/octet-stream'
+    with open(image_path, "rb") as image_file:
+        base64_encoded_data = base64.b64encode(image_file.read()).decode('utf-8')
+    return f"data:{mime_type};base64,{base64_encoded_data}"
+
+def ask_gpt(image_path, prompt):
+    image_url = local_image_to_data_url(image_path)
+    message_text = [
+            {"role": "system", "content": "You are an AI assistant that helps people find information."},
+            {"role": "user", "content": [
+                {"type": "text", "text": prompt},
+                {"type": "image_url", "image_url": {"url": image_url}},
+            ]}
+        ]
+
+    completion = client.chat.completions.create(
+            model=deployment_name,
+            messages=message_text,)
+    output = completion.choices[0].message.content
+    return output
+
+if __name__ == '__main__':
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--test-path', type=str, default='data/sl_data/test.jsonl', help='Path to the JSONL file containing test data')
+    parser.add_argument('--name', type=str, default='original_seq', help='Run name of the testee')
+    parser.add_argument('--figure-dir', type=str, default='exp/figures')
+    parser.add_argument('--save-path', type=str, default='exp/evals', help='Target folder to save the results')
+    parser.add_argument('--repetition', type=int, default=5, help='Number of repetitions for each image')
+    args = parser.parse_args()
+
+    results = []
+    jsonl_path = args.test_path
+    name = args.name
+    figures_dir = f"{args.figure_dir}/{name}/"
+    save_path = f"{args.save_path}/{name}.jsonl"
+
+    with open(jsonl_path, 'r+') as file:
+        test_data = json.load(file)
+    repetition = args.repetition
+    results = []
+    for i in tqdm(range(len(test_data[:800]))):
+        item = test_data[i]
+        for j in range(repetition):
+            img_num = i * repetition + j
+            image_name = f"{img_num:06d}.png"
+            image_path = os.path.join(figures_dir, image_name)
+            if os.path.exists(image_path):
+                description = item['description']
+                try:
+                    score = ask_gpt(image_path, f"The following is a text description of a 3D CAD figure and an image of a CAD instance. Measure if the figure corresponds to the given description, and give a score in the scale of 10. Only return the score. Do not comment on issues such as texture, smoothness and colors.\n description:{description}\n")
+
+                    # "The following is an original image of a CAD instance, a text description on editing and an image of the edited result. Measure if the figure corresponds to the given description, and give a score in the scale of 10. Only return the score. Do not comment on issues such as texture, smoothness and colors.\n description:{description}\n"
+                except Exception as e:
+                    print(img_num)
+                    print(e)
+                    score = -1
+                result = {
+                    "index": img_num,
+                    "gpt_score": score
+                }
+                results.append(result)
+                with open(save_path, 'w+') as file:
+                    json.dump(results, file, indent=4)

CADFusion/src/test/chamfer_dist.py

@@ -0,0 +1,308 @@
+import torch
+import argparse
+import os
+import numpy as np
+from tqdm import tqdm
+import random
+import warnings
+from glob import glob
+from scipy.stats import entropy
+from sklearn.neighbors import NearestNeighbors
+from plyfile import PlyData
+from pathlib import Path
+from multiprocessing import Pool
+from chamfer_distance import ChamferDistance
+
+random.seed(0)
+N_POINTS = 2000
+NUM_TRHEADS = 16
+
+
+def find_files(folder, extension):
+    return sorted([Path(os.path.join(folder, f)) for f in os.listdir(folder) if f.endswith(extension)])
+
+
+def read_ply(path):
+    with open(path, 'rb') as f:
+        plydata = PlyData.read(f)
+        x = np.array(plydata['vertex']['x'])
+        y = np.array(plydata['vertex']['y'])
+        z = np.array(plydata['vertex']['z'])
+        vertex = np.stack([x, y, z], axis=1)
+    return vertex
+
+
+def distChamfer(a, b):
+    x, y = a, b
+    bs, num_points, points_dim = x.size()
+    xx = torch.bmm(x, x.transpose(2, 1))
+    yy = torch.bmm(y, y.transpose(2, 1))
+    zz = torch.bmm(x, y.transpose(2, 1))
+    diag_ind = torch.arange(0, num_points).to(a).long()
+    rx = xx[:, diag_ind, diag_ind].unsqueeze(1).expand_as(xx)
+    ry = yy[:, diag_ind, diag_ind].unsqueeze(1).expand_as(yy)
+    P = (rx.transpose(2, 1) + ry - 2 * zz)
+    return P.min(1)[0], P.min(2)[0]
+
+
+def _pairwise_CD(sample_pcs, ref_pcs, batch_size):
+    N_sample = sample_pcs.shape[0]
+    N_ref = ref_pcs.shape[0]
+    all_cd = []
+    all_emd = []
+    iterator = range(N_sample)
+    matched_gt = []
+    pbar = tqdm(iterator)
+    chamfer_dist = ChamferDistance()
+
+    for sample_b_start in pbar:
+        sample_batch = sample_pcs[sample_b_start]
+
+        cd_lst = []
+        emd_lst = []
+        for ref_b_start in range(0, N_ref, batch_size):
+            ref_b_end = min(N_ref, ref_b_start + batch_size)
+            ref_batch = ref_pcs[ref_b_start:ref_b_end]
+
+            batch_size_ref = ref_batch.size(0)
+            sample_batch_exp = sample_batch.view(1, -1, 3).expand(batch_size_ref, -1, -1)
+            sample_batch_exp = sample_batch_exp.contiguous()
+            
+            dl, dr, idx1, idx2 = chamfer_dist(sample_batch_exp,ref_batch)
+            cd_lst.append((dl.mean(dim=1) + dr.mean(dim=1)).view(1, -1))
+
+        cd_lst = torch.cat(cd_lst, dim=1)
+        all_cd.append(cd_lst)
+
+        hit = np.argmin(cd_lst.detach().cpu().numpy()[0])
+        matched_gt.append(hit)
+        pbar.set_postfix({"cov": len(np.unique(matched_gt)) * 1.0 / N_ref})
+
+    all_cd = torch.cat(all_cd, dim=0)  # N_sample, N_ref
+
+    return all_cd
+
+
+def compute_cov_mmd(sample_pcs, ref_pcs, batch_size):
+    all_dist = _pairwise_CD(sample_pcs, ref_pcs, batch_size)
+    print(all_dist.shape, flush=True)
+    N_sample, N_ref = all_dist.size(0), all_dist.size(1)
+    min_val_fromsmp, min_idx = torch.min(all_dist, dim=1)
+    min_val, _ = torch.min(all_dist, dim=0)
+    mmd = min_val.mean()
+    cov = float(min_idx.unique().view(-1).size(0)) / float(N_ref)
+    cov = torch.tensor(cov).to(all_dist)
+
+    return {
+        # 'med-CD': torch.diagonal(all_dist).median().item(),
+        'avg-CD': torch.diagonal(all_dist).mean().item(),
+        'COV-CD': cov.item(),
+        'MMD-CD': mmd.item()
+    }
+
+
+def jsd_between_point_cloud_sets(sample_pcs, ref_pcs, in_unit_sphere, resolution=28):
+    '''Computes the JSD between two sets of point-clouds, as introduced in the paper ```Learning Representations And Generative Models For 3D Point Clouds```.
+    Args:
+        sample_pcs: (np.ndarray S1xR2x3) S1 point-clouds, each of R1 points.
+        ref_pcs: (np.ndarray S2xR2x3) S2 point-clouds, each of R2 points.
+        resolution: (int) grid-resolution. Affects granularity of measurements.
+    '''
+    sample_grid_var = entropy_of_occupancy_grid(sample_pcs, resolution, in_unit_sphere)[1]
+    ref_grid_var = entropy_of_occupancy_grid(ref_pcs, resolution, in_unit_sphere)[1]
+    return jensen_shannon_divergence(sample_grid_var, ref_grid_var)
+
+
+def entropy_of_occupancy_grid(pclouds, grid_resolution, in_sphere=False):
+    '''Given a collection of point-clouds, estimate the entropy of the random variables
+    corresponding to occupancy-grid activation patterns.
+    Inputs:
+        pclouds: (numpy array) #point-clouds x points per point-cloud x 3
+        grid_resolution (int) size of occupancy grid that will be used.
+    '''
+    epsilon = 10e-4
+    bound = 1 + epsilon
+    if abs(np.max(pclouds)) > bound or abs(np.min(pclouds)) > bound:
+        print(abs(np.max(pclouds)), abs(np.min(pclouds)))
+        warnings.warn('Point-clouds are not in unit cube.')
+
+    if in_sphere and np.max(np.sqrt(np.sum(pclouds ** 2, axis=2))) > bound:
+        warnings.warn('Point-clouds are not in unit sphere.')
+
+    grid_coordinates, _ = unit_cube_grid_point_cloud(grid_resolution, in_sphere)
+    grid_coordinates = grid_coordinates.reshape(-1, 3)
+    grid_counters = np.zeros(len(grid_coordinates))
+    grid_bernoulli_rvars = np.zeros(len(grid_coordinates))
+    nn = NearestNeighbors(n_neighbors=1).fit(grid_coordinates)
+
+    for pc in pclouds:
+        _, indices = nn.kneighbors(pc)
+        indices = np.squeeze(indices)
+        for i in indices:
+            grid_counters[i] += 1
+        indices = np.unique(indices)
+        for i in indices:
+            grid_bernoulli_rvars[i] += 1
+
+    acc_entropy = 0.0
+    n = float(len(pclouds))
+    for g in grid_bernoulli_rvars:
+        p = 0.0
+        if g > 0:
+            p = float(g) / n
+            acc_entropy += entropy([p, 1.0 - p])
+
+    return acc_entropy / len(grid_counters), grid_counters
+
+
+def unit_cube_grid_point_cloud(resolution, clip_sphere=False):
+    '''Returns the center coordinates of each cell of a 3D grid with resolution^3 cells,
+    that is placed in the unit-cube.
+    If clip_sphere it True it drops the "corner" cells that lie outside the unit-sphere.
+    '''
+    grid = np.ndarray((resolution, resolution, resolution, 3), np.float32)
+    spacing = 1.0 / float(resolution - 1) * 2
+    for i in range(resolution):
+        for j in range(resolution):
+            for k in range(resolution):
+                grid[i, j, k, 0] = i * spacing - 0.5 * 2
+                grid[i, j, k, 1] = j * spacing - 0.5 * 2
+                grid[i, j, k, 2] = k * spacing - 0.5 * 2
+
+    if clip_sphere:
+        grid = grid.reshape(-1, 3)
+        grid = grid[np.linalg.norm(grid, axis=1) <= 0.5]
+
+    return grid, spacing
+
+
+def jensen_shannon_divergence(P, Q):
+    if np.any(P < 0) or np.any(Q < 0):
+        raise ValueError('Negative values.')
+    if len(P) != len(Q):
+        raise ValueError('Non equal size.')
+
+    P_ = P / np.sum(P)  # Ensure probabilities.
+    Q_ = Q / np.sum(Q)
+
+    e1 = entropy(P_, base=2)
+    e2 = entropy(Q_, base=2)
+    e_sum = entropy((P_ + Q_) / 2.0, base=2)
+    res = e_sum - ((e1 + e2) / 2.0)
+
+    res2 = _jsdiv(P_, Q_)
+
+    if not np.allclose(res, res2, atol=10e-5, rtol=0):
+        warnings.warn('Numerical values of two JSD methods don\'t agree.')
+
+    return res
+
+
+def _jsdiv(P, Q):
+    '''another way of computing JSD'''
+
+    def _kldiv(A, B):
+        a = A.copy()
+        b = B.copy()
+        idx = np.logical_and(a > 0, b > 0)
+        a = a[idx]
+        b = b[idx]
+        return np.sum([v for v in a * np.log2(a / b)])
+
+    P_ = P / np.sum(P)
+    Q_ = Q / np.sum(Q)
+
+    M = 0.5 * (P_ + Q_)
+
+    return 0.5 * (_kldiv(P_, M) + _kldiv(Q_, M))
+
+
+def downsample_pc(points, n):
+    sample_idx = random.sample(list(range(points.shape[0])), n)
+    return points[sample_idx]
+
+
+def normalize_pc(points):
+    scale = np.max(np.abs(points))  
+    points = points / scale
+    return points
+
+
+def collect_pc(cad_folder):
+    pc_path = find_files(os.path.join(cad_folder, 'ptl'), 'final_pcd.ply')
+    if len(pc_path) == 0:
+        return []
+    pc_path = pc_path[-1] # final pcd
+    pc = read_ply(pc_path)
+    if pc.shape[0] > N_POINTS:
+        pc = downsample_pc(pc, N_POINTS)
+    pc = normalize_pc(pc)
+    return pc
+
+def collect_pc2(cad_folder):
+    pc = read_ply(cad_folder)
+    if pc.shape[0] > N_POINTS:
+        pc = downsample_pc(pc, N_POINTS)
+    pc = normalize_pc(pc)
+    return pc
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--fake", type=str)
+    parser.add_argument("--real", type=str)
+    parser.add_argument("--output", type=str)
+    split = 1
+    args = parser.parse_args()
+    if args.output is None:
+        args.output = args.fake + '_cad_results.txt'
+    chamfer_dist = ChamferDistance()
+    cd = []
+    for i in tqdm(range(952)):
+        fake_pcs = []
+        real_pcs = []
+        for j in range(split):
+            fake_index = i * split + j
+            fake_folder = os.path.join(args.fake, f'{fake_index:06d}')
+            if not os.path.exists(fake_folder):
+                continue
+            else:
+                fake_pc = collect_pc(fake_folder)
+                if len(fake_pc) == 0:
+                    continue
+                fake_pcs.append(fake_pc)
+
+        real_folder = os.path.join(args.real, f'{i:06d}')
+        if not os.path.exists(real_folder):
+            continue
+        else:
+            real_pc = collect_pc(real_folder)
+            if len(real_pc) == 0:
+                continue
+            real_pcs.append(real_pc)
+
+        if len(fake_pcs) == 0 or len(real_pcs) == 0:
+            continue
+        sample_pcs = np.stack(fake_pcs, axis=0)
+        ref_pcs = np.stack(real_pcs, axis=0)
+
+        sample_pcs = torch.tensor(sample_pcs, dtype=torch.float32).cuda()
+        ref_pcs = torch.tensor(ref_pcs, dtype=torch.float32).cuda()
+        print(sample_pcs.shape, ref_pcs.shape)
+        dl, dr, idx1, idx2 = chamfer_dist(sample_pcs, ref_pcs)
+        min_val = (dl.mean(dim=1) + dr.mean(dim=1)).view(1, -1).squeeze(0).min().item()
+        cd.append(min_val)
+
+    cd = np.array(cd)
+    mean = np.mean(cd)
+    median = np.median(cd)
+    print('mean:', mean)
+    print('median:', median)
+
+
+if __name__ == '__main__':
+    import time
+    start_time = time.time()
+    main()
+    end_time = time.time()
+    print(end_time - start_time)

CADFusion/src/test/dist_eval.py

@@ -0,0 +1,351 @@
+import torch
+import argparse
+import os
+import numpy as np
+from tqdm import tqdm
+import random
+import warnings
+from glob import glob
+from scipy.stats import entropy
+from sklearn.neighbors import NearestNeighbors
+from plyfile import PlyData
+from pathlib import Path
+from multiprocessing import Pool
+from chamfer_distance import ChamferDistance
+
+random.seed(0)
+N_POINTS = 2000
+NUM_TRHEADS = 16
+
+
+def find_files(folder, extension):
+    return sorted([Path(os.path.join(folder, f)) for f in os.listdir(folder) if f.endswith(extension)])
+
+
+def read_ply(path):
+    with open(path, 'rb') as f:
+        plydata = PlyData.read(f)
+        x = np.array(plydata['vertex']['x'])
+        y = np.array(plydata['vertex']['y'])
+        z = np.array(plydata['vertex']['z'])
+        vertex = np.stack([x, y, z], axis=1)
+    return vertex
+
+
+def distChamfer(a, b):
+    x, y = a, b
+    bs, num_points, points_dim = x.size()
+    xx = torch.bmm(x, x.transpose(2, 1))
+    yy = torch.bmm(y, y.transpose(2, 1))
+    zz = torch.bmm(x, y.transpose(2, 1))
+    diag_ind = torch.arange(0, num_points).to(a).long()
+    rx = xx[:, diag_ind, diag_ind].unsqueeze(1).expand_as(xx)
+    ry = yy[:, diag_ind, diag_ind].unsqueeze(1).expand_as(yy)
+    P = (rx.transpose(2, 1) + ry - 2 * zz)
+    return P.min(1)[0], P.min(2)[0]
+
+
+def _pairwise_CD(sample_pcs, ref_pcs, batch_size):
+    N_sample = sample_pcs.shape[0]
+    N_ref = ref_pcs.shape[0]
+    all_cd = []
+    all_emd = []
+    iterator = range(N_sample)
+    matched_gt = []
+    pbar = tqdm(iterator)
+    chamfer_dist = ChamferDistance()
+
+    for sample_b_start in pbar:
+        sample_batch = sample_pcs[sample_b_start]
+
+        cd_lst = []
+        emd_lst = []
+        for ref_b_start in range(0, N_ref, batch_size):
+            ref_b_end = min(N_ref, ref_b_start + batch_size)
+            ref_batch = ref_pcs[ref_b_start:ref_b_end]
+
+            batch_size_ref = ref_batch.size(0)
+            sample_batch_exp = sample_batch.view(1, -1, 3).expand(batch_size_ref, -1, -1)
+            sample_batch_exp = sample_batch_exp.contiguous()
+            
+            dl, dr, idx1, idx2 = chamfer_dist(sample_batch_exp,ref_batch)
+            cd_lst.append((dl.mean(dim=1) + dr.mean(dim=1)).view(1, -1))
+
+        cd_lst = torch.cat(cd_lst, dim=1)
+        all_cd.append(cd_lst)
+
+        hit = np.argmin(cd_lst.detach().cpu().numpy()[0])
+        matched_gt.append(hit)
+        pbar.set_postfix({"cov": len(np.unique(matched_gt)) * 1.0 / N_ref})
+
+    all_cd = torch.cat(all_cd, dim=0)  # N_sample, N_ref
+
+    return all_cd
+
+
+def compute_cov_mmd(sample_pcs, ref_pcs, batch_size):
+    all_dist = _pairwise_CD(sample_pcs, ref_pcs, batch_size)
+    print(all_dist.shape, flush=True)
+    N_sample, N_ref = all_dist.size(0), all_dist.size(1)
+    min_val_fromsmp, min_idx = torch.min(all_dist, dim=1)
+    min_val, _ = torch.min(all_dist, dim=0)
+    mmd = min_val.mean()
+    cov = float(min_idx.unique().view(-1).size(0)) / float(N_ref)
+    cov = torch.tensor(cov).to(all_dist)
+
+    return {
+        # 'med-CD': torch.diagonal(all_dist).median().item(),
+        'avg-CD': torch.diagonal(all_dist).mean().item(),
+        'COV-CD': cov.item(),
+        'MMD-CD': mmd.item()
+    }
+
+
+def jsd_between_point_cloud_sets(sample_pcs, ref_pcs, in_unit_sphere, resolution=28):
+    '''Computes the JSD between two sets of point-clouds, as introduced in the paper ```Learning Representations And Generative Models For 3D Point Clouds```.
+    Args:
+        sample_pcs: (np.ndarray S1xR2x3) S1 point-clouds, each of R1 points.
+        ref_pcs: (np.ndarray S2xR2x3) S2 point-clouds, each of R2 points.
+        resolution: (int) grid-resolution. Affects granularity of measurements.
+    '''
+    sample_grid_var = entropy_of_occupancy_grid(sample_pcs, resolution, in_unit_sphere)[1]
+    ref_grid_var = entropy_of_occupancy_grid(ref_pcs, resolution, in_unit_sphere)[1]
+    return jensen_shannon_divergence(sample_grid_var, ref_grid_var)
+
+
+def entropy_of_occupancy_grid(pclouds, grid_resolution, in_sphere=False):
+    '''Given a collection of point-clouds, estimate the entropy of the random variables
+    corresponding to occupancy-grid activation patterns.
+    Inputs:
+        pclouds: (numpy array) #point-clouds x points per point-cloud x 3
+        grid_resolution (int) size of occupancy grid that will be used.
+    '''
+    epsilon = 10e-4
+    bound = 1 + epsilon
+    if abs(np.max(pclouds)) > bound or abs(np.min(pclouds)) > bound:
+        print(abs(np.max(pclouds)), abs(np.min(pclouds)))
+        warnings.warn('Point-clouds are not in unit cube.')
+
+    if in_sphere and np.max(np.sqrt(np.sum(pclouds ** 2, axis=2))) > bound:
+        warnings.warn('Point-clouds are not in unit sphere.')
+
+    grid_coordinates, _ = unit_cube_grid_point_cloud(grid_resolution, in_sphere)
+    grid_coordinates = grid_coordinates.reshape(-1, 3)
+    grid_counters = np.zeros(len(grid_coordinates))
+    grid_bernoulli_rvars = np.zeros(len(grid_coordinates))
+    nn = NearestNeighbors(n_neighbors=1).fit(grid_coordinates)
+
+    for pc in pclouds:
+        _, indices = nn.kneighbors(pc)
+        indices = np.squeeze(indices)
+        for i in indices:
+            grid_counters[i] += 1
+        indices = np.unique(indices)
+        for i in indices:
+            grid_bernoulli_rvars[i] += 1
+
+    acc_entropy = 0.0
+    n = float(len(pclouds))
+    for g in grid_bernoulli_rvars:
+        p = 0.0
+        if g > 0:
+            p = float(g) / n
+            acc_entropy += entropy([p, 1.0 - p])
+
+    return acc_entropy / len(grid_counters), grid_counters
+
+
+def unit_cube_grid_point_cloud(resolution, clip_sphere=False):
+    '''Returns the center coordinates of each cell of a 3D grid with resolution^3 cells,
+    that is placed in the unit-cube.
+    If clip_sphere it True it drops the "corner" cells that lie outside the unit-sphere.
+    '''
+    grid = np.ndarray((resolution, resolution, resolution, 3), np.float32)
+    spacing = 1.0 / float(resolution - 1) * 2
+    for i in range(resolution):
+        for j in range(resolution):
+            for k in range(resolution):
+                grid[i, j, k, 0] = i * spacing - 0.5 * 2
+                grid[i, j, k, 1] = j * spacing - 0.5 * 2
+                grid[i, j, k, 2] = k * spacing - 0.5 * 2
+
+    if clip_sphere:
+        grid = grid.reshape(-1, 3)
+        grid = grid[np.linalg.norm(grid, axis=1) <= 0.5]
+
+    return grid, spacing
+
+
+def jensen_shannon_divergence(P, Q):
+    if np.any(P < 0) or np.any(Q < 0):
+        raise ValueError('Negative values.')
+    if len(P) != len(Q):
+        raise ValueError('Non equal size.')
+
+    P_ = P / np.sum(P)  # Ensure probabilities.
+    Q_ = Q / np.sum(Q)
+
+    e1 = entropy(P_, base=2)
+    e2 = entropy(Q_, base=2)
+    e_sum = entropy((P_ + Q_) / 2.0, base=2)
+    res = e_sum - ((e1 + e2) / 2.0)
+
+    res2 = _jsdiv(P_, Q_)
+
+    if not np.allclose(res, res2, atol=10e-5, rtol=0):
+        warnings.warn('Numerical values of two JSD methods don\'t agree.')
+
+    return res
+
+
+def _jsdiv(P, Q):
+    '''another way of computing JSD'''
+
+    def _kldiv(A, B):
+        a = A.copy()
+        b = B.copy()
+        idx = np.logical_and(a > 0, b > 0)
+        a = a[idx]
+        b = b[idx]
+        return np.sum([v for v in a * np.log2(a / b)])
+
+    P_ = P / np.sum(P)
+    Q_ = Q / np.sum(Q)
+
+    M = 0.5 * (P_ + Q_)
+
+    return 0.5 * (_kldiv(P_, M) + _kldiv(Q_, M))
+
+
+def downsample_pc(points, n):
+    sample_idx = random.sample(list(range(points.shape[0])), n)
+    return points[sample_idx]
+
+
+def normalize_pc(points):
+    scale = np.max(np.abs(points))  
+    points = points / scale
+    return points
+
+
+def collect_pc(cad_folder):
+    pc_path = find_files(os.path.join(cad_folder, 'ptl'), 'final_pcd.ply')
+    if len(pc_path) == 0:
+        return []
+    pc_path = pc_path[-1] # final pcd
+    pc = read_ply(pc_path)
+    if pc.shape[0] > N_POINTS:
+        pc = downsample_pc(pc, N_POINTS)
+    pc = normalize_pc(pc)
+    return pc
+
+def collect_pc2(cad_folder):
+    pc = read_ply(cad_folder)
+    if pc.shape[0] > N_POINTS:
+        pc = downsample_pc(pc, N_POINTS)
+    pc = normalize_pc(pc)
+    return pc
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--fake", type=str)
+    parser.add_argument("--real", type=str)
+    parser.add_argument("--output", type=str)
+    parser.add_argument("--n_test", type=int, default=200)
+    parser.add_argument("--multi", type=int, default=1)    
+    parser.add_argument("--times", type=int, default=10)
+    parser.add_argument("--batch_size", type=int, default=64)
+    args = parser.parse_args()
+
+    print("n_test: {}, multiplier: {}, repeat times: {}".format(args.n_test, args.multi, args.times))
+    if args.output is None:
+        args.output = args.fake + '_cad_results.txt'
+
+    # Load fake pcd
+
+    fake_folders = sorted(glob(args.fake+'/*/'))
+    real_folders = sorted(glob(args.real+'/*/'))
+
+    fake_overlapped = []
+    real_overlapped = []
+    for i in range(800):
+        if f'{args.fake}/{i:06d}/' in fake_folders and f'{args.real}/{i:06d}/' in real_folders:
+            if len(glob(f'{args.fake}/{i:06d}/ptl/*')) > 0 and len(glob(f'{args.real}/{i:06d}/ptl/*')) > 0:
+                fake_overlapped.append(f'{args.fake}/{i:06d}/')
+                real_overlapped.append(f'{args.real}/{i:06d}/')
+    print(len(fake_overlapped), len(real_overlapped))
+
+    fake_folders = fake_overlapped
+    real_folders = real_overlapped
+    
+    sample_pcs = []
+    load_iter = Pool(NUM_TRHEADS).imap(collect_pc, fake_folders)
+    for pc in tqdm(load_iter, total=len(fake_folders)):
+        if len(pc) > 0:
+            sample_pcs.append(pc)
+    sample_pcs = np.stack(sample_pcs, axis=0)
+    print("fake point clouds: {}".format(sample_pcs.shape))
+
+    # Load reference pcd 
+    ref_pcs = [] 
+    load_iter = Pool(NUM_TRHEADS).imap(collect_pc, real_folders)
+    for pc in tqdm(load_iter, total=len(real_folders)):
+        if len(pc) > 0:
+            ref_pcs.append(pc)
+    ref_pcs = np.stack(ref_pcs, axis=0)
+    print("real point clouds: {}".format(ref_pcs.shape))
+    
+    # # Testing
+    fp = open(args.output, "w")
+
+    rand_sample_pcs = sample_pcs
+    rand_ref_pcs = ref_pcs
+
+    jsd = jsd_between_point_cloud_sets(rand_sample_pcs, rand_ref_pcs, in_unit_sphere=False)
+    with torch.no_grad():
+        rand_sample_pcs = torch.tensor(rand_sample_pcs).cuda()
+        rand_ref_pcs = torch.tensor(rand_ref_pcs).cuda()
+        result = compute_cov_mmd(rand_sample_pcs, rand_ref_pcs, batch_size=args.batch_size)
+    result.update({"JSD": jsd})
+
+    print(result)
+    print(result, file=fp)
+    fp.close()
+
+    # Testing
+    # fp = open(args.output, "w")
+    # result_list = []
+    # for i in range(args.times):
+    #     print("iteration {}...".format(i))
+    #     select_idx = random.sample(list(range(len(sample_pcs))), int(args.multi * args.n_test))
+    #     rand_sample_pcs = sample_pcs[select_idx]
+
+    #     select_idx = random.sample(list(range(len(ref_pcs))), args.n_test)
+    #     rand_ref_pcs = ref_pcs[select_idx]
+
+    #     jsd = jsd_between_point_cloud_sets(rand_sample_pcs, rand_ref_pcs, in_unit_sphere=False)
+    #     with torch.no_grad():
+    #         rand_sample_pcs = torch.tensor(rand_sample_pcs).cuda()
+    #         rand_ref_pcs = torch.tensor(rand_ref_pcs).cuda()
+    #         result = compute_cov_mmd(rand_sample_pcs, rand_ref_pcs, batch_size=args.batch_size)
+    #     result.update({"JSD": jsd})
+
+    #     print(result)
+    #     print(result, file=fp)
+    #     result_list.append(result)
+    # avg_result = {}
+    # for k in result_list[0].keys():
+    #     avg_result.update({"avg-" + k: np.mean([x[k] for x in result_list])})
+    # print("average result:")
+    # print(avg_result)
+    # print(avg_result, file=fp)
+    # fp.close()
+
+
+if __name__ == '__main__':
+    import time
+    start_time = time.time()
+    main()
+    end_time = time.time()
+    print(end_time - start_time)

CADFusion/src/test/f1_eval.py

@@ -0,0 +1,74 @@
+import json
+import re
+import argparse
+
+"""
+We did not implement the Hungarian matching algorithm from text2cad, but provided a vanilla matching for f1. It is because
+ 1. We argue that CAD scenarios are too complicated to be evaluated with a simple matching algorithm, especially when performed on the primitive level. Moreover, matching every primitive exactly is against the goal of our framework which attempt to encourage CAD models generate visually correct objects instead of accurate primitives compared to the ground truth.
+ 2. In our exploration, discrepancies on the number of primitives between model generation and the ground truth usually indicates the entire failure of the sketch so that using any of the algorithm does not affect the final evaluation result anyway.
+ 3. Our evaluation is a lower bound of the performance of the model on the matching algorithm, therefore it does not affect the overall integrety of our framework.
+
+We encourage users to implement their own matching algorithm if they want to evaluate the model with a more strict metric.
+"""
+
+parser = argparse.ArgumentParser(description='Evaluate F1 scores for generated sketches.')
+parser.add_argument('--test-path', type=str, default='data/sl_data/test.jsonl', help='Path to the JSONL file containing test data')    
+parser.add_argument('--file_path', type=str, required=True, help='Path to the JSONL file containing generated sketches.')
+args = parser.parse_args()
+file_path = args.file_path
+data_path = args.test_path
+with open(data_path, 'r') as f:
+    data = json.load(f)
+    
+def find_f1(ground_truth, pred, token):
+    num_tok_gt = len(re.findall(token, ground_truth))
+    num_tok_pred = len(re.findall(token, pred))
+    # print(num_tok_gt, num_tok_pred)
+    min_tok = min(num_tok_gt, num_tok_pred)
+    if min_tok <= 0:
+        return -1
+    tok_recall = min_tok / num_tok_gt
+    tok_precision = min_tok / num_tok_pred
+    tok_f1 = 2 * tok_recall * tok_precision / (tok_recall + tok_precision)
+    return tok_f1
+
+with open(file_path, 'r') as f:
+    gen = json.load(f)
+line = []
+arc = []
+circle = []
+ext = []
+for i in range(1000):
+    ground_truth = data[i]['output']
+    pred = gen[i]['output']
+    ext_f1 = find_f1(ground_truth, pred, r'<extrude_end>')
+    if ext_f1 > 0:
+        ext.append(ext_f1)
+
+    skext_gt = ground_truth.split('<extrude_end>')[:-1]
+    skext_pred = pred.split('<extrude_end>')[:-1]
+    min_len_skext = min(len(skext_gt), len(skext_pred))
+    if min_len_skext == 0:
+        continue
+    line_f1 = 0
+    arc_f1 = 0
+    circle_f1 = 0 
+    for gt, pr in zip(skext_gt, skext_pred):
+        line_f1 += find_f1(gt, pr, r'line.*?<curve_end>')
+        arc_f1 += find_f1(gt, pr, r'arc.*?<curve_end>')
+        circle_f1 += find_f1(gt, pr, r'circle.*?<curve_end>')
+        
+    line_f1 = line_f1 / min_len_skext
+    arc_f1 = arc_f1 / min_len_skext
+    circle_f1 = circle_f1 / min_len_skext
+    if line_f1 > 0:
+        line.append(line_f1)
+    if arc_f1 > 0:
+        arc.append(arc_f1)
+    if circle_f1 > 0:
+        circle.append(circle_f1)
+line_avg = sum(line) / len(line)
+arc_avg = sum(arc) / len(arc)
+circle_avg = sum(circle) / len(circle)
+avgf1 = (line_avg + arc_avg + circle_avg) / 3
+print(file_path, line_avg, arc_avg, circle_avg, avgf1, sum(ext) / len(ext))

CADFusion/src/test/generate.ipynb

@@ -0,0 +1,291 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2d243f81",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import argparse\n",
+    "import random\n",
+    "import os\n",
+    "import subprocess\n",
+    "import shutil\n",
+    "\n",
+    "from PIL import Image\n",
+    "from huggingface_hub import login\n",
+    "from utils import MAX_LENGTH, prepare_model_and_tokenizer\n",
+    "from visual_utils.parser import CADparser, write_obj_sample\n",
+    "from IPython.display import clear_output"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b98812ed",
+   "metadata": {},
+   "source": [
+    "### Initializing model and arguments"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 49,
+   "id": "df625563",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "parser = argparse.ArgumentParser()\n",
+    "# parser.add_argument(\"--model-name\", type=str, default=\"llama3\")\n",
+    "parser.add_argument(\"--device-map\", type=str, default='auto')\n",
+    "parser.add_argument(\"--lora-rank\", type=int, default=32)\n",
+    "parser.add_argument(\"--lora-alpha\", type=int, default=32)\n",
+    "parser.add_argument(\"--lora-dropout\", type=float, default=0.05)\n",
+    "parser.add_argument(\"--pretrained-path\", type=str, required=True)\n",
+    "parser.add_argument(\"--top-p\", type=float, default=0.9)\n",
+    "parser.add_argument(\"--temperature\", type=float, default=0.9)\n",
+    "\n",
+    "arguments = ['--pretrained-path', '/home/v-wangruiyu/repos/CADFusion/exp/model_ckpt/CADFusion_v1_1', '--temperature', '0.3']\n",
+    "args = parser.parse_args(arguments)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5624f320",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "login() # put your own hf token to access llama\n",
+    "random.seed(0)\n",
+    "model, tokenizer = prepare_model_and_tokenizer(args)\n",
+    "model.eval()\n",
+    "clear_output()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "86b9cb09",
+   "metadata": {},
+   "source": [
+    "### Custom prompting"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 180,
+   "id": "db06d560",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "description = input(\"Please input a description of a 3D shape: \")\n",
+    "# description = 'The 3D shape is a cylinder.'\n",
+    "\n",
+    "prompt = 'Below is a description of a 3D shape:\\n'\n",
+    "prompt += description\n",
+    "prompt += '\\nGenerate a Computer-Aided Design(CAD) command sequence of the 3D shape:\\n'"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bb16f861",
+   "metadata": {},
+   "source": [
+    "### Inference and rendering"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "59c5f38e",
+   "metadata": {},
+   "source": [
+    "#### Model Inference"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 181,
+   "id": "ab5ff2e8",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Setting `pad_token_id` to `eos_token_id`:128001 for open-end generation.\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "'circle,31,53,31,9,53,31,9,31 <curve_end> <loop_end> circle,31,51,31,11,51,31,11,31 <curve_end> <loop_end> <face_end> circle,31,51,31,11,51,31,11,31 <curve_end> <loop_end> <face_end> <sketch_end> add,0,62,31,31,31,1,0,0,0,0,1,0,-1,0,7,31,31 <extrude_end>'"
+      ]
+     },
+     "execution_count": 181,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "batch = tokenizer(\n",
+    "    prompt,\n",
+    "    return_tensors=\"pt\",\n",
+    ")\n",
+    "batch = {k: v.cuda() for k, v in batch.items()}\n",
+    "\n",
+    "generate_ids = model.generate(\n",
+    "    **batch,\n",
+    "    do_sample=True,\n",
+    "    max_new_tokens=MAX_LENGTH,\n",
+    "    temperature=args.temperature,\n",
+    "    top_p=args.top_p,\n",
+    "    repetition_penalty=1.3,\n",
+    ")\n",
+    "\n",
+    "gen_strs = tokenizer.batch_decode(\n",
+    "    generate_ids,\n",
+    "    skip_special_tokens=True,\n",
+    "    clean_up_tokenization_spaces=False,\n",
+    ")\n",
+    "gen_strs = gen_strs[0][len(prompt):]\n",
+    "gen_strs"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f56d6fcf",
+   "metadata": {},
+   "source": [
+    "#### Render .obj file"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 182,
+   "id": "95498ccb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "out_path = 'visual_cache/gen_obj'\n",
+    "# remove the existing output directory if it exists\n",
+    "if os.path.exists(out_path):\n",
+    "    shutil.rmtree(out_path)\n",
+    "# create the output directory\n",
+    "os.makedirs(out_path, exist_ok=True)\n",
+    "\n",
+    "cad_parser = CADparser(bit=6)\n",
+    "parsed_data = cad_parser.perform(gen_strs)\n",
+    "write_obj_sample(out_path, parsed_data)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "79b5dfaf",
+   "metadata": {},
+   "source": [
+    "#### Render .step, .stl, .ply files\n",
+    "N.B. if the Statistics on Transfer logs do not show up, the model may not have produced renderable outputs. Re-run the inference or change your prompt to see if it gets better results. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8a49694f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "out_path = os.path.abspath(out_path)\n",
+    "py_path = os.path.abspath('../rendering_utils/parser_visual.py')\n",
+    "subprocess.run(['python3', py_path, '--data_folder', out_path, '--single-file'])\n",
+    "py_path = os.path.abspath('../rendering_utils/ptl_sampler.py')\n",
+    "subprocess.run(['python3', py_path, '--in_dir', out_path, '--out_dir', 'ptl', '--single-file'])\n",
+    "# clear_output()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0e0f1fd1",
+   "metadata": {},
+   "source": [
+    "#### Image rendering"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "586f3a91",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "visual_obj_path = 'visual_cache'\n",
+    "output_figure_path = 'visual_cache/figures'\n",
+    "if os.path.exists(output_figure_path):\n",
+    "    shutil.rmtree(output_figure_path)\n",
+    "py_path = os.path.abspath('../rendering_utils/img_renderer.py')\n",
+    "os.makedirs(output_figure_path, exist_ok=True)\n",
+    "try:\n",
+    "    xvfb_process = subprocess.Popen(\n",
+    "        [\"Xvfb\", \":99\", \"-screen\", \"0\", \"640x480x24\"],\n",
+    "        stdout=subprocess.DEVNULL,\n",
+    "        stderr=subprocess.DEVNULL\n",
+    "    )\n",
+    "    print(\"Xvfb started in the background.\")\n",
+    "except FileNotFoundError:\n",
+    "    print(\"Error: Xvfb not found. Please ensure it is installed and in your system's PATH.\")\n",
+    "\n",
+    "os.environ['DISPLAY'] = ':99'\n",
+    "try:\n",
+    "    subprocess.run(\n",
+    "        ['python3', py_path, '--input_dir', visual_obj_path, '--output_dir', output_figure_path]\n",
+    "    )\n",
+    "    print(\"Rendering script completed successfully.\")\n",
+    "finally:\n",
+    "    if xvfb_process.poll() is None: # Check if Xvfb is still running\n",
+    "        xvfb_process.terminate()\n",
+    "        print(\"Xvfb terminated.\")\n",
+    "    else:\n",
+    "        print(\"Xvfb already exited.\")\n",
+    "        \n",
+    "del os.environ['DISPLAY']\n",
+    "clear_output()\n",
+    "\n",
+    "input_image_path = os.path.join(output_figure_path, 'gen_ob.png')\n",
+    "if os.path.exists(input_image_path):\n",
+    "    img = Image.open(input_image_path)\n",
+    "    img.show()\n",
+    "else:\n",
+    "    print(f\"{input_image_path} does not exist.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c78fed0f",
+   "metadata": {},
+   "source": [
+    "#### Files retrieval\n",
+    "By default, the produced step, stl, obj and ply files are stored under the visual_cache folder. You can save them to your custom places for further use. Do not put them in the cache folder as they will be deleted after the next run."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "cdfs",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.23"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

CADFusion/src/test/inference.py

@@ -0,0 +1,106 @@
+import argparse
+import json
+import random
+
+from huggingface_hub import login
+from tqdm import tqdm
+from utils import MAX_LENGTH, prepare_model_and_tokenizer
+
+login()
+
+random.seed(0)
+
+def conditional_sample(args):
+    model, tokenizer = prepare_model_and_tokenizer(args)
+    
+    model.eval()
+    with open(args.in_path, 'r', encoding='utf-8') as file:  
+        data = json.load(file)
+        
+    print(data[0])
+    data = [item for item in data if item['description'] != 'null']
+
+    global_count=0
+    responses = []  
+    if args.full:
+        data=data
+    else:
+        random.shuffle(data)
+        data = data[:args.sample_len]
+    
+    for item in tqdm(data):
+        prompts = []
+        for _ in range(args.num_samples):
+            prompt = 'Below is a description of a 3D shape:\n'
+            prompt += item['description']  
+            prompt += '\nGenerate a Computer-Aided Design(CAD) command sequence of the 3D shape:\n'     
+
+            prompts.append(prompt)
+
+        outputs = []
+
+        while len(outputs) < args.num_samples:
+            batch_prompts = prompts[len(outputs) : len(outputs) + args.batch_size]
+
+            batch = tokenizer(
+                list(batch_prompts),
+                return_tensors="pt",
+            )
+            batch = {k: v.cuda() for k, v in batch.items()}
+
+            generate_ids = model.generate(
+                **batch,
+                do_sample=True,
+                max_new_tokens=MAX_LENGTH,
+                temperature=args.temperature,
+                top_p=args.top_p,
+                repetition_penalty=1.3,
+            )
+
+            gen_strs = tokenizer.batch_decode(
+                generate_ids,
+                skip_special_tokens=True,
+                clean_up_tokenization_spaces=False,
+            )
+
+            outputs.extend(gen_strs)
+            print(f"Generated {len(outputs)}/{args.num_samples}samples.")
+
+            for prompt, output in zip(prompts, outputs):
+                result = {
+                    'index': global_count,
+                    # 'pic_name': item['pic_name'],
+                    'ground_truth': item['command_sequence'],
+                    'description': item['description'],
+                    'prompt': prompt,
+                    'output': output[len(prompt):]
+                }
+                if 'original_seq' in item.keys():
+                    result['original_seq'] = item['original_seq']
+                responses.append(result)
+                global_count += 1
+
+        with open(args.out_path, "w+") as f:
+            json.dump(responses, f, indent=4)
+
+                
+                    
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-name", type=str, default="llama3")
+    parser.add_argument("--lora-rank", type=int, default=32)
+    parser.add_argument("--lora-alpha", type=int, default=32)
+    parser.add_argument("--lora-dropout", type=float, default=0.05)
+    parser.add_argument("--sample-len", type=int, default=100)
+    parser.add_argument("--pretrained-path", type=str, required=True)
+    parser.add_argument("--num-samples", type=int, default=500)
+    parser.add_argument("--batch-size", type=int, default=32)
+    parser.add_argument("--in-path", type=str, default="test_description.json")
+    parser.add_argument("--out-path", type=str, default="cad_samples.jsonl")
+    parser.add_argument("--temperature", type=float, default=0.9)
+    parser.add_argument("--device-map", type=str, default='auto')
+    parser.add_argument("--top-p", type=float, default=0.9)
+    parser.add_argument("--full", action="store_true", default=False)
+    args = parser.parse_args()
+
+    conditional_sample(args)

CADFusion/src/test/utils.py

@@ -0,0 +1,86 @@
+import torch
+import transformers
+from peft import LoraConfig, PeftModel, get_peft_model
+
+IGNORE_INDEX = -100
+MAX_LENGTH = 512
+DEFAULT_PAD_TOKEN = "[PAD]"
+DEFAULT_EOS_TOKEN = "</s>"
+DEFAULT_BOS_TOKEN = "<s>"
+DEFAULT_UNK_TOKEN = "<unk>"
+
+def smart_tokenizer_and_embedding_resize(
+    special_tokens_dict,
+    llama_tokenizer,
+    model,
+):
+    """Resize tokenizer and embedding.
+
+    Note: This is the unoptimized version that may make your embedding size not be divisible by 64.
+    """
+    num_new_tokens = llama_tokenizer.add_special_tokens(special_tokens_dict)
+    model.resize_token_embeddings(len(llama_tokenizer))
+
+    if num_new_tokens > 0:
+        input_embeddings = model.get_input_embeddings().weight.data
+        output_embeddings = model.get_output_embeddings().weight.data
+
+        input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(
+            dim=0, keepdim=True
+        )
+        output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(
+            dim=0, keepdim=True
+        )
+
+        input_embeddings[-num_new_tokens:] = input_embeddings_avg
+        output_embeddings[-num_new_tokens:] = output_embeddings_avg
+
+def prepare_model_and_tokenizer(args):
+    model_id = "meta-llama/Meta-Llama-3-8B"
+    print(f"Model size: {model_id}")
+    if hasattr(args, 'device_map'):
+        device_map = args.device_map
+    else:
+        device_map = 'auto'
+    pipeline = transformers.pipeline("text2text-generation",
+                                        model=model_id, model_kwargs={"torch_dtype": torch.float32}, device_map=device_map)
+    tokenizer = pipeline.tokenizer
+    base_model = pipeline.model
+
+    special_tokens_dict = dict()
+    if tokenizer.pad_token is None:
+        special_tokens_dict["pad_token"] = DEFAULT_PAD_TOKEN
+    if tokenizer.eos_token is None:
+        special_tokens_dict["eos_token"] = DEFAULT_EOS_TOKEN
+    if tokenizer.bos_token is None:
+        special_tokens_dict["bos_token"] = DEFAULT_BOS_TOKEN
+    if tokenizer.unk_token is None:
+        special_tokens_dict["unk_token"] = DEFAULT_UNK_TOKEN
+
+    smart_tokenizer_and_embedding_resize(
+        special_tokens_dict=special_tokens_dict,
+        llama_tokenizer=tokenizer,
+        model=base_model,
+    )
+    
+    peft_config = LoraConfig(
+        r=args.lora_rank,
+        lora_alpha=args.lora_alpha,
+        lora_dropout=args.lora_dropout,
+        bias="none",
+        task_type="CAUSAL_LM",
+    )
+    
+    tokenizer.padding_side = 'left'
+    peftmodel = get_peft_model(base_model, peft_config)
+    if args.pretrained_path:
+        # load a previous checkpoint if the path is given
+        model = PeftModel.from_pretrained(base_model, args.pretrained_path, device_map=device_map)
+        peft_state_dict = {f"{k}": v for k, v in model.state_dict().items()}
+        peftmodel.load_state_dict(peft_state_dict)
+        
+        for name, param in peftmodel.named_parameters():
+            if "lora" in name:  # Check if "lora" is in the parameter's name
+                param.requires_grad = True  
+    peftmodel.print_trainable_parameters()
+    return peftmodel, tokenizer

CADFusion/src/test/visual_utils/parser.py

@@ -0,0 +1,478 @@
+import numpy as np
+from collections import OrderedDict
+import re
+from pathlib import Path
+import argparse
+import os
+import json
+import math
+
+# hyperparameters from SkexGen project
+SKETCH_R = 1
+RADIUS_R = 1
+EXTRUDE_R = 1.0
+SCALE_R = 1.4
+OFFSET_R = 0.9
+PIX_PAD = 4
+CMD_PAD = 3
+COORD_PAD = 4
+EXT_PAD = 1
+EXTRA_PAD = 1
+R_PAD = 2
+
+
+class CADparser:
+    """Parse CAD sequence to CAD object."""
+
+    def __init__(self, bit):
+        self.vertex_dict = OrderedDict()
+        self.bit = bit
+
+    def perform(self, cad_seq):
+        # divide into sketch and extrude
+        sketches, extrudes = self.get_SE(cad_seq)
+        if sketches is None or extrudes is None:
+            return None
+        # sequentially parse each pair of SE into obj
+        se_datas = []
+        for sketch, extrude in zip(sketches, extrudes):
+            extrude_param, scale, offset = self.parse_extrude(extrude)
+            if extrude_param is None or scale is None or offset is None:
+                return None
+            vertex_str, se_str = self.parse_sketch(sketch, scale, offset)
+            if vertex_str is None or se_str is None:
+                return None
+            se_datas.append(
+                {"vertex": vertex_str, "curve": se_str, "extrude": extrude_param}
+            )
+            self.vertex_dict.clear()
+
+        return se_datas
+
+    def parse_sketch(self, sketch, scale, offset):
+        faces = self.get_faces(sketch)
+        if len(faces) == 0:
+            return None, None
+        se_str = ""
+        for face_idx, face in enumerate(faces):  # each face
+            face_str = "face\n"
+            loops = self.get_loops(face)
+            if len(loops) == 0:
+                return None, None
+            for loop_idx, loop in enumerate(loops):  # each loop
+                curves = self.get_curves(loop)
+                if len(curves) == 0:
+                    return None, None
+                next_curves = curves[1:]
+                next_curves += curves[:1]
+                cur_str = []
+                for curve, next_curve in zip(curves, next_curves):  # each curve
+                    if not self.obj_curve(curve, next_curve, cur_str, scale, offset):
+                        return None, None
+                loop_str = ""
+                for c in cur_str:
+                    loop_str += f"{c}\n"
+                if loop_idx == 0:
+                    face_str += f"out\n{loop_str}\n"
+                else:
+                    face_str += f"in\n{loop_str}\n"
+            se_str += face_str
+        vertex_str = self.convert_vertices()
+        return vertex_str, se_str
+
+    def parse_extrude(self, extrude):
+        ext = extrude.split(",")
+        if len(ext) != 18:
+            return None, None, None
+
+        # operation str to int
+        ext_op = {"add": 1, "cut": 2, "intersect": 3}.get(ext[0], None)
+        if ext_op is None:
+            return None, None, None
+        # dequantize ext_v, ext_T, scale and offset
+        ext_v, ext_T, scale, offset = self.dequantize_extrude_params(ext)
+        # get ext_R
+        ext_R = np.array(ext[6:15], dtype=int)
+
+        extrude_param = {"value": ext_v, "T": ext_T, "R": ext_R, "op": ext_op}
+        return extrude_param, scale, offset
+
+    def obj_curve(self, curve, next_curve, cur_str, scale, offset):
+        cur = curve.split(",")
+        next_cur = next_curve.split(",")
+        if cur[0] == "circle":
+            if len(cur) != 9:
+                return False
+            p1, p2, p3, p4 = self.dequantize_circle_points(
+                cur, next_cur, scale, offset)
+            center = np.asarray([0.5 * (p1[0] + p2[0]), 0.5 * (p3[1] + p4[1])])
+            radius = (np.linalg.norm(p1 - p2) + np.linalg.norm(p3 - p4)) / 4.0
+
+            center = center * scale + offset
+            radius = radius * scale
+
+            center_idx = self.save_vertex(center[0], center[1], "p")
+            radius_idx = self.save_vertex(radius, 0.0, "r")
+            cur_str.append(f"c {center_idx} {radius_idx}")
+        elif cur[0] == "arc":
+            if len(cur) != 5:
+                return False
+            if (
+                cur[1:3] == cur[3:5]
+                or cur[1:3] == next_cur[1:3]
+                or cur[3:5] == next_cur[3:5]
+            ):  # invalid arc
+                return False
+            start_v, mid_v, end_v = self.dequantize_arc_points(
+                cur, next_cur, scale, offset
+            )
+            try:
+                center, _, _, _ = find_arc_geometry(start_v, mid_v, end_v)
+            except Exception:
+                return False
+            start_v = start_v * scale + offset
+            mid_v = mid_v * scale + offset
+            end_v = end_v * scale + offset
+            center = center * scale + offset
+
+            center_idx = self.save_vertex(center[0], center[1], "p")
+            start_idx = self.save_vertex(start_v[0], start_v[1], "p")
+            mid_idx = self.save_vertex(mid_v[0], mid_v[1], "p")
+            end_idx = self.save_vertex(end_v[0], end_v[1], "p")
+            cur_str.append(f"a {start_idx} {mid_idx} {center_idx} {end_idx}")
+        elif cur[0] == "line":
+            if len(cur) != 3:
+                return False
+            if cur[1:3] == next_cur[1:3]:
+                return False
+            start_v, end_v = self.dequantize_line_points(
+                cur, next_cur, scale, offset)
+            start_v = start_v * scale + offset
+            end_v = end_v * scale + offset
+
+            start_idx = self.save_vertex(start_v[0], start_v[1], "p")
+            end_idx = self.save_vertex(end_v[0], end_v[1], "p")
+            cur_str.append(f"l {start_idx} {end_idx}")
+        else:
+            return False
+        return True
+
+    def get_SE(self, cad_seq):
+        # sketches: 1) between sequence start and sketch_end,
+        sketches_from_start = re.findall(r"^(.+?)(?=<sketch_end>)", cad_seq)
+        # sketches: 2) between extrude_end and sketch_end
+        sketches_after_extrude = re.findall(
+            r"(?<=<extrude_end>)(.+?)(?=<sketch_end>)", cad_seq
+        )
+        sketches = [x.strip() for x in sketches_from_start] + [
+            x.strip() for x in sketches_after_extrude
+        ]
+        # extrudes: between sketch_end and extrude_end
+        extrudes = [
+            x.strip() for x in re.findall(r"<sketch_end>(.+?)<extrude_end>", cad_seq)
+        ]
+        if len(sketches) != len(extrudes):
+            return None, None
+        return sketches, extrudes
+
+    def get_faces(self, sketch):
+        faces = sketch.split("<face_end>")
+        return [x.strip() for x in faces if x.strip() != ""]
+
+    def get_loops(self, face):
+        loops = face.split("<loop_end>")
+        return [x.strip() for x in loops if x.strip() != ""]
+
+    def get_curves(self, loop):
+        curves = loop.split("<curve_end>")
+        return [x.strip() for x in curves if x.strip() != ""]
+
+    def dequantize_circle_points(self, curve, next_curve, scale, offset):
+        p1 = dequantize_verts(
+            np.array(curve[1:3], dtype=int),
+            n_bits=self.bit,
+            min_range=-SKETCH_R,
+            max_range=SKETCH_R,
+            add_noise=False,
+        )
+        p2 = dequantize_verts(
+            np.array(curve[3:5], dtype=int),
+            n_bits=self.bit,
+            min_range=-SKETCH_R,
+            max_range=SKETCH_R,
+            add_noise=False,
+        )
+        p3 = dequantize_verts(
+            np.array(curve[5:7], dtype=int),
+            n_bits=self.bit,
+            min_range=-SKETCH_R,
+            max_range=SKETCH_R,
+            add_noise=False,
+        )
+        p4 = dequantize_verts(
+            np.array(curve[7:9], dtype=int),
+            n_bits=self.bit,
+            min_range=-SKETCH_R,
+            max_range=SKETCH_R,
+            add_noise=False,
+        )
+        return p1, p2, p3, p4
+
+    def dequantize_arc_points(self, curve, next_curve, scale, offset):
+        start_v = dequantize_verts(
+            np.array(curve[1:3], dtype=int),
+            n_bits=self.bit,
+            min_range=-SKETCH_R,
+            max_range=SKETCH_R,
+            add_noise=False,
+        )
+        mid_v = dequantize_verts(
+            np.array(curve[3:5], dtype=int),
+            n_bits=self.bit,
+            min_range=-SKETCH_R,
+            max_range=SKETCH_R,
+            add_noise=False,
+        )
+        end_v = dequantize_verts(
+            np.array(next_curve[1:3], dtype=int),
+            n_bits=self.bit,
+            min_range=-SKETCH_R,
+            max_range=SKETCH_R,
+            add_noise=False,
+        )
+        return start_v, mid_v, end_v
+
+    def dequantize_line_points(self, curve, next_curve, scale, offset):
+        start_v = dequantize_verts(
+            np.array(curve[1:3], dtype=int),
+            n_bits=self.bit,
+            min_range=-SKETCH_R,
+            max_range=SKETCH_R,
+            add_noise=False,
+        )
+        end_v = dequantize_verts(
+            np.array(next_curve[1:3], dtype=int),
+            n_bits=self.bit,
+            min_range=-SKETCH_R,
+            max_range=SKETCH_R,
+            add_noise=False,
+        )
+        return start_v, end_v
+
+    def dequantize_extrude_params(self, extrude):
+        ext_v = dequantize_verts(
+            np.array(extrude[1:3], dtype=int),
+            n_bits=self.bit,
+            min_range=-EXTRUDE_R,
+            max_range=EXTRUDE_R,
+            add_noise=False,
+        )
+        ext_T = dequantize_verts(
+            np.array(extrude[3:6], dtype=int),
+            n_bits=self.bit,
+            min_range=-EXTRUDE_R,
+            max_range=EXTRUDE_R,
+            add_noise=False,
+        )
+        scale = dequantize_verts(
+            np.array(extrude[15], dtype=int),
+            n_bits=self.bit,
+            min_range=0.0,
+            max_range=SCALE_R,
+            add_noise=False,
+        )
+        offset = dequantize_verts(
+            np.array(extrude[16:18], dtype=int),
+            n_bits=self.bit,
+            min_range=-OFFSET_R,
+            max_range=OFFSET_R,
+            add_noise=False,
+        )
+        return ext_v, ext_T, scale, offset
+
+    def save_vertex(self, h_x, h_y, text):
+        unique_key = f"{text}:x{h_x}y{h_y}"
+        index = 0
+        for key in self.vertex_dict.keys():
+            # Vertex location already exist in dict
+            if unique_key == key:
+                return index
+            index += 1
+        # Vertex location does not exist in dict
+        self.vertex_dict[unique_key] = [h_x, h_y]
+        return index
+
+    def convert_vertices(self):
+        """Convert all the vertices to .obj format"""
+        vertex_strings = ""
+        for pt in self.vertex_dict.values():
+            # e.g. v 0.123 0.234 0.345 1.0
+            vertex_string = f"v {pt[0]} {pt[1]}\n"
+            vertex_strings += vertex_string
+        return vertex_strings
+
+
+def find_arc_geometry(a, b, c):
+    A = b[0] - a[0]
+    B = b[1] - a[1]
+    C = c[0] - a[0]
+    D = c[1] - a[1]
+
+    E = A*(a[0] + b[0]) + B*(a[1] + b[1])
+    F = C*(a[0] + c[0]) + D*(a[1] + c[1])
+
+    G = 2.0*(A*(c[1] - b[1])-B*(c[0] - b[0]))
+
+    if G == 0:
+        raise Exception("zero G")
+
+    p_0 = (D*E - B*F) / G
+    p_1 = (A*F - C*E) / G
+
+    center = np.array([p_0, p_1])
+    radius = np.linalg.norm(center - a)
+
+    angles = []
+    for xx in [a, b, c]:
+        angle = angle_from_vector_to_x(xx - center)
+        angles.append(angle)
+
+    ab = b-a
+    ac = c-a
+    cp = np.cross(ab, ac)
+    if cp >= 0:
+        start_angle_rads = angles[0]
+        end_angle_rads = angles[2]
+    else:
+        start_angle_rads = angles[2]
+        end_angle_rads = angles[0]
+
+    return center, radius, start_angle_rads, end_angle_rads
+
+
+def angle_from_vector_to_x(vec):
+    assert vec.size == 2
+    # We need to find a unit vector
+    angle = 0.0
+
+    l = np.linalg.norm(vec)
+    uvec = vec/l
+
+    # 2 | 1
+    # -------
+    # 3 | 4
+    if uvec[0] >= 0:
+        if uvec[1] >= 0:
+            # Qadrant 1
+            angle = math.asin(uvec[1])
+        else:
+            # Qadrant 4
+            angle = 2.0*math.pi - math.asin(-uvec[1])
+    else:
+        if vec[1] >= 0:
+            # Qadrant 2
+            angle = math.pi - math.asin(uvec[1])
+        else:
+            # Qadrant 3
+            angle = math.pi + math.asin(-uvec[1])
+    return angle
+
+
+def dequantize_verts(verts, n_bits=8, min_range=-0.5, max_range=0.5, add_noise=False):
+    """Convert quantized vertices to floats."""
+    range_quantize = 2**n_bits - 1
+    verts = verts.astype("float32")
+    verts = verts * (max_range - min_range) / range_quantize + min_range
+    return verts
+
+
+def write_obj_sample(save_folder, data):
+    for idx, write_data in enumerate(data):
+        obj_name = Path(save_folder).stem + "_" + \
+            str(idx).zfill(3) + "_param.obj"
+        obj_file = Path(save_folder) / obj_name
+        extrude_param = write_data["extrude"]
+        vertex_strings = write_data["vertex"]
+        curve_strings = write_data["curve"]
+
+        """Write an .obj file with the curves and verts"""
+        if extrude_param["op"] == 1:  # 'add'
+            set_op = "NewBodyFeatureOperation"
+        elif extrude_param["op"] == 2:  # 'cut'
+            set_op = "CutFeatureOperation"
+        elif extrude_param["op"] == 3:  # 'cut'
+            set_op = "IntersectFeatureOperation"
+
+        with open(obj_file, "w") as fh:
+            # Write Meta info
+            fh.write("# WaveFront *.obj file\n")
+            fh.write("# ExtrudeOperation: " + set_op + "\n")
+            fh.write("\n")
+
+            # Write vertex and curve
+            fh.write(vertex_strings)
+            fh.write("\n")
+            fh.write(curve_strings)
+            fh.write("\n")
+
+            # Write extrude value
+            extrude_string = "Extrude "
+            for value in extrude_param["value"]:
+                extrude_string += str(value) + " "
+            fh.write(extrude_string)
+            fh.write("\n")
+
+            # Write refe plane value
+            p_orig = parse3d_sample(extrude_param["T"])
+            x_axis = parse3d_sample(extrude_param["R"][0:3])
+            y_axis = parse3d_sample(extrude_param["R"][3:6])
+            z_axis = parse3d_sample(extrude_param["R"][6:9])
+            fh.write("T_origin " + p_orig)
+            fh.write("\n")
+            fh.write("T_xaxis " + x_axis)
+            fh.write("\n")
+            fh.write("T_yaxis " + y_axis)
+            fh.write("\n")
+            fh.write("T_zaxis " + z_axis)
+
+
+def parse3d_sample(point3d):
+    x = point3d[0]
+    y = point3d[1]
+    z = point3d[2]
+    return str(x) + " " + str(y) + " " + str(z)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--in-path", type=str, required=True)
+    parser.add_argument("--out-path", type=str, required=True)
+    args = parser.parse_args()
+
+    # with open(args.in_path, "r") as f:
+        # data = f.readlines()
+    with open(args.in_path, 'r') as file:  
+        data = file.read()
+ 
+    data = json.loads(data)
+
+    num_valid_str = 0
+    for idx, item in enumerate(data):
+        try:
+            cad_parser = CADparser(bit=6)
+            # print(idx)
+            if type(item) == str:
+                parsed_data = cad_parser.perform(item)
+            elif type(item) == dict:
+                parsed_data = cad_parser.perform(item['output'])
+            else:
+                raise ValueError("Invalid data type")
+            out_path = os.path.join(args.out_path, str(idx).zfill(6))
+            os.makedirs(out_path, exist_ok=True)
+            if parsed_data is not None:
+                num_valid_str += 1
+                write_obj_sample(out_path, parsed_data)
+        except Exception as e:
+            print(e)
+            pass
+    print(f"Number of valid CAD strings: {num_valid_str}/{len(data)}")

CADFusion/src/train/CAD_dataset.py

@@ -0,0 +1,89 @@
+import json
+import os
+import torch
+import random
+import transformers
+
+from dataclasses import dataclass
+from torch.utils.data import Dataset
+from utils import IGNORE_INDEX, MAX_LENGTH
+
+class CADDataset(Dataset):
+    def __init__(self, json_fn, cutoff=True, llama_tokenizer=None):
+        if not os.path.exists(json_fn):
+            raise ValueError(f"{json_fn} does not exist")
+        self.inputs = json.load(open(json_fn, "r"))
+        print(len(self.inputs))
+        self.inputs = [item for item in self.inputs if 'null' not in item['description']]
+        random.shuffle(self.inputs)
+        if cutoff:
+            self.inputs = self.inputs[:18953]
+        print(len(self.inputs))
+        self.llama_tokenizer = llama_tokenizer
+
+    def __len__(self):
+        return len(self.inputs)
+
+    def __getitem__(self, index):  
+        item = self.inputs[index]  
+        seq = item['command_sequence']  
+        des = item['description']
+        val = self.tokenize(seq, des)
+        return val
+    
+    
+    def tokenize(self, seq, des):  
+        tokens, prompt_length = self.conditional_generation_task(seq=seq, des=des)  
+        input_ids = tokens.input_ids[0]  
+        labels = tokens.input_ids[0].clone()  # Clone the input_ids for labels  
+        # Set the labels for the prompt part to IGNORE_INDEX so they are ignored in loss calculation  
+        labels[:prompt_length] = IGNORE_INDEX  
+        input_id_lens = label_lens = (  
+            tokens.input_ids.ne(self.llama_tokenizer.pad_token_id).sum().item()  
+        )  
+        return dict(  
+            input_ids=input_ids,  
+            input_id_lens=input_id_lens,  
+            labels=labels,  
+            label_lens=label_lens,  
+        )  
+
+    
+    def conditional_generation_task(self, seq, des):  
+        prompt = 'Below is a description of a 3D shape:\n'  
+        prompt += des  
+        prompt += '\nGenerate a Computer-Aided Design(CAD) command sequence of the 3D shape:\n'  
+        full_text = prompt + seq + self.llama_tokenizer.eos_token  
+        tokens = self.llama_tokenizer(  
+            full_text,  
+            max_length=MAX_LENGTH,  
+            return_tensors="pt",  
+            truncation=True,  
+        )  
+        prompt_length = len(self.llama_tokenizer(prompt)['input_ids'])  
+        return tokens, prompt_length  
+    
+    
+@dataclass
+class DataCollatorForSupervisedDataset(object):
+    """Collate examples for supervised fine-tuning."""
+
+    tokenizer: transformers.PreTrainedTokenizer
+
+    def __call__(self, instances):
+        input_ids, labels = tuple(
+            [instance[key].clone().detach() for instance in instances]
+            for key in ("input_ids", "labels")
+        )
+        # force left padding
+        reversed_sequences = [torch.flip(input_id, [0]) for input_id in input_ids]
+        input_ids = torch.nn.utils.rnn.pad_sequence(reversed_sequences, batch_first=True, padding_value=self.tokenizer.pad_token_id)
+        input_ids = torch.flip(input_ids, [0, 1])
+        labels = torch.nn.utils.rnn.pad_sequence(
+            labels, batch_first=True, padding_value=IGNORE_INDEX
+        )
+        return dict(
+            input_ids=input_ids,
+            labels=labels,
+            attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
+        )

CADFusion/src/train/dpo.py

@@ -0,0 +1,79 @@
+import argparse
+import os
+import torch
+import json
+import random
+import transformers
+from huggingface_hub import login
+
+login() # put your huggingface token here
+os.environ["WANDB_PROJECT"] = "CADFusion_VF"
+
+from datasets import Dataset
+from trl import DPOTrainer, DPOConfig
+from utils import prepare_model_and_tokenizer
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--run-name", type=str, required=True)
+parser.add_argument("--lora-rank", type=int, default=32)
+parser.add_argument("--lora-alpha", type=int, default=32)
+parser.add_argument("--lora-dropout", type=float, default=0.05)
+parser.add_argument("--sample-cutoff", default=100000, type=int)
+parser.add_argument("--pretrained-path", type=str, required=True)
+parser.add_argument("--data-path", type=str, required=True)
+parser.add_argument("--output-path", type=str, required=True)
+parser.add_argument("--num-epochs", type=int, default=3)
+parser.add_argument("--batch-size", type=int, default=2)
+parser.add_argument("--eval-freq", default=1000, type=int)
+parser.add_argument("--save-freq", default=500, type=int)
+parser.add_argument("--debug", action="store_true", default=False)
+args = parser.parse_args()
+
+
+
+with open(args.data_path, 'r') as f:
+    raw_data = json.load(f)
+    
+random.shuffle(raw_data)
+
+if len(raw_data) > args.sample_cutoff + 100:
+    ds = {
+        "train": Dataset.from_list(raw_data[:args.sample_cutoff]),
+        "val": Dataset.from_list(raw_data[-100:])
+    }
+else:
+    ds = {
+        "train": Dataset.from_list(raw_data[:-100]),
+        "val": Dataset.from_list(raw_data[-100:])
+        }
+
+llama_model, llama_tokenizer = prepare_model_and_tokenizer(args)
+
+for name, param in llama_model.named_parameters():
+    if "lora" in name:  # Check if "lora" is in the parameter's name
+        param.requires_grad = True
+        
+training_args = DPOConfig(
+    run_name=args.run_name,
+    learning_rate=1.41e-5, 
+    per_device_train_batch_size=2,
+    per_device_eval_batch_size=args.batch_size,
+    report_to="wandb",
+    num_train_epochs=args.num_epochs,
+    do_eval=True,
+    eval_steps=args.eval_freq,
+    save_steps=args.save_freq,
+    output_dir=args.output_path
+    )
+
+trainer = DPOTrainer(
+    llama_model,
+    None,
+    args=training_args,
+    train_dataset=ds['train'],
+    eval_dataset=ds['val'],
+    tokenizer=llama_tokenizer,
+)
+trainer.save_model()
+trainer.train()
+trainer.save_model()

CADFusion/src/train/llama_finetune.py

@@ -0,0 +1,127 @@
+
+import argparse
+import os
+import torch
+import transformers
+
+from CAD_dataset import CADDataset, DataCollatorForSupervisedDataset
+from huggingface_hub import login
+from pathlib import Path
+from peft import LoraConfig, get_peft_model
+from transformers import Trainer, TrainingArguments
+from utils import prepare_model_and_tokenizer
+
+login() # put your huggingface token here
+
+def setup_datasets(args, llama_tokenizer, transform_args={}):
+    datasets = {
+        "train": CADDataset(
+            args.data_path,
+            llama_tokenizer=llama_tokenizer,
+        ),
+        "val": CADDataset(
+            args.eval_data_path,
+            llama_tokenizer=llama_tokenizer,
+        ),
+    }
+    return datasets
+
+
+def setup_training_args(args):
+    output_dir = args.expdir / args.run_name
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    if args.debug:
+        os.environ["WANDB_DISABLED"] = "True"
+    os.environ["ACCELERATE_MIXED_PRECISION"] = "no"
+    training_args = TrainingArguments(
+        fsdp=False,
+        fp16=False,
+        bf16=False,
+        do_eval=True,
+        gradient_checkpointing=False,
+        ddp_find_unused_parameters=False,
+        num_train_epochs=args.num_epochs,
+        eval_steps=args.eval_freq,
+        save_steps=args.save_freq,
+        logging_steps=10,
+        evaluation_strategy="steps",
+        per_device_train_batch_size=args.batch_size,
+        per_device_eval_batch_size=args.batch_size,
+        learning_rate=args.lr,
+        lr_scheduler_type=args.lr_scheduler,
+        warmup_steps=args.num_warmup_steps,
+        weight_decay=args.weight_decay,
+        gradient_accumulation_steps=args.grad_accum,
+        output_dir=output_dir,
+        run_name=args.run_name,
+        report_to="wandb",
+        dataloader_num_workers=8,
+        remove_unused_columns=False,
+        # label_names=["cad_ids"],  # this is to make trainer behave as expected
+    )
+    return training_args
+
+
+def setup_trainer(args):
+    training_args = setup_training_args(args)
+    if args.device_map == 'accelerate':
+        args.device_map = {'': training_args.local_rank}
+    model, llama_tokenizer = prepare_model_and_tokenizer(args)
+
+    datasets = setup_datasets(args, llama_tokenizer)
+
+    data_collator = DataCollatorForSupervisedDataset(
+        tokenizer=llama_tokenizer,
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=datasets["train"],
+        eval_dataset=datasets["val"],
+        data_collator=data_collator,
+    )
+
+    return trainer
+
+
+def main(args):
+    trainer = setup_trainer(args)
+
+    if args.resume_dir is not None:
+        train_result = trainer.train(resume_from_checkpoint=args.resume_dir)
+    else:
+        train_result = trainer.train()
+
+    print(train_result)
+    trainer.save_state()
+    trainer.save_model()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--run-name", type=str, required=True)
+    parser.add_argument("--expdir", type=Path, default="exp")
+    parser.add_argument("--model-name", default="llama3")
+    parser.add_argument("--lora-rank", type=int, default=32)
+    parser.add_argument("--lora-alpha", type=int, default=32)
+    parser.add_argument("--lora-dropout", type=float, default=0.05)
+    parser.add_argument("--data-path", type=Path, default="data/train.json")
+    parser.add_argument("--eval-data-path", type=Path, default="data/eval.json")
+    parser.add_argument("--pretrained-path", type=Path, default=None)
+    parser.add_argument("--num-epochs", type=int, default=40)
+    parser.add_argument("--batch-size", type=int, default=1)
+    parser.add_argument("--grad-accum", type=int, default=1)
+    parser.add_argument("--lr", type=float, default=1e-4)
+    parser.add_argument("--lr-scheduler", type=str, default="cosine")
+    parser.add_argument("--num-warmup-steps", type=int, default=100)
+    parser.add_argument("--weight-decay", type=float, default=0.0)
+    parser.add_argument("--eval-freq", default=1000, type=int)
+    parser.add_argument("--save-freq", default=50000, type=int)
+    parser.add_argument("--device-map", type=str, default='auto')
+    parser.add_argument("--resume-dir", type=Path, default=None)
+    parser.add_argument("--debug", action="store_true", default=False)
+    args = parser.parse_args()
+    os.environ["WANDB_PROJECT"] = "CADFusion_SL"
+    main(args)

CADFusion/src/train/utils.py

@@ -0,0 +1,86 @@
+import torch
+import transformers
+from peft import LoraConfig, PeftModel, get_peft_model
+
+IGNORE_INDEX = -100
+MAX_LENGTH = 512
+DEFAULT_PAD_TOKEN = "[PAD]"
+DEFAULT_EOS_TOKEN = "</s>"
+DEFAULT_BOS_TOKEN = "<s>"
+DEFAULT_UNK_TOKEN = "<unk>"
+
+def smart_tokenizer_and_embedding_resize(
+    special_tokens_dict,
+    llama_tokenizer,
+    model,
+):
+    """Resize tokenizer and embedding.
+
+    Note: This is the unoptimized version that may make your embedding size not be divisible by 64.
+    """
+    num_new_tokens = llama_tokenizer.add_special_tokens(special_tokens_dict)
+    model.resize_token_embeddings(len(llama_tokenizer))
+
+    if num_new_tokens > 0:
+        input_embeddings = model.get_input_embeddings().weight.data
+        output_embeddings = model.get_output_embeddings().weight.data
+
+        input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(
+            dim=0, keepdim=True
+        )
+        output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(
+            dim=0, keepdim=True
+        )
+
+        input_embeddings[-num_new_tokens:] = input_embeddings_avg
+        output_embeddings[-num_new_tokens:] = output_embeddings_avg
+
+def prepare_model_and_tokenizer(args):
+    model_id = "meta-llama/Meta-Llama-3-8B"
+    print(f"Model size: {model_id}")
+    if hasattr(args, 'device_map'):
+        device_map = args.device_map
+    else:
+        device_map = 'auto'
+    pipeline = transformers.pipeline("text2text-generation",
+                                        model=model_id, model_kwargs={"torch_dtype": torch.float32}, device_map=device_map)
+    tokenizer = pipeline.tokenizer
+    base_model = pipeline.model
+
+    special_tokens_dict = dict()
+    if tokenizer.pad_token is None:
+        special_tokens_dict["pad_token"] = DEFAULT_PAD_TOKEN
+    if tokenizer.eos_token is None:
+        special_tokens_dict["eos_token"] = DEFAULT_EOS_TOKEN
+    if tokenizer.bos_token is None:
+        special_tokens_dict["bos_token"] = DEFAULT_BOS_TOKEN
+    if tokenizer.unk_token is None:
+        special_tokens_dict["unk_token"] = DEFAULT_UNK_TOKEN
+
+    smart_tokenizer_and_embedding_resize(
+        special_tokens_dict=special_tokens_dict,
+        llama_tokenizer=tokenizer,
+        model=base_model,
+    )
+    
+    peft_config = LoraConfig(
+        r=args.lora_rank,
+        lora_alpha=args.lora_alpha,
+        lora_dropout=args.lora_dropout,
+        bias="none",
+        task_type="CAUSAL_LM",
+    )
+    
+    tokenizer.padding_side = 'left'
+    peftmodel = get_peft_model(base_model, peft_config)
+    if args.pretrained_path:
+        # load a previous checkpoint if the path is given
+        model = PeftModel.from_pretrained(base_model, args.pretrained_path, device_map=device_map)
+        peft_state_dict = {f"{k}": v for k, v in model.state_dict().items()}
+        peftmodel.load_state_dict(peft_state_dict)
+        
+        for name, param in peftmodel.named_parameters():
+            if "lora" in name:  # Check if "lora" is in the parameter's name
+                param.requires_grad = True  
+    peftmodel.print_trainable_parameters()
+    return peftmodel, tokenizer