diff --git a/make_headcase.py b/make_headcase.py index 28aa694..62e09d7 100644 --- a/make_headcase.py +++ b/make_headcase.py @@ -14,6 +14,9 @@ from blender_code import blender_carve_model_template +cwd, _ = os.path.split(__file__) +DEFAULT_CUSTOMIZATIONS = os.path.join(cwd, "stls", "default_customizations.stl") + def _call_blender(code): """Call blender, while running the given code. If the filename doesn't exist, @@ -29,6 +32,7 @@ def _call_blender(code): def meshlab_filter(ms): + """Apply mesh filters to clean and process a 3D model using PyMeshLab.""" # "Transform: Move, Translate, Center" ms.apply_filter(filter_name="compute_matrix_from_translation") # "Transform: Rotate" @@ -86,6 +90,7 @@ def meshlab_filter(ms): def meshlab_filter_pre2022(ms): + """Apply mesh filters to clean and process a 3D model using PyMeshLab (pre-2022).""" # "Transform: Move, Translate, Center" ms.apply_filter(filter_name="transform_translate_center_set_origin") # "Transform: Rotate" @@ -139,6 +144,35 @@ def meshlab_filter_pre2022(ms): def model_clean(infile, outfile): + """ + Clean and process a 3D model file. + + Parameters + ---------- + infile : str + The input file path of the 3D model. It can be a zip file containing the model + or a direct path to the model file. + outfile : str + The output file path to save the cleaned and processed 3D model. + + Notes + ----- + This function cleans and processes a 3D model file using PyMeshLab. + If the input file is a zip file, it will be extracted to a temporary directory + before processing. The cleaned and processed model will be saved to the specified + output file. + + The function checks the version of PyMeshLab installed and applies the appropriate + mesh filters accordingly. + + If the input file is a zip file, the temporary directory will be deleted after + processing. + + Examples + -------- + >>> model_clean("input.zip", "output.ply") + """ + clean_tmp = False if infile.endswith("zip"): path = mkdtemp() @@ -165,6 +199,16 @@ def model_clean(infile, outfile): def align_scan(infile, outfile): + """ + Automatically aligns a head scan and saves the aligned scan as an STL file. + + Parameters + ---------- + infile : str + The path to the input scan file. + outfile : str + The path to save the aligned scan as an STL file. + """ from cortex import formats from autocase3d.fmin_autograd import fit_xfm_autograd @@ -175,10 +219,45 @@ def align_scan(infile, outfile): def gen_case( - scanfile, outfile, workdir=None, casetype="s32", nparts=4, expand_head_model=0.1 + scanfile, + outfile, + workdir=None, + casetype="s32", + nparts=4, + expand_head_model=0.1, + customizations=DEFAULT_CUSTOMIZATIONS, ): - cwd, _ = os.path.split(__file__) - customizations = os.path.join(cwd, "stls", "default_customizations.stl") + """ + Generate a headcase. + + Parameters + ---------- + scanfile : str + Path to the cleaned and aligned head model. + outfile : str + Path to the output file where the generated head case will be saved. + workdir : str, optional + Path to the working directory. If not provided, a temporary directory will be + created and deleted after processing. + casetype : str, optional + Type of head case to generate. + Possible values are 's32', 's64', 'n32', 'meg_ctf275'. + Default is 's32'. + nparts : int, optional + Number of parts to divide the head case into. Possible values are 2 or 4. + Default is 4. + expand_head_model : float, optional + Factor (in mm) to expand the head model by. Default is 0.1. + customizations : str, optional + Path to the customizations file to remove additional parts from the headcase. + Default is `default_customizations.stl` in the `stls` folder. + + Examples + -------- + >>> gen_case("02aligned.stl", "head_case.zip", casetype="s64", nparts=2) + """ + + customizations = os.path.abspath(customizations) casefile = dict( s32="s32.stl", s64="s64.stl", n32="n32.stl", meg_ctf275="meg_ctf275.stl" ) @@ -189,16 +268,17 @@ def gen_case( workdir = mkdtemp() cleanup = True - _call_blender( - blender_carve_model_template.format( - preview=casefile, - scan=scanfile, - customizations=customizations, - tempdir=workdir, - nparts=nparts, - shrinking_factor=expand_head_model, - ) + blender_params = dict( + preview=casefile, + scan=scanfile, + customizations=customizations, + tempdir=workdir, + nparts=nparts, + shrinking_factor=expand_head_model, ) + print("Generating head model by calling Blender with the following parameters:") + print(blender_params) + _call_blender(blender_carve_model_template.format(**blender_params)) pieces = { 2: ["back.stl", "front.stl"], @@ -213,6 +293,7 @@ def gen_case( def pymeshlab_version(): + """Return the version of PyMeshLab installed.""" out = sp.check_output( [ "python", @@ -225,7 +306,48 @@ def pymeshlab_version(): return version -def pipeline(infile, outfile, casetype="s32", nparts=4, workdir=None): +def pipeline( + infile, + outfile, + casetype="s32", + nparts=4, + workdir=None, + customizations=DEFAULT_CUSTOMIZATIONS, +): + """ + Run the pipeline to generate a head case from a head model. + + Parameters + ---------- + infile : str + Path to the input file containing the head model. It can either be a zip file + generated by the Structure Sensor or an obj file containing the head model. + outfile : str + Path to the output file containing the generated head case. + casetype : str, optional + Type of head case, default is "s32" (Siemens 32ch). Possible values are + "s32", "s64", "n32", "meg_ctf275". + nparts : int, optional + Number of parts, default is 4. + workdir : str, optional + Path to the working directory, default is None. + customizations : dict, optional + Customizations for the head case, default is `default_customizations.stl`. + + Notes + ----- + This function performs the following steps: + 1. If `workdir` is provided, creates the working directory if it does not exist. + 2. Cleans the head model by calling `model_clean` function. + 3. Aligns the cleaned head model by calling `align_scan` function. + 4. Generates the head case by calling `gen_case` function. + 5. If `workdir` is not provided, removes the working directory. + + Examples + -------- + >>> pipeline("Model.zip", "Headcase.zip", casetype="s32", nparts=4) + """ + if workdir is not None: working_dir = os.path.abspath(workdir) os.makedirs(working_dir, exist_ok=True) @@ -239,7 +361,14 @@ def pipeline(infile, outfile, casetype="s32", nparts=4, workdir=None): print("Aligning head model") align_scan(cleaned, aligned) print("Making head case") - gen_case(aligned, outfile, working_dir, casetype=casetype, nparts=nparts) + gen_case( + aligned, + outfile, + working_dir, + casetype=casetype, + nparts=nparts, + customizations=customizations, + ) if workdir is None: shutil.rmtree(working_dir) @@ -293,6 +422,15 @@ def pipeline(infile, outfile, casetype="s32", nparts=4, workdir=None): help="Only generate the headcase given the input stl file. This assumes that " "the input stl contains a head model that is already cleaned and aligned.", ) + parser.add_argument( + "--customizations-file", + type=str, + default=DEFAULT_CUSTOMIZATIONS, + help="File containing additional shapes that will be removed from the headcase " + "after carving out the head model. For example, this file is used to carve out " + "space near the ears and the nose bridge. This customization file can be edited " + f"to fine-tune the headcase. The default file is {DEFAULT_CUSTOMIZATIONS}", + ) parser.add_argument( "--workdir", type=str, @@ -311,10 +449,24 @@ def pipeline(infile, outfile, casetype="s32", nparts=4, workdir=None): casetype = args.headcoil nparts = args.nparts workdir = args.workdir + customizations = args.customizations_file generate_headcase_only = args.generate_headcase_only if generate_headcase_only: print("Making head case") - gen_case(infile, outfile, casetype=casetype, nparts=nparts) + gen_case( + infile, + outfile, + casetype=casetype, + nparts=nparts, + customizations=customizations, + ) else: - pipeline(infile, outfile, casetype=casetype, nparts=nparts, workdir=workdir) + pipeline( + infile, + outfile, + casetype=casetype, + nparts=nparts, + workdir=workdir, + customizations=customizations, + )