HOOPS AI: Use the Dataset Explorer to navigate the dataset
The dataset module provides a comprehensive framework for exploring, navigating, and loading CAD model datasets for machine learning applications. It consists of two primary components that work together to simplify data handling:
DatasetExplorer - For exploring and querying dataset contents
DatasetLoader - For loading and preparing datasets for machine learning training
These components are designed to work with the processed data from the cadaccess and cadencoder modules, as well as the outputs from the flow pipeline system. They provide high-level abstractions that allow users to focus on machine learning tasks rather than data handling complexities.
DatasetExplorer
The DatasetExplorer class (dataset_explorer.py) provides methods for exploring and querying datasets stored in Zarr format (.dataset) with accompanying metadata (.infoset) in Parquet files. This class focuses on data discovery, filtering, and statistical analysis.
Key Methods
Data Discovery and Metadata
available_groups() -> set: Returns the set of available dataset groups (faces, edges, file, etc.)get_descriptions(table_name: str, key_id: Optional[int] = None, use_wildchar: Optional[bool] = False) -> pd.DataFrame: Retrieves metadata descriptions (labels, face types, edge types, etc.)get_parquet_info_by_code(file_id_code: int): Returns rows from the Parquet file for a specific file ID codeget_file_info_all() -> pd.DataFrame: Returns all file info from the Parquet metadata
Data Distribution Analysis
create_distribution(key: str, bins: int = 10, group: str = "faces") -> Dict[str, Any]: Computes histograms of data distributions using Dask for parallel processing
Data Filtering and Selection
get_file_list(group: str, where: Callable[[xr.Dataset], xr.DataArray]) -> List[str]: Returns file IDs matching a boolean filter conditionfile_dataset(file_id_code: int, group: str) -> xr.Dataset: Returns a subset of the dataset for a specific filebuild_membership_matrix(group: str, key: str, bins_or_categories: Union[int, List, np.ndarray], as_counts: bool = False) -> tuple[np.ndarray, np.ndarray, np.ndarray]: Builds a file-by-bin membership matrix for stratified splittingdecode_file_id_code(code: int) -> str: Converts an integer file ID code to the original string identifier
[1]:
import hoops_ai
import os
hoops_ai.set_license(hoops_ai.use_test_license())
ℹ️ Using TEST LICENSE (expires February 8th, 2026 - 9 days remaining)
For production use, obtain your own license from Tech Soft 3D
HOOPS AI version : 1.0.0-b2dev12
======================================================================
✓ HOOPS AI License: Valid (TEST LICENSE)
======================================================================
[2]:
from hoops_ai.dataset import DatasetExplorer
import pathlib
# Define paths
flow_name = "cadsynth_aag"
flow_root_dir = pathlib.Path.cwd().parent.joinpath("packages", "flows", flow_name)
print(flow_root_dir)
parquet_file_path = str(flow_root_dir.joinpath(f"{flow_name}.infoset"))
merged_store_path = str(flow_root_dir.joinpath(f"{flow_name}.dataset"))
parquet_file_attribs = str(flow_root_dir.joinpath(f"{flow_name}.attribset"))
explorer = DatasetExplorer(merged_store_path=merged_store_path, parquet_file_path=parquet_file_path, parquet_file_attribs=parquet_file_attribs)
explorer.print_table_of_contents()
C:\Users\LuisSalazar\Documents\MAIN\MLProject\repo\HOOPS-AI-tutorials\packages\flows\cadsynth_aag
C:\Users\LuisSalazar\miniconda3\envs\hoops_ai_gpu_dev\lib\contextlib.py:126: UserWarning: Creating scratch directories is taking a surprisingly long time. (1.02s) This is often due to running workers on a network file system. Consider specifying a local-directory to point workers to write scratch data to a local disk.
next(self.gen)
[DatasetExplorer] Default local cluster started: <Client: 'tcp://127.0.0.1:64734' processes=1 threads=16, memory=7.45 GiB>
--- Dataset Table of Contents ---
EDGES_GROUP:
EDGE_CONVEXITIES_DATA: Shape: (11724557,), Dims: ('edge',), Size: 11724557
EDGE_DIHEDRAL_ANGLES_DATA: Shape: (11724557,), Dims: ('edge',), Size: 11724557
EDGE_INDICES_DATA: Shape: (11724557,), Dims: ('edge',), Size: 11724557
EDGE_LENGTHS_DATA: Shape: (11724557,), Dims: ('edge',), Size: 11724557
EDGE_TYPES_DATA: Shape: (11724557,), Dims: ('edge',), Size: 11724557
EDGE_U_GRIDS_DATA: Shape: (11724557, 5, 6), Dims: ('edge', 'dim_x', 'component'), Size: 351736710
FILE_ID_CODE_EDGES_DATA: Shape: (11724557,), Dims: ('edge',), Size: 11724557
FACEFACE_GROUP:
A3_DISTANCE_DATA: Shape: (137844937, 64), Dims: ('faceface_flat', 'bins_d'), Size: 8822075968
D2_DISTANCE_DATA: Shape: (137844937, 64), Dims: ('faceface_flat', 'bins_a'), Size: 8822075968
EXTENDED_ADJACENCY_DATA: Shape: (137844937,), Dims: ('faceface_flat',), Size: 137844937
FACE_PAIR_EDGES_PATH_DATA: Shape: (137844937, 32), Dims: ('faceface_flat', 'dim_path'), Size: 4411037984
FACE_X_DATA: Shape: (137844937,), Dims: ('faceface_flat',), Size: 137844937
FACE_Y_DATA: Shape: (137844937,), Dims: ('faceface_flat',), Size: 137844937
FILE_ID_CODE_FACEFACE_DATA: Shape: (137844937,), Dims: ('faceface_flat',), Size: 137844937
FACES_GROUP:
FACE_AREAS_DATA: Shape: (4492347,), Dims: ('face',), Size: 4492347
FACE_INDICES_DATA: Shape: (4492347,), Dims: ('face',), Size: 4492347
FACE_LABELS_DATA: Shape: (4492347,), Dims: ('face',), Size: 4492347
FACE_LOOPS_DATA: Shape: (4492347,), Dims: ('face',), Size: 4492347
FACE_NEIGHBORSCOUNT_DATA: Shape: (4492347,), Dims: ('face',), Size: 4492347
FACE_TYPES_DATA: Shape: (4492347,), Dims: ('face',), Size: 4492347
FACE_UV_GRIDS_DATA: Shape: (4492347, 5, 5, 7), Dims: ('face', 'uv_x', 'uv_y', 'component'), Size: 786160725
FILE_ID_CODE_FACES_DATA: Shape: (4492347,), Dims: ('face',), Size: 4492347
FILE_GROUP:
DURATION_DGLCONVERT_DATA: Shape: (162412,), Dims: ('file',), Size: 162412
FILE_ID_CODE_FILE_DATA: Shape: (162412,), Dims: ('file',), Size: 162412
SIZE_CADFILE_DATA: Shape: (162412,), Dims: ('file',), Size: 162412
SIZE_DGLFILE_DATA: Shape: (162412,), Dims: ('file',), Size: 162412
GRAPH_GROUP:
DESTINATION_DATA: Shape: (11724557,), Dims: ('edge',), Size: 11724557
FILE_ID_CODE_GRAPH_DATA: Shape: (11724557,), Dims: ('edge',), Size: 11724557
NUM_NODES_DATA: Shape: (11724557,), Dims: ('edge',), Size: 11724557
SOURCE_DATA: Shape: (11724557,), Dims: ('edge',), Size: 11724557
==================================
Columns in file_info:
name description subset id table_name
0 00001d44b8eb37fe4bb6ec4e89746ea3 ...Temp\Cadsynth_aag\step\00021388.stp train 0 file_info
1 000093c3b74c8076fbba20bf8613e2a2 ..._aag\step\20221121_154647_2127.step validation 1 file_info
2 00012dc0e22bd178d7c6a12436734130 ...Temp\Cadsynth_aag\step\00067737.stp validation 2 file_info
3 00013001df401fe06d5213ef6fc9e581 ...Temp\Cadsynth_aag\step\00058353.stp validation 3 file_info
4 00015d0e0e285c5089a485360b34db1a ...Temp\Cadsynth_aag\step\00048571.stp train 4 file_info
5 0001841a18705e7d132a7bf07aa2bc19 ..._aag\step\20221124_154714_8374.step train 5 file_info
6 00018ba5cd4389699cc374231f667176 ...Temp\Cadsynth_aag\step\00016855.stp train 6 file_info
7 0001c365725a060ad6cd4aade3a3527f ...Temp\Cadsynth_aag\step\00055043.stp validation 7 file_info
8 0001d760a84ddcdabbd0edaa1a5fd2b0 ...aag\step\20221124_154714_15394.step test 8 file_info
9 00024cff62d011c4a78b608c73d08f63 ...Temp\Cadsynth_aag\step\00014960.stp train 9 file_info
... ... ... ... ... ...
162402 fffbf8786af40d631fb540e4c2ce799f ...Temp\Cadsynth_aag\step\00014282.stp validation 162402 file_info
162403 fffc41e365c1823f5f3f42625d82d1c1 ...aag\step\20221124_154714_11719.step test 162403 file_info
162404 fffc5e96b9879cde780f2703088f26d0 ...Temp\Cadsynth_aag\step\00044702.stp train 162404 file_info
162405 fffc766d355768cc5a8ff7318fc09d1e ...Temp\Cadsynth_aag\step\00026869.stp train 162405 file_info
162406 fffc948efb43ac02ae6f2c11d2818c63 ..._aag\step\20221124_154714_5834.step validation 162406 file_info
162407 fffe278b406958997fff4c1f76b72ee0 ..._aag\step\20221123_142528_2390.step test 162407 file_info
162408 fffe3fba779b5fa5543c4a82c3cfbabb ...Temp\Cadsynth_aag\step\00091954.stp train 162408 file_info
162409 fffee3edd46107ca808a958811dcb9a8 ...aag\step\20221124_154714_13628.step validation 162409 file_info
162410 ffff1ac0abc28d3ede6b7bfb6434cdf4 ...Temp\Cadsynth_aag\step\00065860.stp train 162410 file_info
162411 ffffdc1d0cb629065026fa55a97ca314 ...Temp\Cadsynth_aag\step\00016373.stp train 162411 file_info
[ ]:
[3]:
groups = explorer.available_groups()
print(groups)
{'edges', 'file', 'faces', 'graph', 'faceface'}
[4]:
face_type_description = explorer.get_descriptions("face_types")
print(type(face_type_description), face_type_description)
<class 'pandas.core.frame.DataFrame'> id name description table_name
0 0 Plane not set face_types
1 1 Cylinder not set face_types
2 2 Cone not set face_types
3 14 Extrusion not set face_types
4 3 Sphere not set face_types
5 4 Torus not set face_types
6 5 Nurbs not set face_types
7 13 Revolution not set face_types
[ ]:
[5]:
# Get and print meta data information
file_id = 25
df_info = explorer.get_parquet_info_by_code(file_id)
#print(type(df_info), df_info)
[6]:
# Access various dataset groups
file_datasetGroup = {grp: explorer.file_dataset(file_id_code=file_id, group=grp) for grp in groups}
print(f"Datasets (Table of Content) for file ID '{file_id}':")
for grp, ds in file_datasetGroup.items():
for name, da in ds.data_vars.items():
print(f" [{grp}] DATA: {name}, Shape: {da.shape}, Dims: {da.dims}, Size: {da.size}")
print()
file_dataset = file_datasetGroup["faces"]
print("type of file_data_arrays", type(file_dataset))
#print the areas of each face
array_areas = file_dataset["face_areas"].data.compute()
print("type of array_areas", type(array_areas))
print("brep surfaces", array_areas.shape)
Datasets (Table of Content) for file ID '25':
[edges] DATA: edge_convexities, Shape: (128,), Dims: ('edge',), Size: 128
[edges] DATA: edge_dihedral_angles, Shape: (128,), Dims: ('edge',), Size: 128
[edges] DATA: edge_indices, Shape: (128,), Dims: ('edge',), Size: 128
[edges] DATA: edge_lengths, Shape: (128,), Dims: ('edge',), Size: 128
[edges] DATA: edge_types, Shape: (128,), Dims: ('edge',), Size: 128
[edges] DATA: edge_u_grids, Shape: (128, 5, 6), Dims: ('edge', 'dim_x', 'component'), Size: 3840
[edges] DATA: file_id_code_edges, Shape: (128,), Dims: ('edge',), Size: 128
[file] DATA: duration_dglconvert, Shape: (1,), Dims: ('file',), Size: 1
[file] DATA: file_id_code_file, Shape: (1,), Dims: ('file',), Size: 1
[file] DATA: size_cadfile, Shape: (1,), Dims: ('file',), Size: 1
[file] DATA: size_dglfile, Shape: (1,), Dims: ('file',), Size: 1
[faces] DATA: face_areas, Shape: (45,), Dims: ('face',), Size: 45
[faces] DATA: face_indices, Shape: (45,), Dims: ('face',), Size: 45
[faces] DATA: face_labels, Shape: (45,), Dims: ('face',), Size: 45
[faces] DATA: face_loops, Shape: (45,), Dims: ('face',), Size: 45
[faces] DATA: face_neighborscount, Shape: (45,), Dims: ('face',), Size: 45
[faces] DATA: face_types, Shape: (45,), Dims: ('face',), Size: 45
[faces] DATA: face_uv_grids, Shape: (45, 5, 5, 7), Dims: ('face', 'uv_x', 'uv_y', 'component'), Size: 7875
[faces] DATA: file_id_code_faces, Shape: (45,), Dims: ('face',), Size: 45
[graph] DATA: destination, Shape: (128,), Dims: ('edge',), Size: 128
[graph] DATA: file_id_code_graph, Shape: (128,), Dims: ('edge',), Size: 128
[graph] DATA: num_nodes, Shape: (128,), Dims: ('edge',), Size: 128
[graph] DATA: source, Shape: (128,), Dims: ('edge',), Size: 128
[faceface] DATA: a3_distance, Shape: (2025, 64), Dims: ('faceface_flat', 'bins_d'), Size: 129600
[faceface] DATA: d2_distance, Shape: (2025, 64), Dims: ('faceface_flat', 'bins_a'), Size: 129600
[faceface] DATA: extended_adjacency, Shape: (2025,), Dims: ('faceface_flat',), Size: 2025
[faceface] DATA: face_pair_edges_path, Shape: (2025, 32), Dims: ('faceface_flat', 'dim_path'), Size: 64800
[faceface] DATA: face_x, Shape: (2025,), Dims: ('faceface_flat',), Size: 2025
[faceface] DATA: face_y, Shape: (2025,), Dims: ('faceface_flat',), Size: 2025
[faceface] DATA: file_id_code_faceface, Shape: (2025,), Dims: ('faceface_flat',), Size: 2025
type of file_data_arrays <class 'xarray.core.dataset.Dataset'>
type of array_areas <class 'numpy.ndarray'>
brep surfaces (45,)
[ ]:
[7]:
# this requires to be a bit familiar with pandas and dask.
uv_grid_data = file_dataset["face_uv_grids"].data.compute()
print("numpy array shape", uv_grid_data.shape)
numpy array shape (45, 5, 5, 7)
[8]:
#print(uv_grid_data)
[9]:
explorer.get_file_info_all()
[9]:
| name | description | subset | id | table_name | |
|---|---|---|---|---|---|
| 0 | 00001d44b8eb37fe4bb6ec4e89746ea3 | C:\Temp\Cadsynth_aag\step\00021388.stp | train | 0 | file_info |
| 1 | 000093c3b74c8076fbba20bf8613e2a2 | C:\Temp\Cadsynth_aag\step\20221121_154647_2127... | validation | 1 | file_info |
| 2 | 00012dc0e22bd178d7c6a12436734130 | C:\Temp\Cadsynth_aag\step\00067737.stp | validation | 2 | file_info |
| 3 | 00013001df401fe06d5213ef6fc9e581 | C:\Temp\Cadsynth_aag\step\00058353.stp | validation | 3 | file_info |
| 4 | 00015d0e0e285c5089a485360b34db1a | C:\Temp\Cadsynth_aag\step\00048571.stp | train | 4 | file_info |
| ... | ... | ... | ... | ... | ... |
| 162407 | fffe278b406958997fff4c1f76b72ee0 | C:\Temp\Cadsynth_aag\step\20221123_142528_2390... | test | 162407 | file_info |
| 162408 | fffe3fba779b5fa5543c4a82c3cfbabb | C:\Temp\Cadsynth_aag\step\00091954.stp | train | 162408 | file_info |
| 162409 | fffee3edd46107ca808a958811dcb9a8 | C:\Temp\Cadsynth_aag\step\20221124_154714_1362... | validation | 162409 | file_info |
| 162410 | ffff1ac0abc28d3ede6b7bfb6434cdf4 | C:\Temp\Cadsynth_aag\step\00065860.stp | train | 162410 | file_info |
| 162411 | ffffdc1d0cb629065026fa55a97ca314 | C:\Temp\Cadsynth_aag\step\00016373.stp | train | 162411 | file_info |
162412 rows × 5 columns
[10]:
# Visualization libraries
import matplotlib.pyplot as plt
def print_distribution_info(dist, title="Distribution"):
"""Helper function to print and visualize distribution data."""
list_filecount = list()
for i, bin_files in enumerate(dist['file_id_codes_in_bins']):
list_filecount.append(bin_files.size)
dist['file_count'] =list_filecount
# Visualization with matplotlib
fig, ax = plt.subplots(figsize=(12, 4))
bin_centers = 0.5 * (dist['bin_edges'][1:] + dist['bin_edges'][:-1])
ax.bar(bin_centers, dist['file_count'], width=(dist['bin_edges'][1] - dist['bin_edges'][0]),
alpha=0.7, color='steelblue', edgecolor='black', linewidth=1)
# Add file count annotations
for i, count in enumerate(dist['file_count']):
if count > 0: # Only annotate non-empty bins
ax.text(bin_centers[i], count + 0.5, f"{count}",
ha='center', va='bottom', fontsize=8)
ax.set_xlabel('Value')
ax.set_ylabel('Count')
ax.set_title(f'{title} Histogram')
ax.grid(True, linestyle='--', alpha=0.7)
plt.tight_layout()
plt.show()
[11]:
import time
start_time = time.time()
face_dist = explorer.create_distribution(key="face_labels", bins=None, group="faces")
print(f"Face labels distribution created in {(time.time() - start_time):.2f} seconds\n")
print_distribution_info(face_dist, title="Labels")
Face labels distribution created in 3.28 seconds
[12]:
print(explorer.get_descriptions("face_labels"))
id name description table_name
14 0 no-label not set face_labels
15 7 2sides_through_step not set face_labels
16 11 rectangular_blind_slot not set face_labels
17 18 circular blind step not set face_labels
18 5 6sides_passage not set face_labels
19 9 rectangular_blind_step not set face_labels
20 12 rectangular_pocket not set face_labels
21 13 triangular_pocket not set face_labels
22 14 6sides_pocket not set face_labels
23 21 circular end pocket not set face_labels
24 23 blind hole not set face_labels
25 2 triangular_through_slot not set face_labels
26 17 through hole not set face_labels
27 3 rectangular_passage not set face_labels
28 16 circular through slot not set face_labels
29 4 triangular_passage not set face_labels
30 6 rectangular_through_step not set face_labels
31 10 triangular_blind_step not set face_labels
32 24 fillet not set face_labels
33 1 rectangular_through_slot not set face_labels
34 15 chamfer not set face_labels
35 19 horizontal circular end blind slot not set face_labels
36 22 o-ring not set face_labels
37 8 slanted_through_step not set face_labels
38 20 vertical circular end blind slot not set face_labels
[13]:
start_time = time.time()
dist = explorer.create_distribution(key="num_nodes", bins=12, group="graph")
print(f"Face labels distribution created in {(time.time() - start_time):.2f} seconds\n")
print_distribution_info(dist, title="Brep facecount Distribution")
Face labels distribution created in 4.67 seconds
Gather files that fulfilled a given condition. Filter
[ ]:
[14]:
start_time = time.time()
# condition
label_is_pipefittings = lambda ds: ds['face_labels'] == 15
filelist = explorer.get_file_list(group="faces", where=label_is_pipefittings)
print(f"Filtering completed in {(time.time() - start_time):.2f} seconds")
print(f"Found {len(filelist)} files with face_labels == 15 (chamfer)\n")
print(filelist)
Filtering completed in 0.31 seconds
Found 32234 files with face_labels == 15 (chamfer)
[ 10 12 15 ... 162392 162398 162401]
Query data for single file
[15]:
def demo_query_single_file(explorer, file_id):
"""Show how to access and query dataset details for a single file."""
print("=== Single File Dataset Access ===")
import time
# Get and print parquet info
df_info = explorer.get_parquet_info_by_code(file_id)
print("Files info:")
for column in df_info.columns:
print(f"Column: {column}")
for value in df_info[column]:
print(f" {value}")
print()
# Access various dataset groups
groups = ["faces", "file", "edges", "graph"]
datasets = {grp: explorer.file_dataset(file_id_code=file_id, group=grp) for grp in groups}
print(f"Datasets for file ID '{file_id}':")
for grp, ds in datasets.items():
for name, da in ds.data_vars.items():
print(f" [{grp}] VARIABLE: {name}, Shape: {da.shape}, Dims: {da.dims}, Size: {da.size}")
print()
# Query uv grids data for a specific face
start_time = time.time()
uv_grid_data = datasets["faces"]["face_uv_grids"].isel(face=2)
print("uv_grids data for face index 2:")
np_uvgrid = uv_grid_data.data.compute()
print(f"Query took {(time.time() - start_time):.2f} seconds\n")
[16]:
demo_query_single_file(explorer,file_id=4500)
=== Single File Dataset Access ===
Files info:
Column: name
07207dfd094fe0ebe9368ded4c271b23
Column: description
C:\Temp\Cadsynth_aag\step\20221124_154714_17096.step
Column: subset
test
Column: id
4500
Column: table_name
file_info
Datasets for file ID '4500':
[faces] VARIABLE: face_areas, Shape: (38,), Dims: ('face',), Size: 38
[faces] VARIABLE: face_indices, Shape: (38,), Dims: ('face',), Size: 38
[faces] VARIABLE: face_labels, Shape: (38,), Dims: ('face',), Size: 38
[faces] VARIABLE: face_loops, Shape: (38,), Dims: ('face',), Size: 38
[faces] VARIABLE: face_neighborscount, Shape: (38,), Dims: ('face',), Size: 38
[faces] VARIABLE: face_types, Shape: (38,), Dims: ('face',), Size: 38
[faces] VARIABLE: face_uv_grids, Shape: (38, 5, 5, 7), Dims: ('face', 'uv_x', 'uv_y', 'component'), Size: 6650
[faces] VARIABLE: file_id_code_faces, Shape: (38,), Dims: ('face',), Size: 38
[file] VARIABLE: duration_dglconvert, Shape: (1,), Dims: ('file',), Size: 1
[file] VARIABLE: file_id_code_file, Shape: (1,), Dims: ('file',), Size: 1
[file] VARIABLE: size_cadfile, Shape: (1,), Dims: ('file',), Size: 1
[file] VARIABLE: size_dglfile, Shape: (1,), Dims: ('file',), Size: 1
[edges] VARIABLE: edge_convexities, Shape: (96,), Dims: ('edge',), Size: 96
[edges] VARIABLE: edge_dihedral_angles, Shape: (96,), Dims: ('edge',), Size: 96
[edges] VARIABLE: edge_indices, Shape: (96,), Dims: ('edge',), Size: 96
[edges] VARIABLE: edge_lengths, Shape: (96,), Dims: ('edge',), Size: 96
[edges] VARIABLE: edge_types, Shape: (96,), Dims: ('edge',), Size: 96
[edges] VARIABLE: edge_u_grids, Shape: (96, 5, 6), Dims: ('edge', 'dim_x', 'component'), Size: 2880
[edges] VARIABLE: file_id_code_edges, Shape: (96,), Dims: ('edge',), Size: 96
[graph] VARIABLE: destination, Shape: (96,), Dims: ('edge',), Size: 96
[graph] VARIABLE: file_id_code_graph, Shape: (96,), Dims: ('edge',), Size: 96
[graph] VARIABLE: num_nodes, Shape: (96,), Dims: ('edge',), Size: 96
[graph] VARIABLE: source, Shape: (96,), Dims: ('edge',), Size: 96
uv_grids data for face index 2:
Query took 1.17 seconds
[ ]:
