Tracking/SVTOB_GDML/translate_stl.py

0001 #!/usr/bin/env python3
0002
0003 """
0004 translate_stl.py
0005
0006 Translate all points in an STL mesh by a specified displacement
0007 (dx, dy, dz) and write the result to a new STL file.
0008
0009 Usage:
0010     python translate_stl.py input.stl output.stl --dx 10 --dy -5 --dz 2.5
0011
0012 The displacement values use the same coordinate units as the STL file.
0013 """
0014
0015 import argparse
0016 import sys
0017 from pathlib import Path
0018
0019 import numpy as np
0020 import trimesh
0021
0022
0023 def load_mesh(path):
0024     """Load an STL file as a Trimesh."""
0025
0026     mesh = trimesh.load_mesh(
0027         path,
0028         force="mesh",
0029         process=True
0030     )
0031
0032     if not isinstance(mesh, trimesh.Trimesh):
0033         raise ValueError(
0034             "{} did not load as a single mesh.".format(path)
0035         )
0036
0037     return mesh
0038
0039
0040 def translate_mesh(mesh, dx, dy, dz):
0041     """
0042     Translate every vertex by (dx, dy, dz).
0043
0044     Returns a copy of the original mesh.
0045     """
0046
0047     translated = mesh.copy()
0048
0049     displacement = np.array(
0050         [dx, dy, dz],
0051         dtype=float
0052     )
0053
0054     translated.vertices += displacement
0055
0056     return translated
0057
0058
0059 def main():
0060
0061     parser = argparse.ArgumentParser(
0062         description=(
0063             "Translate an STL mesh by a specified "
0064             "X, Y and Z displacement."
0065         )
0066     )
0067
0068     parser.add_argument(
0069         "input",
0070         help="Input STL file"
0071     )
0072
0073     parser.add_argument(
0074         "output",
0075         help="Output STL file"
0076     )
0077
0078     parser.add_argument(
0079         "--dx",
0080         type=float,
0081         default=0.0,
0082         help="Displacement in X (default: 0)"
0083     )
0084
0085     parser.add_argument(
0086         "--dy",
0087         type=float,
0088         default=0.0,
0089         help="Displacement in Y (default: 0)"
0090     )
0091
0092     parser.add_argument(
0093         "--dz",
0094         type=float,
0095         default=0.0,
0096         help="Displacement in Z (default: 0)"
0097     )
0098
0099     args = parser.parse_args()
0100
0101     # Load mesh
0102     try:
0103         mesh = load_mesh(args.input)
0104
0105     except Exception as exc:
0106         print(
0107             "Error loading mesh: {}".format(exc),
0108             file=sys.stderr
0109         )
0110         sys.exit(1)
0111
0112     print("Input mesh:")
0113     print("  vertices:   {}".format(len(mesh.vertices)))
0114     print("  faces:      {}".format(len(mesh.faces)))
0115     print("  watertight: {}".format(mesh.is_watertight))
0116     print("  volume:     {}".format(abs(mesh.volume)))
0117     print(
0118         "  centroid:   ({:.6g}, {:.6g}, {:.6g})".format(
0119             mesh.centroid[0],
0120             mesh.centroid[1],
0121             mesh.centroid[2]
0122         )
0123     )
0124
0125     print()
0126     print("Displacement:")
0127     print("  dx: {}".format(args.dx))
0128     print("  dy: {}".format(args.dy))
0129     print("  dz: {}".format(args.dz))
0130
0131     # Translate mesh
0132     translated = translate_mesh(
0133         mesh,
0134         args.dx,
0135         args.dy,
0136         args.dz
0137     )
0138
0139     print()
0140     print(
0141         "New centroid: ({:.6g}, {:.6g}, {:.6g})".format(
0142             translated.centroid[0],
0143             translated.centroid[1],
0144             translated.centroid[2]
0145         )
0146     )
0147
0148     # Write output
0149     output_path = Path(args.output)
0150
0151     try:
0152         translated.export(str(output_path))
0153
0154     except Exception as exc:
0155         print(
0156             "Error writing output: {}".format(exc),
0157             file=sys.stderr
0158         )
0159         sys.exit(1)
0160
0161     print()
0162     print("Wrote:")
0163     print("  {}".format(output_path))
0164
0165
0166 if __name__ == "__main__":
0167     main()