ENH,DOC allow passing customizations, add docstrings (#16)

* ENH,DOC allow passing customizations, add docstrings

* Be verbose with blender params

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,
+        )