fa1f75109bfe4e7c9c98b7d8ff29120b

Indexing a CAD Dataset with HOOPS Embeddings

This notebook walks through the full indexing pipeline: loading a CAD dataset, computing embeddings in batch, building a FAISS index, and saving it to disk.

The index produced here is consumed directly by `demo_HOOPS_Embeddings_retrieval.ipynb <./demo_HOOPS_Embeddings_retrieval.ipynb>`__ for similarity search.

Pre-trained vs. custom model: The pre-trained SIGNAL model bundled with the HOOPS AI release is used here. To index with a model trained on your own data, register your checkpoint in the same way — see `demo_HOOPS_EMBEDDINGS_training.ipynb <./demo_HOOPS_EMBEDDINGS_training.ipynb>`__.

[1]:

import hoops_ai
import os
import sys

license_key = os.environ.get("HOOPS_AI_LICENSE")
if not license_key:
    sys.exit("HOOPS_AI_LICENSE environment variable is required.")

hoops_ai.set_license(license_key, validate=True)

------------------------------------------------------------
HOOPS AI
------------------------------------------------------------
  Platform      : Windows 11
  Architecture  : AMD64
  Python        : 3.12.10
------------------------------------------------------------
  Core          : hoops-ai             1.1.0  (build: ed23c844 2026-06-12T14:14:13Z)
  CAD Access    : hoops-exchange       26.2.0  (build: 1e11169 2026-06-12T10:38:16Z)
  Conversion    : hoops-converter      26.1.1  (build: 00dc9f6 2026-06-12T10:22:46Z)
  Insights      : hoops-web-viewer     26.1.1  (build: d30058f 2026-06-12T10:22:25Z)
------------------------------------------------------------
======================================================================
[OK] HOOPS AI License: Valid
======================================================================

Dataset

This notebook uses the TMCAD dataset (Truly Mechanical CAD Dataset), introduced in:

Zou, Q., & Zhu, L. (2025). Bringing attention to CAD: Boundary representation learning via transformer. Computer-Aided Design, 103940. Elsevier.

Download links:

TMCAD v2 is a cleaned and verified collection of ~10,000 CAD models in STEP (.stp) format across 10 mechanical part categories (bearing, bolt-screw, bracket, flange, gear, nut, shaft, coupling, pulley, spring). Released under the GPL-3.0 license.

[2]:

from hoops_ai.storage import CADFileRetriever, LocalStorageProvider
import pathlib

path_tmcad_v2_dataset = os.environ.get("PATH_DATASET_TMCAD")
if not path_tmcad_v2_dataset:
    sys.exit("PATH_DATASET_TMCAD environment variable is required.")

cad_sources = pathlib.Path(path_tmcad_v2_dataset)
if not cad_sources.exists():
    sys.exit(f"PATH_DATASET_TMCAD path does not exist: {cad_sources}")

retriever = CADFileRetriever(storage_provider=LocalStorageProvider(directory_path=cad_sources), formats=[".stp", ".step", ".iges", ".igs"])

# Get files using the library's retriever
cad_files = retriever.get_file_list()
print(len(cad_files), "files found." )

output_dir = pathlib.Path.cwd().joinpath("out", "tmcad")
output_dir.joinpath("images").mkdir(parents=True, exist_ok=True)       # needed by my_specs

my_specs = {"generate_images": True, "images_out_dir": output_dir.joinpath("images")}

10896 files found.

1. Register the model

Register the model checkpoint with HOOPSEmbeddings before use. The model_name string is used to instantiate the embedder in the next cell.

[4]:

from hoops_ai.ml.embeddings import HOOPSEmbeddings

# HOOPS AI release 1.0
#HOOPSEmbeddings.register_model(
#    model_name="HOOPS Embeddings",
#    checkpoint_path=str(pathlib.Path.cwd().parent.joinpath("packages","trained_ml_models", "ts3d_1M_hoops_embeddings.ckpt"))
#)

# HOOPS AI release 1.1
HOOPSEmbeddings.register_model(
    model_name="HOOPS Embeddings SIGNAL preview",
    checkpoint_path=str(pathlib.Path.cwd().parent.joinpath("packages","trained_ml_models", "ts3d_2M_hoops_embeddings_SIGNAL-preview.ckpt"))
)

print(HOOPSEmbeddings.list_available_models())

['HOOPS Embeddings SIGNAL preview']

2. Compute embeddings and build the index

Embed a batch of CAD files, load the resulting EmbeddingBatch into a CADSearch index, and save the index to disk.

[5]:

SIGNAL_embeddings_model = HOOPSEmbeddings(model="HOOPS Embeddings SIGNAL preview")

print(f"Using model: {SIGNAL_embeddings_model.model_name} with dimension: {SIGNAL_embeddings_model.embedding_dim}")

Successfully loaded model from checkpoint: C:\Users\LuisSalazar.LY-LS-LEGION\Documents\repos\HOOPS-AI-tutorials\packages\trained_ml_models\ts3d_2M_hoops_embeddings_SIGNAL-preview.ckpt
Using model: HOOPS Embeddings SIGNAL preview with dimension: 2048

C:\Users\LuisSalazar.LY-LS-LEGION\Documents\hoops-ai-packages-test\.venv\Lib\site-packages\pytorch_lightning\utilities\parsing.py:136: pytorch_lightning save_hyperparameters: frame introspection failed. Nuitka-compiled frames have empty locals, which breaks PL's collect_init_args. Pass an explicit dict to save_hyperparameters() for correct hparams under Nuitka.

[ ]:

from hoops_ai.ml.embeddings import Embedding, EmbeddingBatch

demoEmbeddings = SIGNAL_embeddings_model.embed_shape_batch(cad_files, num_workers=12, show_progress=True, specifications = my_specs)

# Inspect results
print(f"Successfully embedded: {len(demoEmbeddings.ids)} files")
print(f"Embedding matrix shape: {demoEmbeddings.values.shape}")  # (n_files, dim)
print(f"Failed: {demoEmbeddings.metadata['failed_count']}")
print(f"Model used: {demoEmbeddings.model}")

Successfully embedded: 821 files
Embedding matrix shape: (821, 2048)
Failed: 0
Model used: CUSTOM:HOOPS Embeddings SIGNAL preview

[7]:

from hoops_ai.ml import CADSearch

emb_searcher = CADSearch(shape_model=SIGNAL_embeddings_model)

emb_searcher.index_shape(demoEmbeddings)

[8]:

output_dir.joinpath("vectorstores").mkdir(parents=True, exist_ok=True) # needed by save_shape_index
emb_searcher.save_shape_index(output_dir.joinpath("vectorstores","TMCAD_SIGNAL.faiss"))

Output

The index is now saved under out/tmcad/:

out/tmcad/
├── images/               ← optional per-part renders (generated_images=True)
│   ├── bearing/
│   └── ...
└── vectorstores/
    ├── TMCAD_SIGNAL.faiss   ← vector index
    └── TMCAD_SIGNAL.meta    ← id-to-path mapping

Next steps