This code uses generative deep learning models to understand the skeletal anatomy of lizards and some snakes (Squamata). This code is forked and modified from gattia/NSM following the terms of the GNU Affero GPL 3.0 License. See Original NSM Documentation.
Figure 1: Traversing an isomap of the NSM trained model latent space using travelling salesman and k-nearest neighbors. Video animation made using isomap_video.py
# Create and activate conda environment
conda create -n NSM python=3.10
conda activate NSM
# Install pytorch and dependencies
conda install pytorch=2.5.1 torchvision=0.20.1 torchaudio=2.5.1 pytorch-cuda=12.4 -c pytorch -c nvidia -c conda-forge -c defaults
# Install NSM
mkdir NSM
cd NSM
git clone https://github.com/aubricot/nsm.git
cd nsm
python -m pip install -r requirements.txt
pip install .
Update sections of train_model.py commented with # TO DO: to update PROJECT_NAME, ENTITY_NAME, RUN_NAME, folder_vtk, N_TRAIN, N_TEST, N_VAL. These variables point to where your data was collected, where it is saved, and where outputs should go. Adjust model training hyperparameters in vertebrae_config.json. See python script and config files for details and save before running using commands below.
conda activate NSM
cd NSM/nsm
python train_model.py
NSM provides a convenient model loader that simplifies loading pre-trained Neural Shape Models. For trained models, you'll have:
experiment_dir/model_params_config.json- Configuration saved during trainingexperiment_dir/model/2000.pth- Model weights at epoch 2000experiment_dir/latent_codes/2000.pth- Latent codes at epoch 2000
import json, torch
from NSM.models import TriplanarDecoder
# Load config file
with open(config_path, 'r') as f:
config = json.load(f)
# Get model weights and latent codes
latent_ckpt = torch.load(LC_PATH, map_location=device)
latent_codes = latent_ckpt['latent_codes']['weight'].detach().cpu()
# Build model
triplane_args = {
'latent_dim': config['latent_size'],
'n_objects': config['objects_per_decoder'],
'conv_hidden_dims': config['conv_hidden_dims'],
'conv_deep_image_size': config['conv_deep_image_size'],
'conv_norm': config['conv_norm'],
'conv_norm_type': config['conv_norm_type'],
'conv_start_with_mlp': config['conv_start_with_mlp'],
'sdf_latent_size': config['sdf_latent_size'],
'sdf_hidden_dims': config['sdf_hidden_dims'],
'sdf_weight_norm': config['weight_norm'],
'sdf_final_activation': config['final_activation'],
'sdf_activation': config['activation'],
'sdf_dropout_prob': config['dropout_prob'],
'sum_sdf_features': config['sum_conv_output_features'],
'conv_pred_sdf': config['conv_pred_sdf'],
}
model = TriplanarDecoder(**triplane_args)
model_ckpt = torch.load(MODEL_PATH, map_location=device)
model.load_state_dict(model_ckpt['model'])
device = config.get("device", "cuda:0")
model.to(device)
model.eval()After loading a trained model, you can generate meshes from manipulated/new latent vectors. The example below generates the mean mesh shape based on model training data.
from NSM.mesh import create_mesh
import pyvista as pv
# Get the mean of the latent codes
latents_np = latent_codes.numpy()
latent_mean = np.mean(latents_np, axis=0)
# Convert the mean latent code to a pytorch tensor
new_latent = torch.tensor(new_latent_np, dtype=torch.float32).unsqueeze(0).to(device)
# Create a mesh from the latent tensor
mesh_out = create_mesh(
decoder=model, latent_vector=new_latent, n_pts_per_axis=n_pts_per_axis,
voxel_origin=voxel_origin, voxel_size=voxel_size, path_original_mesh=None,
offset=offset, scale=scale, icp_transform=icp_transform,
objects=objects, verbose=False, device=device
)
# Ensure mesh is PyVista Polydata (.vtk)
if isinstance(mesh_out, list):
mesh_out = mesh_out[0]
if not isinstance(mesh_out, pv.PolyData):
mesh_pv = mesh_out.extract_geometry()
else:
mesh_pv = mesh_out
# Write to file
mesh_pv.save(output_path)This code is forked and modified from https://github.com/gattia/NSM following the terms of the GNU Affero GPL 3.0 License and NSM License.