Add additive manufacturing example - from shape deviation prediction and compensation

examples/additive_manufacturing/compensation/README.md

@@ -0,0 +1,85 @@
+
+
+# PyTorch version of deformation predictor & compensation
+
+## Key requirments
+
+1. ``Torch_Geometric 2.5.1 or above``: PyTorch based geometric/graph neural network library
+
+   - https://pytorch-geometric.readthedocs.io/en/latest/install/installation.html#installation-via-anaconda
+
+   - conda install pyg=*=*cu* -c pyg
+
+2. ``pip install trimesh``
+
+3. ``pip install matplotlib``
+
+4. ``pip install pandas``
+
+5. ``pip install hydra-core --upgrade --pre``
+
+6. ``PyTorch3D``: PyTorch based 3D computer vision library 
+
+   - Check requirements from official install page: https://github.com/facebookresearch/pytorch3d/blob/main/INSTALL.md
+   - when tested, Pytorch3D requires Python 3.8, 3.9 or 3.10
+
+   - ``pip install -U iopath``
+
+   - Install directly from the source ``pip install "git+https://github.com/facebookresearch/pytorch3d.git@stable" ``
+
+7. ``pip install torch-cluster``
+
+To test in customized CUDA environment, install compatible torch version compatible with cudatoolkit, i.e.
+
+``pip install torch==2.2.1 torchvision==0.17.1 torchaudio==2.2.1 --index-url https://download.pytorch.org/whl/cu121``
+
+ Refer to: 
+https://pytorch.org/get-started/previous-versions/
+
+Other dependencies for development: 
+
+- ``open3d``: pip install open3d, tested version 0.18.0
+- ``torch-cluster``: conda install pytorch-cluster -c pyg
+
+
+
+## Dataset
+- Currently available: 
+  - Bar repository [link not working yet](https://duckduckgo.com)
+  - Molded-fiber repository [link not working yet](https://duckduckgo.com)
+
+- Sample input data folder format: 
+
+   -  input_data.txt: logs for each row, the build geometry folder 
+
+      - /part_folder_i:
+
+         - cad_<part_id>.txt: contains 3 columns for point location 
+
+         - scan_red<part_id>.csv: contains 3 columns for point location 
+
+- Post-processing: 
+
+    - https://github.azc.ext.hp.com/Shape-Compensation/Shape_compensator
+
+
+## Training
+
+- To test running with cpu ``Connfig.yaml`` setting (not recommended): 
+
+    - `` cuda: False ``
+    - ``use_distributed: False``
+    - ``use_multigpu: False``
+- Gpu training: set params listed above to True
+
+- There are two training codes that need to run in sequential manner.
+1. ``train_dis.py``: This code trains the discriminator (predict part deformations with its position and geometry) 
+2. ``train_gen.py``: This code trains the generator (compensate part geometry)
+
+## Inference
+
+- Supported 3D formats: 
+    - Stereolitography (STL)
+    - Wavefront file (OBJ)
+- How to run: 
+  - ``python inference.py`` 

examples/additive_manufacturing/compensation/conf/config.yaml

@@ -0,0 +1,64 @@
+# ignore_header_test
+# ruff: noqa: E402
+
+# © Copyright 2023 HP Development Company, L.P.
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+general:
+  seed: 1234
+  random_sample: True
+  cuda: True
+  use_distributed: False
+  sync_batch: True
+  use_multigpu: False # Default: False for D, True for G
+
+train_dis_options:
+  model_path: './pretrained/11parts_lr-3/pred_model_0000.pth'
+  log_dir: './pretrained/11parts_lr-3/'
+  save_path: './pretrained/11parts_lr-3/'
+  num_epoch: 15001
+  num_batch: 2
+  learning_rate: 0.0001
+  pretrain: False
+  num_points: 190000
+
+train_gen_options:
+  num_points: 190000
+  pred_model_path: "./pretrained/ocardo_iso_p500k/pred_model_3000.pth"   # For D
+  gen_model_path: "./pretrained/ocardo_iso_p500k/pred_model_3000.pth"    # For G
+  log_dir: './pretrained/ocardo_iso_p500k/'
+  save_path: './pretrained/ocardo_iso_p500k/'
+  num_epoch: 50001
+  num_batch: 1
+  learning_rate: 0.001
+
+inference_options:
+  seed: 1234
+  num_points: 20000   # 'Num of points to use'
+  data_path: '/home/chenle/codes/DL_prediction_compensation-master/data/molded_fiber/10/cad' # for other dataset: 'input_data_bar_sample','molded_fiber'
+  discriminator_path: './pretrained/pretrained_os/pred_model_3000.pth'   # 'discriminator model path'
+  generator_path: './pretrained/pretrained_os/gen_model_46500.pth'  # 'generator model path'
+  save_path: './output/test'  # 'save output path'
+  save_extra: True  # 'exports prediction and additional data csv'
+
+
+data_options:
+  dataset_name: "Ocardo" # choices=['Ocardo', 'Bar']
+  data_path: '/home/chenle/codes/DL_prediction_compensation-master/data/molded_fiber' # for other dataset: 'input_data_bar_sample','molded_fiber'
+  cad_format: 'txt'
+  scan_format: 'csv'

examples/additive_manufacturing/compensation/data_proces/ply_sampling.py

@@ -0,0 +1,148 @@
+# ignore_header_test
+# ruff: noqa: E402
+
+# © Copyright 2023 HP Development Company, L.P.
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import os
+
+import numpy as np
+import open3d as o3d
+import pandas as pd
+import torch
+import torch_geometric
+import trimesh
+
+
+def generate_mesh_train(
+    ply_path,
+    scan_pcd_path,
+    save_csv=True,
+    save_mesh_path=None,
+    part_name="bar",
+    part_id="3",
+    export_format="pth",
+    filter_dist=False,
+):
+    """
+    A PLY file is a computer file format for storing 3D data as a collection of polygons.
+    PLY stands for Polygon File Format, and it's also known as the Stanford Triangle Format.
+    PLY files are used to store 3D data from 3D scanners.
+
+    This function load a CAD file in PLY format, or STL format with trimesh:
+         i.e. <trimesh.Trimesh(vertices.shape=(point_cnt, 3), faces.shape=(cnt, 3), name=`ply_path`)>
+
+    Load the raw scan file sampled points in PCD format, then save the updated scan mesh in OBJ format.
+
+    Parameters:
+        - ply_path = os.path.join(root_data_path, "data_pipeline_bar/remesh98.ply")
+        - scan_pcd_path = os.path.join(root_data_path, "data_pipeline_bar/bar_98/scan/98_SAMPLED_POINTS_aligned.pcd")
+        - save_mesh_path = "test_data_pipeline"
+
+    Return:
+        Saved scan mesh path
+    """
+    os.makedirs(save_mesh_path, exist_ok=True)
+
+    # Load cad mesh from PLY file
+    cad_mesh = trimesh.load(ply_path)
+
+    # Centralize the coordinates
+    cad_pts = torch.FloatTensor(np.asarray(cad_mesh.vertices)) - torch.FloatTensor(
+        cad_mesh.bounds.mean(0)
+    )
+
+    # Load raw scan file in PCD, o3d function to read PointCloud from file
+    scan_pts = o3d.io.read_point_cloud(scan_pcd_path)
+
+    # Centralize the coordinates
+    scan_pts = torch.FloatTensor(np.asarray(scan_pts.points)) - torch.FloatTensor(
+        cad_mesh.bounds.mean(0)
+    )
+
+    # Fined one-to-one matching
+    idx1, idx2 = torch_geometric.nn.knn(scan_pts, cad_pts, 1)
+    new_vert = scan_pts[idx2]
+
+    if filter_dist:
+        dist = torch.sqrt(torch.sum(torch.pow(cad_pts - new_vert, 2), 1))
+        filt = dist > 1.2
+        new_vert[filt] = cad_pts[filt]
+
+    # Updates the scan coordinates to the original CAD mesh
+    scan_mesh = cad_mesh
+    vertices = new_vert + torch.FloatTensor(cad_mesh.bounds.mean(0))
+    scan_mesh.vertices = vertices
+
+    if export_format == "obj":
+        scan_mesh.export(os.path.join(save_mesh_path, "data_out.obj"))
+    elif export_format == "pth":
+        torch.save(vertices, os.path.join(save_mesh_path, f"{part_id}/{part_name}.pth"))
+    else:
+        print("Export format should be OBJ or PTH")
+        exit()
+
+    if save_csv:
+        # save the original CAD points with centralize coordinates
+        np.savetxt(
+            os.path.join(save_mesh_path, f"{part_id}/{part_name}_cad.csv"), cad_pts
+        )
+        # save the mapped scan_pts points with centralize coordinates
+        np.savetxt(
+            os.path.join(save_mesh_path, f"{part_id}/{part_name}_scan.csv"), new_vert
+        )
+
+    return os.path.join(save_mesh_path, "data_out.obj")
+
+
+def generate_mesh_eval(cad_path, comp_out_path, export_path, view=False):
+    """
+    Function to load a 3D object pair (Original design file v.s. Scanned printed / Compensated part),
+        - CAD design in format of OBJ or STL
+        - Scanned printed, or compensated part points, in CSV or TXT
+    Export the Scanned in mesh, OBJ format
+
+    Parameters:
+    - object_name = "bar"
+    - part_id = 5
+    - cad_path = "%s_%d/cad/%s_%d_uptess.obj" % (object_name, part_id, object_name, part_id)
+    - comp_out_path = comp/out__%02d.csv" % (part_id)
+
+    Return:
+        Saved scan mesh path
+    """
+    os.makedirs(export_path, exist_ok=True)
+
+    # Sample design CAD name
+    cad_mesh = trimesh.load(cad_path)
+
+    # Sample scanned printed file, or generated compensated file, in CSV or TXT
+    # change the reading format, if data was saved with other separators, " ", ","
+    scan_pts = pd.read_csv(comp_out_path, sep=",").values
+
+    # Define the new vertices as the scanned printed points coordinates
+    new_vert = torch.FloatTensor(scan_pts)
+
+    # Define the mesh from the Design CAD
+    scan_mesh = cad_mesh
+
+    # Export new mesh
+    scan_mesh.vertices = new_vert
+    scan_mesh.export(os.path.join(export_path, "export_out.obj"))
+    if view:
+        scan_mesh.show()