Start v2

Dockerfile

@@ -2,10 +2,12 @@
 FROM docker.io/fnndsc/mni-conda-base:unofficial
 
 LABEL org.opencontainers.image.authors="Jennings Zhang <[email protected]>" \
-      org.opencontainers.image.title="pl-fetal-cp-surface-extract" \
-      org.opencontainers.image.description="Fetal brain MRI CP surface extraction using CIVET marching-cubes"
+      org.opencontainers.image.title="pl-fetal-surface-extract" \
+      org.opencontainers.image.description="Fetal brain MRI surface extraction using CIVET marching-cubes"
 
-WORKDIR /usr/local/src/pl-fetal-cp-surface-extract
+WORKDIR /usr/local/src/pl-fetal-surface-extract
+
+RUN conda install -c conda-forge -y numpy=1.22.3
 
 COPY requirements.txt .
 RUN pip install -r requirements.txt

README.md

@@ -35,15 +35,30 @@ Moreover, `sphere_mesh` guarantees a spherical topology.
 
 ## Surface Extraction Algorithm
 
-1. Proprocess mask using `mincmorph` to fill in disconnected voxels (improve mask quality)
+1. Preprocess mask using `mincmorph` to fill in disconnected voxels (improve mask quality)
 2. Marching-cubes -> spherical topology surface mesh with unknown number of triangles
 3. Sphere-to-sphere interpolation -> resample mesh to standard connectivity of 81,920 triangles, preserving morphology
-4. A little smoothing using `adapt_object_mesh`.
-
+4. Calculate distance error
+5. _If_ maximum distance error is too large, _then_ redo marching-cubes with subsampling.
+6. Calculate smoothness error
+7. Run `adapt_object_mesh` to achieve desired smoothness. Number of iterations is predicted by a linear model.
+8. Run `mincresample` on the mask with trilinear interpolation to increase mask resolution.
+9. Run `surface_fit` to fully converge surface to mask boundary.
+
+## Pipeline
+
+1. https://github.com/FNNDSC/ep-premc-mincmorph
+2. https://github.com/FNNDSC/pl-nums2mask
+3. https://github.com/FNNDSC/pl-subdiv-minc TODO
+4. pl-fetal-cp-surface-extract
+5. pl-extracted-surface-asp TODO
+
+<!--
 While the upstream
 [marching_cube.pl](https://github.com/aces/surface-extraction/blob/master/scripts/marching_cubes.pl.in)
 script uses ASP (`surface_fit`) post-processing to fully converge the surface to the volume boundary,
 without the extra step the accuracy is nonetheless sufficient.
+-->
 
 ## Installation
 
@@ -68,11 +83,6 @@ individually and in parallel.
 
 ### Options
 
-#### `--subsample`
-
-It is necessary to use the `--subsample` option for fetal brains 29 GA
-and older to avoid "bridge" errors between narrow (<1 voxel) sulcal walls.
-
 #### `--mincmorph-iterations`
 
 Number of `mincmorph` iterations to perform on the mask before marching-cubes.
@@ -84,15 +94,6 @@ _Garbage in, garbage out_.
 extract_cp --mincmorph-iterations 10 /incoming /outgoing
 ```
 
-#### `--adapt_object_mesh`
-
-Arguments to pass to `adapt_object_mesh`, which does mesh smoothing.
-Use a larger value if the results are bumpy/voxelated in appearance.
-
-```shell
-extract_cp --adapt_object_mesh 0,100,0,0 /incoming /outgoing
-```
-
 #### `--inflate_to_sphere_implicit`
 
 Arguments to pass to `inflate_to_sphere_implicit`. The default value `(200, 200)`

extract_cp/__main__.py

@@ -13,27 +13,29 @@
 
 parser = ArgumentParser(description='Fetal brain MRI CP inner surface extraction',
                         formatter_class=ArgumentDefaultsHelpFormatter)
+# -------------------- Input and Output Options --------------------
 parser.add_argument('-s', '--side', default='auto', choices=SIDE_OPTIONS,
                     help='brain hemisphere side. "auto" => infer from file name')
 parser.add_argument('-p', '--pattern', default='**/*.mnc',
                     help='pattern for file names to include')
-parser.add_argument('--subsample', action='store_true',
-                    help='Use -subsample option for spehre_mesh, use with narrow sulci')
-parser.add_argument('--mincmorph-iterations', dest='mincmorph_iterations', type=int, default=5,
-                    help='Number of mincmorph iterations. Mincmorph is a mask preprocessing step '
-                         'which repairs disconnected voxels. A larger value may improve results '
-                         'for messy masks, but at the cost of accuracy.')
-parser.add_argument('--adapt_object_mesh', dest='adapt_object_mesh', type=str, default='0,50,0,0',
-                    help='Parameters for adapt_object_mesh, which does mesh smoothing.')
-parser.add_argument('--inflate_to_sphere_implicit', dest='inflate_to_sphere_implicit', type=str, default='200,200',
-                    help='Parameters for inflate_to_sphere_implicit. Larger values are necessary '
-                         'for larger brain size.')
+parser.add_argument('-t', '--threads', type=int, default=0,
+                    help='number of threads to use (pass 0 to use number of visible CPU cores)')
 parser.add_argument('-k', '--keep-mask', dest='keep_mask', action='store_true',
                     help='Copy input mask file to output directory')
 parser.add_argument('--no-fail', dest='no_fail', action='store_true',
                     help='Exit normally even when failed to process a subject')
 parser.add_argument('-V', '--version', action='version',
-                    version=f'$(prog)s {__version__}')
+                    version=f'%(prog)s {__version__}')
+# --------------------   Algorithm Parameters   --------------------
+parser.add_argument('--inflate_to_sphere_implicit', dest='inflate_to_sphere_implicit', type=str, default='500,500',
+                    help='Parameters for inflate_to_sphere_implicit. Larger values are necessary '
+                         'for larger brain size.')
+parser.add_argument('--distance-threshold', dest='distance_threshold', default=1.0, type=float,
+                    help='Maximum distance error to allow without using subsampling')
+parser.add_argument('--target-smoothness', dest='target_smoothness', type=float, default=0.2,
+                    help='Target mean smoothness error for how many iterations of adapt_object_mesh to perform.')
+parser.add_argument('--max-smooth-iterations', dest='max_smooth_iterations', type=int, default=200,
+                    help='Maximum allowed number of smoothing iterations using adapt_object_mesh.')
 
 
 @chris_plugin(
@@ -49,7 +51,8 @@ def main(options: Namespace, inputdir: Path, outputdir: Path):
     print(DISPLAY_TITLE, file=sys.stderr, flush=True)
     print(params, file=sys.stderr, flush=True)
 
-    with ThreadPoolExecutor(max_workers=len(os.sched_getaffinity(0))) as pool:
+    proc = len(os.sched_getaffinity(0)) if options.threads <= 0 else options.threads
+    with ThreadPoolExecutor(max_workers=proc) as pool:
         mapper = PathMapper.file_mapper(inputdir, outputdir, glob=options.pattern, suffix='._81920.obj')
         results = pool.map(lambda t: extract_surface(*t, params=params), mapper)
 

extract_cp/extract_surface.py

@@ -1,28 +1,36 @@
+import dataclasses
+import json
 import os
+import shlex
 import shutil
 from pathlib import Path
-from typing import Optional, Sequence, BinaryIO
+from typing import Optional, Sequence, BinaryIO, TypedDict
 from civet.extraction.hemisphere import Side, HemisphereMask
 import subprocess as sp
 from loguru import logger
 
+import numpy as np
+import numpy.typing as npt
+
 from extract_cp.params import SideStr, SIDE_OPTIONS, Parameters
+from extract_cp.qc import surfdisterr, smtherr
 
-__log_prefix = b'[' + os.path.basename(__file__).encode(encoding='utf-8') + b']$> '
+_LOG_PREFIX = b'[' + os.path.basename(__file__).encode(encoding='utf-8') + b']'
 
 
 def extract_surface(mask: Path, surface: Path, params: Parameters):
     log_path = surface.with_suffix('.extraction.log')
+    chosen_side = __pick_side(mask, params.side)
     try:
-        chosen_side = __pick_side(mask, params.side)
-        logger.info('Processing {} to {}, log: {}', mask, surface, log_path)
+        logger.info('Surface extraction {} to {} ({}), log: {}', mask, surface, chosen_side, log_path)
         with log_path.open('wb') as log:
-            HemisphereMask(mask)\
-                .smoothen_using_mincmorph(iterations=params.mincmorph_iterations)\
-                .just_sphere_mesh(chosen_side, subsample=params.subsample)\
-                .adapt_object_mesh(*params.adapt_object_mesh)\
-                .interpolate_with_sphere(chosen_side, *params.inflate_to_sphere_implicit)\
-                .save(surface, shell=__curry_log(log))
+            _ExtractSurface(
+                mask=mask,
+                surface=surface,
+                side=chosen_side,
+                params=params,
+                output_sink=log
+            ).run()
         if params.keep_mask:
             shutil.copy(mask, surface.with_name(mask.name))
         logger.info('Completed {}', surface)
@@ -31,6 +39,143 @@ def extract_surface(mask: Path, surface: Path, params: Parameters):
         raise e
 
 
+class StepOutcome(TypedDict):
+    name: str
+    disterr_abs_max: float
+    smtherr_mean: float
+
+
+@dataclasses.dataclass(frozen=True)
+class _ExtractSurface:
+    mask: Path
+    surface: Path
+    side: Side
+    params: Parameters
+    output_sink: BinaryIO
+    outcomes: list[StepOutcome] = dataclasses.field(default_factory=list)
+
+    def run(self):
+        self._marching_cubes_until_disterr_ok()
+        self.conclude()
+
+    def conclude(self):
+        """
+        Finishing touches
+        """
+        self.output_sink.flush()
+        self._write_steps()
+
+    def _marching_cubes_until_disterr_ok(self):
+        """
+        Runs marching-cubes (``sphere_mesh``), saving the surface.
+        The created surface has an inconsistent number of triangles.
+        If maximum absolute distance error exceeds threshold, then
+        retry marching-cubes with subsampling.
+        """
+        self._message('marching-cubes first attempt')
+        self._marching_cubes(subsample=False)
+        first_outcome = self._evaluate('marching-cubes (sphere_mesh)')
+        disterr = first_outcome['disterr_abs_max']
+        if disterr <= self.params.distance_threshold:
+            self._message(f'max(surfdisterr)={disterr:.3f} is ok')
+            return
+
+        self._message(f'max(surfdisterr)={disterr:.3f} is bad')
+        self._message('marching-cubes second attempt **with subsampling**')
+        self._marching_cubes(subsample=True)
+
+        second_outcome = self._evaluate('marching-cubes (sphere_mesh) WITH SUBSAMPLING')
+        second_disterr = second_outcome['disterr_abs_max']
+        percent_change = (second_disterr - disterr) / disterr * 100
+        self._message('marching-cubes with subsampling improvement: '
+                      f'max(surfdisterr) {disterr:.3f} -> {second_disterr:.3f} '
+                      f'({percent_change:.1g}% change)')
+
+        if second_disterr > disterr:
+            self._message('!!!WARNING!!! marching-cubes with subsampling produced a surface '
+                          'with greater distance error, this is unexpected.')
+
+        smth_first = first_outcome['smtherr_mean']
+        smth_second = second_outcome['smtherr_mean']
+        smth_percent_change = (smth_second - smth_first) / smth_first * 100
+        self._message(f'affected mean(smtherr): {smth_first:.3f} -> {smth_second:.3f} '
+                      f'({smth_percent_change:.1g}% change)')
+
+    def _marching_cubes(self, subsample: bool):
+        HemisphereMask(self.mask)\
+            .just_sphere_mesh(self.side, subsample=subsample)\
+            .save(self.surface, shell=self._logged_runner)
+
+    def _evaluate(self, name: str) -> StepOutcome:
+        """
+        (Re)-calculate smtherr and disterr, writing both data to files and append
+        a summary to ``self.outcomes``. This summary is also returned.
+        """
+        outcome = StepOutcome(
+            name=name,
+            smtherr_mean=self._smtherr_mean(),
+            disterr_abs_max=self._disterr_abs_max()
+        )
+        self.outcomes.append(outcome)
+        return outcome
+
+    def _write_surfdisterr(self):
+        surfdisterr(self.surface, self.mask, self._disterr_file)
+
+    def _load_smtherr(self) -> npt.NDArray[np.float32]:
+        data = smtherr(self.surface)
+        np.savetxt(self._smtherr_file, data)
+        return data
+
+    def _load_surfdisterr(self) -> npt.NDArray[np.float32]:
+        self._write_surfdisterr()
+        return np.loadtxt(self._disterr_file, dtype=np.float32)
+
+    def _disterr_abs_max(self) -> float:
+        data = self._load_surfdisterr()
+        return float(np.abs(data).max())
+
+    def _smtherr_mean(self) -> float:
+        data = self._load_smtherr()
+        return float(data.mean())
+
+    def _write_steps(self):
+        with self._outcomes_file.open('w') as f:
+            json.dump(self.outcomes, f, indent=2)
+
+    @property
+    def _disterr_file(self) -> Path:
+        return self.surface.with_suffix('.disterr.txt')
+
+    @property
+    def _smtherr_file(self) -> Path:
+        return self.surface.with_suffix('.smtherr.txt')
+
+    @property
+    def _outcomes_file(self) -> Path:
+        return self.surface.with_suffix('.extraction.steps.json')
+
+    @property
+    def _logged_runner(self):
+        def run_with_log(cmd: Sequence[str | os.PathLike]) -> None:
+            self.output_sink.write(_LOG_PREFIX)
+            self.output_sink.write(b' COMMAND $>')
+            cmdstr = shlex.join(map(str, cmd))
+            self.output_sink.write(cmdstr.encode('utf-8'))
+            self.output_sink.write(b'\n')
+            self.output_sink.flush()
+            sp.run(cmd, stderr=self.output_sink, stdout=self.output_sink, check=True)
+
+        return run_with_log
+
+    def _message(self, *args):
+        self.output_sink.write(_LOG_PREFIX)
+        self.output_sink.write(b' MESSAGE  >>>  ')
+        for msg in map(str, args):
+            self.output_sink.write(msg.encode('utf-8'))
+        self.output_sink.write(b'  <<<\n')
+
+
 def __pick_side(mask: Path, side: SideStr) -> Optional[Side]:
     if side == 'left':
         return Side.LEFT
@@ -46,13 +191,3 @@ def __pick_side(mask: Path, side: SideStr) -> Optional[Side]:
     if side == 'none':
         return None
     raise ValueError(f'side must be one of: {SIDE_OPTIONS}')
-
-
-def __curry_log(log: BinaryIO):
-    def run_with_log(cmd: Sequence[str | os.PathLike]) -> None:
-        log.write(__log_prefix)
-        log.write(str(cmd).encode('utf-8'))
-        log.write(b'\n')
-        log.flush()
-        sp.run(cmd, stderr=log, stdout=log, check=True)
-    return run_with_log

extract_cp/params.py

@@ -8,50 +8,28 @@
 @dataclass(frozen=True)
 class Parameters:
     """
-    Magic constants which set how many iterations should be performed by
-    internal subroutines. These values are primarily dependent on brain size.
+    Provides typing for the expected values of a `argparse.Namespace`.
     """
     side: SideStr
-    mincmorph_iterations: int
-    adapt_object_mesh: tuple[int, int, int, int]
     inflate_to_sphere_implicit: tuple[int, int]
+    distance_threshold: float
+    target_smoothness: float
+    max_smooth_iterations: int
     keep_mask: bool
-    subsample: bool
 
     def __post_init__(self):
-        if not len(self.adapt_object_mesh) == 4:
-            raise ValueError('adapt_object_mesh takes 4 parameters')
         if not len(self.inflate_to_sphere_implicit) == 2:
             raise ValueError('inflate_to_sphere_implicit takes 2 parameters')
 
-    @classmethod
-    def new(
-            cls,
-            side: SideStr,
-            mincmorph_iterations: int,
-            adapt_object_mesh: str,
-            inflate_to_sphere_implicit: str,
-            keep_mask: bool,
-            subsample: bool
-    ) -> 'Parameters':
-        return cls(
-            side=side,
-            mincmorph_iterations=mincmorph_iterations,
-            adapt_object_mesh=_str2ints(adapt_object_mesh),
-            inflate_to_sphere_implicit=_str2ints(inflate_to_sphere_implicit),
-            keep_mask=keep_mask,
-            subsample=subsample
-        )
-
     @classmethod
     def from_obj(cls, options) -> 'Parameters':
-        return cls.new(
+        return cls(
             side=options.side,
-            mincmorph_iterations=options.mincmorph_iterations,
-            adapt_object_mesh=options.adapt_object_mesh,
-            inflate_to_sphere_implicit=options.inflate_to_sphere_implicit,
+            inflate_to_sphere_implicit=_str2ints(options.inflate_to_sphere_implicit),
+            distance_threshold=options.distance_threshold,
+            target_smoothness=options.target_smoothness,
+            max_smooth_iterations=options.max_smooth_iterations,
             keep_mask=options.keep_mask,
-            subsample=options.subsample
         )