From bd4407f72b2034da2e966cedb9ce011ebab20e00 Mon Sep 17 00:00:00 2001
From: Matteo Visconti di Oleggio Castello <mvdoc@berkeley.edu>
Date: Sat, 20 Apr 2024 16:17:42 -0700
Subject: [PATCH] ENH allow passing shrinking factor, store blender script in
 workdir

---
 blender_code.py  |  2 +-
 make_headcase.py | 30 ++++++++++++++++++++++++------
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/blender_code.py b/blender_code.py
index ec20a49..1af087b 100644
--- a/blender_code.py
+++ b/blender_code.py
@@ -40,7 +40,7 @@ def readstl(path, name):
 bpy.data.objects['scan'].modifiers['Displace'].name = 'shrinking'
 bpy.data.objects['scan'].modifiers['shrinking'].direction = 'NORMAL'
 bpy.data.objects['scan'].modifiers['shrinking'].mid_level = 0
-bpy.data.objects['scan'].modifiers['shrinking'].strength = 0.1
+bpy.data.objects['scan'].modifiers['shrinking'].strength = {shrinking_factor}
 
 try:
     readstl('{customizations}', 'customizations')
diff --git a/make_headcase.py b/make_headcase.py
index b7ba44e..28aa694 100644
--- a/make_headcase.py
+++ b/make_headcase.py
@@ -1,4 +1,5 @@
 """Generate an MRI-compatible headcase from a 3D head model acquired with a Structure Sensor."""
+
 import argparse
 import os
 import shlex
@@ -173,7 +174,9 @@ def align_scan(infile, outfile):
     formats.write_stl(outfile, new_pts, new_polys)
 
 
-def gen_case(scanfile, outfile, casetype="s32", nparts=4):
+def gen_case(
+    scanfile, outfile, workdir=None, casetype="s32", nparts=4, expand_head_model=0.1
+):
     cwd, _ = os.path.split(__file__)
     customizations = os.path.join(cwd, "stls", "default_customizations.stl")
     casefile = dict(
@@ -181,14 +184,19 @@ def gen_case(scanfile, outfile, casetype="s32", nparts=4):
     )
     casefile = os.path.join(cwd, "stls", casefile[casetype])
 
-    tempdir = mkdtemp()
+    cleanup = False
+    if workdir is None:
+        workdir = mkdtemp()
+        cleanup = True
+
     _call_blender(
         blender_carve_model_template.format(
             preview=casefile,
             scan=scanfile,
             customizations=customizations,
-            tempdir=tempdir,
+            tempdir=workdir,
             nparts=nparts,
+            shrinking_factor=expand_head_model,
         )
     )
 
@@ -198,9 +206,10 @@ def gen_case(scanfile, outfile, casetype="s32", nparts=4):
     }
     with zipfile.ZipFile(outfile, mode="w") as pkg:
         for fn in pieces[nparts]:
-            pkg.write(os.path.join(tempdir, fn), fn)
+            pkg.write(os.path.join(workdir, fn), fn)
 
-    shutil.rmtree(tempdir)
+    if cleanup:
+        shutil.rmtree(workdir)
 
 
 def pymeshlab_version():
@@ -230,7 +239,7 @@ 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, casetype=casetype, nparts=nparts)
+    gen_case(aligned, outfile, working_dir, casetype=casetype, nparts=nparts)
 
     if workdir is None:
         shutil.rmtree(working_dir)
@@ -250,6 +259,15 @@ def pipeline(infile, outfile, casetype="s32", nparts=4, workdir=None):
         type=str,
         help="output headcase model (*.zip)",
     )
+    parser.add_argument(
+        "--expand-head-model",
+        type=float,
+        default="0.1",
+        help="Expand the head model by this amount (in mm) before generating the "
+        "headcase. The default (0.1 mm) should work for most cases. If the resulting "
+        "headcase is too tight, one can try increasing this value. It is not "
+        "recommended to pass a value greater than 1 mm.",
+    )
     parser.add_argument(
         "--headcoil",
         "-c",