Geometry, AI & 3D Printing · Part 05

Mesh Augmentation with PyTorch3D

Good augmentation changes nuisance variables while preserving the meaning of a training example. On meshes, topology, orientation, scale, attributes, and sampling density make that contract explicit.

SimplexGen Research · 16 minute read
A central triangle mesh branching into rotated, scaled, noisy, cropped, remeshed, and point-sampled training variants
Augmentation should produce a family of valid observations—not a family of corrupted labels.

Start With SemanticsWhich Changes Preserve the Label?

Rotation is safe for category recognition but wrong when “upright” or world alignment is the target. Uniform scale is safe after normalization but wrong for physical-size prediction. Mirror reflection may invalidate handed parts, text, chirality, or winding. Topology-changing augmentation is useful for robustness only when labels can be transferred faithfully.

Batch StructureMeshes Are Ragged

Different meshes have different vertex and face counts. PyTorch3D’s Meshes structure exposes list, padded, and packed representations so GPU operations can work without pretending every object has identical topology. Keep the face indices paired with their own vertex tensor after every transform.

CodeA Topology-Preserving Augmenter

import torch
from pytorch3d.structures import Meshes
from pytorch3d.transforms import random_rotations
from pytorch3d.ops import sample_points_from_meshes

def augment_meshes(meshes: Meshes, noise_sigma=0.003):
    device = meshes.device
    rotations = random_rotations(len(meshes), device=device)
    out_verts = []

    for i, (verts, normals) in enumerate(
        zip(meshes.verts_list(), meshes.verts_normals_list())
    ):
        center = verts.mean(0, keepdim=True)
        centered = verts - center

        # Similarity transform: topology and local shape survive.
        scale = torch.empty(1, device=device).uniform_(0.9, 1.1)
        transformed = (centered @ rotations[i].T) * scale

        # Small normal-direction noise perturbs the surface rather than
        # adding unrelated XYZ jitter.
        amplitude = torch.randn(
            len(verts), 1, device=device
        ) * noise_sigma
        transformed = transformed + normals * amplitude
        out_verts.append(transformed)

    augmented = Meshes(verts=out_verts, faces=meshes.faces_list())
    points, normals = sample_points_from_meshes(
        augmented, num_samples=4096, return_normals=True
    )
    return augmented, points, normals

This code intentionally leaves UVs and vertex attributes out. In a real pipeline, clone or transform every aligned attribute, and rotate vector labels such as normals, axes, or keypoint frames with the geometry.

Surface SamplingDo Not Confuse Vertices With Shape

Vertices reflect tessellation decisions. Uniform surface samples reflect triangle area and are usually a better input for Chamfer loss, point encoders, or cross-topology comparison. Resample after the geometric transform and record the RNG seed when exact reproducibility matters.

Augmentation FamiliesUse Different Risks Deliberately

  • Rigid: rotation and translation; preserves all intrinsic geometry.
  • Similarity: rigid plus uniform scale; preserves angles and normalized shape.
  • Affine: anisotropic scale and shear; useful only if shape deformation is a nuisance variable.
  • Surface: normal noise, smooth displacement, spectral perturbation; can teach scan robustness.
  • Sampling: point resampling, face dropout, partial views; changes observation rather than the source object.
  • Topology: decimation, subdivision, remeshing; powerful but requires attribute and label transfer.

Losses as GuardsReject Bad Augmentations

PyTorch3D exposes Chamfer distance, edge loss, Laplacian smoothing, and normal-consistency losses. Use them as acceptance checks, not just training objectives. Reject a proposed variant if its surface drift, inverted-face count, component count, or semantic landmark error exceeds your contract.

from pytorch3d.loss import (
    chamfer_distance,
    mesh_edge_loss,
    mesh_normal_consistency,
)

surface_error, _ = chamfer_distance(points_aug, points_source)
quality = (
    surface_error
    + 0.05 * mesh_edge_loss(augmented)
    + 0.02 * mesh_normal_consistency(augmented)
)

Dataset HygieneSplit Before You Augment

All variants of one source belong to one dataset split. If a remeshed or rotated copy leaks into validation, metrics measure source recognition rather than generalization. Hash the source identity and derive deterministic augmentation seeds from that identity plus epoch.

Primary references

← Part 04Part 06 →