Created
April 17, 2023 14:40
-
-
Save NicerNewerCar/8bf17e0674a9e2f62eed5141756aed23 to your computer and use it in GitHub Desktop.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| from stl import mesh | |
| import os, math | |
| import click | |
| import pandas as pd | |
| def read_stl_file(filename: str) -> mesh: | |
| """Read in an stl file and return a mesh object""" | |
| stl_mesh = mesh.Mesh.from_file(filename) | |
| return stl_mesh | |
| def save_stl_file(mesh: mesh, filename: str) -> None: | |
| """Save the mesh object to a new stl file""" | |
| mesh.save(filename) | |
| def read_tra_file(filename: str) -> pd.DataFrame: | |
| """Read in the tra file into a dataframe""" | |
| df = pd.read_csv(filename, sep=",", header=None) | |
| # if there is a 7th column, remove it | |
| if len(df.columns) == 7: | |
| df = df.drop(6, axis=1) | |
| # remove any Nan values | |
| df = df.dropna() | |
| # add the column names x y z yaw pitch roll | |
| df.columns = ["x", "y", "z", "yaw", "pitch", "roll"] | |
| return df | |
| def apply_tra(mesh: mesh, tra: pd.DataFrame, output_dir: str) -> None: | |
| """Apply the positions to the mesh and save the new mesh to a file""" | |
| # get the initial position from the dataframe | |
| x0 = tra.iloc[0]["x"] | |
| y0 = tra.iloc[0]["y"] | |
| z0 = tra.iloc[0]["z"] | |
| roll0 = tra.iloc[0]["roll"] | |
| pitch0 = tra.iloc[0]["pitch"] | |
| yaw0 = tra.iloc[0]["yaw"] | |
| # move the mesh to the initial position | |
| mesh.translate([x0, y0, z0]) | |
| mesh.rotate([roll0, pitch0, yaw0], math.radians(1)) | |
| # save the initial position | |
| save_stl_file(mesh, os.path.join(output_dir, "0.stl")) | |
| # move the mesh to the rest of the positions | |
| for i in range(1, len(tra)): | |
| x = tra.iloc[i]["x"] | |
| y = tra.iloc[i]["y"] | |
| z = tra.iloc[i]["z"] | |
| roll = tra.iloc[i]["roll"] | |
| pitch = tra.iloc[i]["pitch"] | |
| yaw = tra.iloc[i]["yaw"] | |
| mesh.translate([x - x0, y - y0, z - z0]) | |
| mesh.rotate([roll - roll0, pitch - pitch0, yaw - yaw0], math.radians(1)) | |
| save_stl_file(mesh, os.path.join(output_dir, str(i) + ".stl")) | |
| x0 = x | |
| y0 = y | |
| z0 = z | |
| roll0 = roll | |
| pitch0 = pitch | |
| yaw0 = yaw | |
| @click.command() | |
| @click.option("--stl", help="The stl file to transform") | |
| @click.option("--tra", help="The tra file to apply") | |
| @click.option("--output_dir", help="The directory to save the transformed stl files") | |
| def main(stl, tra, output_dir): | |
| if stl is None: | |
| print("Please provide an stl file") | |
| return | |
| if tra is None: | |
| print("Please provide a tra file") | |
| return | |
| if output_dir is None: | |
| print("Please provide an output directory") | |
| return | |
| if not os.path.exists(output_dir): | |
| os.makedirs(output_dir) | |
| print("Reading stl file: {}".format(stl)) | |
| mesh = read_stl_file(stl) | |
| print("Reading tra file: {}".format(tra)) | |
| tra = read_tra_file(tra) | |
| print("Applying tra file to stl file") | |
| apply_tra(mesh, tra, output_dir) | |
| print("Done") |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment