stl_orienter.py

#!/usr/bin/env python3

from argparse import ArgumentParser

parser = ArgumentParser(description='Orient STL files')
parser.add_argument('--path', help='Directory containing STL files or single STL file',default='.')
parser.add_argument('--recursive', action='store_true', help='Recursively search for STL files')
parser.add_argument('--output', help='Output directory, if none is given, the original files will be overwritten')
parser.add_argument('--param-dir', help='Directory containing the parameter files',default='.orient')
parser.add_argument('--no-apply', help='Skip applying the parameters', action='store_true')
parser.add_argument('--no-param', help='Skip creating new parameters', action='store_true')

parser.add_argument('--openscad', help='Path to openscad', default='openscad')
args = parser.parse_args()

template = lambda absolute_filename,offset=(0,0,0): f"""
echo("{absolute_filename}");
/* [ROTATION] */
RX = 0; //[-180:45:180]
RY = 0; //[-180:45:180]
RZ = 0; //[-180:45:180]
/* [TRANSLATION] */
TX = {offset[0]};
TY = {offset[1]};
TZ = {offset[2]};
rotate([RX,RY,RZ])
translate([TX,TY,TZ])
import("{absolute_filename}");
"""

import os.path

# region Get files
files = []
if os.path.isdir(args.path):
    if args.recursive:
        for root, dirs, fs in os.walk(args.path):
            for f in fs:
                if f.endswith('.stl'):
                    files.append(os.path.join(root, f))
    else:
        for f in os.listdir(args.path):
            if f.endswith('.stl'):
                files.append(os.path.join(args.path, f))
else:
    files.append(args.path)
print(f'Found {len(files)} STL files')
# endregion

# region Look for parameter files

file_to_param = {}
for file in files:
    # if file exists in the param dir under the same tree structure, use it
    param_file = os.path.join(args.param_dir, file)+'.json'
    if os.path.isfile(param_file):
        file_to_param[file] = param_file

from more_itertools import ilen
print(f'Found {ilen(file_to_param.values())} parameter files')
# endregion

# region Create missing parameter files

from stl import mesh
import gc

if not args.no_param:
    print('Creating missing parameter files')
    for file in files:
        if file not in file_to_param:
            print(f'Creating parameter file for "{file}"')
            # get the model offset
            # TODO: scan the file for the bounding box instead of loading the whole file
            m = mesh.Mesh.from_file(os.path.abspath(file))
            offset = tuple(-(m.min_+m.max_)/2)
            del m
            gc.collect()

            template_filename = os.path.join(args.param_dir, file)+'.scad'
            # build the directory structure
            os.makedirs(os.path.dirname(template_filename), exist_ok=True)
            # create the template file
            with open(template_filename, 'w') as f:
                f.write(template(os.path.abspath(file),offset))
            # call openscad
            os.system(f'{args.openscad} "{os.path.abspath(template_filename)}" --viewall 2>/dev/null')
            # remove the template file
            os.remove(template_filename)

            # look for the params file
            file_to_param[file] = os.path.join(args.param_dir, file)+'.json'
            if os.path.isfile(file_to_param[file]):
                print(f'Found parameter file for "{file}" @ "{file_to_param[file]}"')
            else:
                print(f'Error: Could not find the params file for "{file}" @ "{file_to_param[file]}"')
                file_to_param[file] = None
# endregion

# region Apply orientation

if not args.no_apply:
    print('Applying orientation')
    for file, param in file_to_param.items():
        print(f'Try {file}')
        template_filename = os.path.join(args.param_dir, file)+'.scad'
        # we assume the directory structure is already created
        # create the template file
        with open(template_filename, 'w') as f:
            f.write(template(os.path.abspath(file)))
        # call openscad
        if args.output:
            output_filename = os.path.join(args.output, file)
            os.makedirs(os.path.dirname(output_filename), exist_ok=True)
        else:
            output_filename = file
        os.system(f'{args.openscad} -p "{file_to_param[file]}" -P "design default values" -o "{output_filename}" "{template_filename}" --viewall 2>/dev/null')
        # remove the template file
        os.remove(template_filename)

        if args.output:
            if os.path.isfile(output_filename):
                print(f'Found output file for "{file}" @ "{output_filename}"')
            else:
                print(f'Error: Could not find the output file for "{file}" @ "{output_filename}"')
        else:
            print(f'Assuming "{output_filename}" was overwritten')