Modifications to support infant pipeline and building samseg atlases (f…

…reesurfer#924)

* [FIX] making infant stream compatible with latest niftireg (1.3.9)

* [WIP] adding binaries needed for infant stream

* [WIP] more robust cmake setings

* [WIP] cmake support for INFANT_MODULE AND MINIMAL

* [WIP] reverting niftyreg tweak now that version is correct

* [WIP] adding infant stream python reqs

* [WIP] neurodocker entrypoint

* [FIX] entrypoint

* [FIX] entrypoint perm fix

* [WIP] updating infant py reqs

* [WIP] install infant entrypoint

* [FIX] fs dir structure

* [FIX] harden against top-level files in subject dir

* [ENH] to support samseg vis

* [WIP] inital attempt at recalc script

* [WIP] cleanup and argparse

* [FIX] std:: namespace ref

* [WIP] helpful debug msg

* [WIP] comments and cleanup

* [TXT] helpful debug/processing msg

* [FIX] semicolons

* [WIP] infant entrypoint for aws initial test

* [WIP] match fs env vars and infant input fname reqs

* [WIP] installing infant entrypoint script for aws

* [ENH] New env var: SSCNN_MODEL_DIR

Used to overide the default location of sscnn_skullstrip model files

* [FIX] indent

* [WIP] exposing params to comand line

* [WIP] syntax

* [WIP] tweaks to support additional command line params

* [WIP] connecting numIterations and edgeCollapseFactor params

* [WIP] syntax

* [ENH] making `prepareAtlasDirectory.py` callable from command line

* [GIT] removing `recompute_atlas_probs` (superceded by `gems_compute_atlas_probs`)

gems/Executables/kvlAtlasMeshBuilder.cxx

@@ -88,14 +88,15 @@ AtlasMeshBuilder
 ::SetUp( const std::vector< LabelImageType::ConstPointer >& labelImages,
          const CompressionLookupTable*  compressionLookupTable,
          const itk::Size< 3>&  initialSize, 
-         const std::vector< double >& initialStiffnesses )
+         const std::vector< double >& initialStiffnesses,
+         const unsigned int maximumNumberOfIterations)
 {
   m_LabelImages = labelImages;
   m_CompressionLookupTable = compressionLookupTable;
   m_InitialSize = initialSize;
   m_InitialStiffnesses = initialStiffnesses;
   m_Mesher->SetUp( m_LabelImages, m_CompressionLookupTable, m_InitialSize, m_InitialStiffnesses );
-
+  m_MaximumNumberOfIterations = maximumNumberOfIterations;
 }
 
 

gems/Executables/kvlAtlasMeshBuilder.h

@@ -171,7 +171,8 @@ public :
   void SetUp( const std::vector< LabelImageType::ConstPointer >& labelImages,
               const CompressionLookupTable*  compressionLookupTable,
               const itk::Size< 3>&  initialSize, 
-              const std::vector< double >& initialStiffnesses );
+              const std::vector< double >& initialStiffnesses,
+              const unsigned int maximumNumberOfIterations );
 
   // Get label images
   const std::vector< LabelImageType::ConstPointer >&  GetLabelImages() const

gems/Executables/kvlBuildAtlasMesh.cxx

@@ -115,7 +115,7 @@ int main( int argc, char** argv )
   // Sanity check on input
   if ( argc < 8 )
     {
-    std::cerr << "Usage: " << argv[ 0 ] << " numberOfUpsamplingSteps meshSizeX meshSizeY meshSizeZ stiffness logDirectory fileName1 [ fileName2 ... ]" << std::endl;
+    std::cerr << "Usage: " << argv[ 0 ] << " numberOfUpsamplingSteps meshSizeX meshSizeY meshSizeZ stiffness numberOfIterations edgeCollapseFactor logDirectory fileName1 [ fileName2 ... ]" << std::endl;
 
     return -1;
     }
@@ -126,7 +126,7 @@ int main( int argc, char** argv )
 
   // Retrieve the input parameters
   std::ostringstream  inputParserStream;
-  for ( int argumentNumber = 1; argumentNumber < 7; argumentNumber++ ) 
+  for ( int argumentNumber = 1; argumentNumber < 9; argumentNumber++ ) 
     {
     inputParserStream << argv[ argumentNumber ] << " ";
     }
@@ -136,16 +136,33 @@ int main( int argc, char** argv )
   unsigned int  meshSizeY;
   unsigned int  meshSizeZ;
   double  stiffness;
+  unsigned int  numberOfIterations;
+  double  edgeCollapseEncouragmentFactor;
   std::string  logDirectory;
-  inputStream >> numberOfUpsamplingSteps >>  meshSizeX >> meshSizeY >> meshSizeZ >> stiffness >> logDirectory;
-
+  inputStream >> \
+    numberOfUpsamplingSteps >> \
+    meshSizeX >> meshSizeY >> meshSizeZ >> \
+    stiffness >> \
+    numberOfIterations >> \
+    edgeCollapseEncouragmentFactor >> \
+    logDirectory;
+
+  std::cout << "kvlBuildAtlasMesh Command line params:" << std::endl;
+  std::cout << "  numberOfUpsamplingSteps:        " << numberOfUpsamplingSteps << std::endl;
+  std::cout << "  meshSizeX:                      " << meshSizeX << std::endl;
+  std::cout << "  meshSizeY:                      " << meshSizeY << std::endl;
+  std::cout << "  meshSizeZ:                      " << meshSizeZ << std::endl;
+  std::cout << "  stiffness:                      " << stiffness << std::endl;
+  std::cout << "  numberOfIterations:             " << numberOfIterations << std::endl;
+  std::cout << "  edgeCollapseEncouragmentFactor: " << edgeCollapseEncouragmentFactor << std::endl;
+  std::cout << "  logDirectory:                   " << logDirectory << std::endl;
 
   // Read the input images
   typedef kvl::CompressionLookupTable::ImageType  LabelImageType;
   std::vector< LabelImageType::ConstPointer >  labelImages;
-  for ( int argumentNumber = 7; argumentNumber < argc; argumentNumber++ )
+  for ( int argumentNumber = 9; argumentNumber < argc; argumentNumber++ )
     {
-    std::cout << "Reading input image: " << argv[ argumentNumber ] << std::endl
+    std::cout << "Reading input image: " << argv[ argumentNumber ] << std::endl;
     // Read the input image
     typedef itk::ImageFileReader< LabelImageType >  ReaderType;
     ReaderType::Pointer  reader = ReaderType::New();
@@ -174,7 +191,7 @@ int main( int argc, char** argv )
   kvl::AtlasMeshBuilder::Pointer  builder = kvl::AtlasMeshBuilder::New();
   const itk::Size< 3 >  initialSize = { meshSizeX,  meshSizeY, meshSizeZ };
   std::vector< double >  initialStiffnesses( numberOfUpsamplingSteps+1, stiffness );
-  builder->SetUp( labelImages, lookupTable, initialSize, initialStiffnesses );
+  builder->SetUp( labelImages, lookupTable, initialSize, initialStiffnesses, numberOfIterations);
   builder->SetVerbose( false );
 
   // Add some observers/callbacks
@@ -252,11 +269,15 @@ int main( int argc, char** argv )
       {
       std::cerr << "Couldn't read mesh from file " << explicitStartCollectionFileName << std::endl;
       return -1;
+      } 
+    else
+      {
+      std::cout << "explicitStartCollection found; reading from: " << explicitStartCollectionFileName << std::endl;
       }
     }
 
   // If edgeCollapseEncouragmentFactor.txt exists in the current directory, read it's content
-  double  edgeCollapseEncouragmentFactor = 1.0;
+  //double  edgeCollapseEncouragmentFactor = 1.0;
   const std::string  edgeCollapseEncouragmentFactorFileName = "edgeCollapseEncouragmentFactor.txt";
   //if ( itksys::SystemTools::FileExists( edgeCollapseEncouragmentFactorFileName.c_str(), true ) )
   //  {

infant/CMakeLists.txt

@@ -32,6 +32,7 @@ add_subdirectory(labelfusion)
 
 # Entrypoint for containers
 install(PROGRAMS docker/infant-container-entrypoint.bash DESTINATION bin)
+install(PROGRAMS docker/infant-container-entrypoint-aws.bash DESTINATION bin)
 
 # install external niftyreg binaries
 if(MARTINOS_BUILD)

infant/docker/infant-container-entrypoint-aws.bash

@@ -0,0 +1,55 @@
+#!/bin/bash
+
+# TODO: remove `FS_INFANT_MODEL` (no longer needed?)
+echo "==============================================================="
+echo "ENVIRONMENT VARIABLES"
+echo ""
+echo "AWS_BATCH_JOB_ID:                  $AWS_BATCH_JOB_ID"
+echo "AWS_BATCH_JQ_NAME:                 $AWS_BATCH_JQ_NAME"
+echo "AWS_BATCH_CE_NAME:                 $AWS_BATCH_CE_NAME"
+echo "---------------------------------------------------------------"
+echo "FREESURFER_HOME:                   $FREESURFER_HOME"
+echo "FS_INFANT_MODEL:                   $FS_INFANT_MODEL"
+echo "SUBJECTS_DIR:                      $SUBJECTS_DIR"
+echo "FS_SUB_NAME:                       $FS_SUB_NAME"
+echo "SSCNN_MODEL_DIR:                   $SSCNN_MODEL_DIR"
+echo "---------------------------------------------------------------"
+echo "FS_NIFTI_INPUT_S3_FILEPATH:        $FS_NIFTI_INPUT_S3_FILEPATH"
+echo "FS_NIFTI_INPUT_LOCAL_FILEPATH:     $FS_NIFTI_INPUT_LOCAL_FILEPATH"
+echo "FS_OUTPUT_S3_FILEPATH:             $FS_OUTPUT_S3_FILEPATH"
+echo "==============================================================="
+
+# infant pipeline input must be ${SUBJECTS_DIR}:${FS_SUB_NAME}/mprage.nii.gz
+if [ -n "${SUBJECTS_DIR}" ] && [ -n "${FS_SUB_NAME}" ]; then
+  echo "---------------------------------------------------------------"
+  echo "SUBJECTS_DIR and FS_SUB_NAME detected. Attempting to make dir"
+  echo "mkdir -p ${SUBJECTS_DIR}/${FS_SUB_NAME}"
+  mkdir -p ${SUBJECTS_DIR}/${FS_SUB_NAME}
+  echo "---------------------------------------------------------------"
+fi
+
+if [ -n "${FS_NIFTI_INPUT_S3_FILEPATH}" ] && [ -n "${FS_NIFTI_INPUT_LOCAL_FILEPATH}" ]; then
+  echo "---------------------------------------------------------------"
+  echo "FS_NIFTI_INPUT_S3_FILEPATH and FS_NIFTI_INPUT_LOCAL_FILEPATH detected. Attempting to copy file locally"
+  echo "aws s3 cp ${FS_NIFTI_INPUT_S3_FILEPATH} ${FS_NIFTI_INPUT_LOCAL_FILEPATH}"
+  aws s3 cp $FS_NIFTI_INPUT_S3_FILEPATH $FS_NIFTI_INPUT_LOCAL_FILEPATH
+  echo "---------------------------------------------------------------"
+fi
+
+# Symlink the volume_mounted model files to where FreeSurfer expects them
+# PW 2021/11/18 No longer needed, since `SSCNN_MODEL_DIR` can now be used in `sscnn_skullstrip`
+# -----------------------------------------------------------------------
+#mkdir -p $FREESURFER_HOME/average/sscnn_skullstripping
+#ln -s $FS_INFANT_MODEL/sscnn_skullstrip/cor_sscnn.h5 $FREESURFER_HOME/average/sscnn_skullstripping/cor_sscnn.h5
+#ln -s $FS_INFANT_MODEL/sscnn_skullstrip/ax_sscnn.h5 $FREESURFER_HOME/average/sscnn_skullstripping/ax_sscnn.h5
+#ln -s $FS_INFANT_MODEL/sscnn_skullstrip/sag_sscnn.h5 $FREESURFER_HOME/average/sscnn_skullstripping/sag_sscnn.h5 
+
+eval "$@"
+
+if [ -n "${FS_OUTPUT_S3_FILEPATH}" ]; then
+  echo "---------------------------------------------------------------"
+  echo "FS_OUTPUT_S3_FILEPATH detected. Attempting to copy the subjects_dir to s3:"
+  echo "aws s3 cp --recursive ${SUBJECTS_DIR}/${FS_SUB_NAME} ${FS_OUTPUT_S3_FILEPATH}"
+  aws s3 cp --recursive ${SUBJECTS_DIR}/${FS_SUB_NAME}  ${FS_OUTPUT_S3_FILEPATH}
+  echo "---------------------------------------------------------------"
+fi

infant/docker/infant-container-entrypoint.bash

@@ -1,9 +1,22 @@
 #!/bin/bash
 
+# TODO: remove `FS_INFANT_MODEL` (no longer needed?)
+echo "==============================================================="
+echo "ENVIRONMENT VARIABLES"
+echo ""
+echo "FREESURFER_HOME:                   $FREESURFER_HOME"
+echo "FS_INFANT_MODEL:                   $FS_INFANT_MODEL"
+echo "SUBJECTS_DIR:                      $SUBJECTS_DIR"
+echo "FS_SUB_NAME:                       $FS_SUB_NAME"
+echo "SSCNN_MODEL_DIR:                   $SSCNN_MODEL_DIR"
+echo "==============================================================="
+
 # Symlink the volume_mounted model files to where FreeSurfer expects them
-mkdir -p $FREESURFER_HOME/average/sscnn_skullstripping
-ln -s $FS_INFANT_MODEL/sscnn_skullstrip/cor_sscnn.h5 $FREESURFER_HOME/average/sscnn_skullstripping/cor_sscnn.h5
-ln -s $FS_INFANT_MODEL/sscnn_skullstrip/ax_sscnn.h5 $FREESURFER_HOME/average/sscnn_skullstripping/ax_sscnn.h5
-ln -s $FS_INFANT_MODEL/sscnn_skullstrip/sag_sscnn.h5 $FREESURFER_HOME/average/sscnn_skullstripping/sag_sscnn.h5 
+# PW 2021/11/18 No longer needed, since `SSCNN_MODEL_DIR` can now be used in `sscnn_skullstrip`
+# -----------------------------------------------------------------------
+#mkdir -p $FREESURFER_HOME/average/sscnn_skullstripping
+#ln -s $FS_INFANT_MODEL/sscnn_skullstrip/cor_sscnn.h5 $FREESURFER_HOME/average/sscnn_skullstripping/cor_sscnn.h5
+#ln -s $FS_INFANT_MODEL/sscnn_skullstrip/ax_sscnn.h5 $FREESURFER_HOME/average/sscnn_skullstripping/ax_sscnn.h5
+#ln -s $FS_INFANT_MODEL/sscnn_skullstrip/sag_sscnn.h5 $FREESURFER_HOME/average/sscnn_skullstripping/sag_sscnn.h5 
 
 eval "$@"

python/freesurfer/samseg/prepareAtlasDirectory.py

@@ -1,12 +1,16 @@
+#!/usr/bin/env python3
+
 from freesurfer.samseg import initVisualizer, requireNumpyArray, gems
 from freesurfer.samseg.io import GMMparameter, kvlReadCompressionLookupTable, \
-                                 kvlWriteCompressionLookupTable, kvlWriteSharedGMMParameters
+                                 kvlWriteCompressionLookupTable, kvlWriteSharedGMMParameters, \
+                                 kvlReadSharedGMMParameters
 from freesurfer.samseg.merge_alphas import kvlGetMergingFractionsTable, kvlMergeAlphas
 import numpy as np
 from scipy import ndimage
 import matplotlib.pyplot as plt
 import os, shutil
-
+import sys
+import argparse
 
 def readAndSimplifyCompressionLookupTable( compressionLookupTableFileName, 
                                            uninterestingStructureSearchStrings=None ):
@@ -254,5 +258,39 @@ def prepareAtlasDirectory( directoryName,
 
     return
 
-
-
+def parse_args(args):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-a", "--atlasdir", required=True,
+                        help="The atlas directory to create")
+    parser.add_argument("-m", "--mesh", required=True,
+                        help="The filename of the mesh collection to use (output of kvlBuildAtlasMesh)")
+    parser.add_argument("-c", "--compression_lut", required=True,
+                        help="The compression lookup table to use (output of kvlBuildAtlasMesh)")
+    parser.add_argument("-g", "--shared_gmm_params", required=True,
+                        help="The filename of the shared GMM parameters to use")
+    parser.add_argument("-t", "--template", required=True,
+                        help="The filename of the template nifti to use")
+    parser.add_argument("--show_figs", default=False)
+    return parser.parse_args()
+
+def main(argv):
+    args = parse_args(argv)
+    if not os.path.exists(args.mesh):
+        print("ERROR: Can't find the mesh file " + args.mesh)
+    if not os.path.exists(args.compression_lut):
+        print("ERROR: Can't find the compression LUT file " + args.compression_lut)
+    if not os.path.exists(args.shared_gmm_params):
+        print("ERROR: Can't find the shared GMM params file " + args.shared_gmm_params)
+    if not os.path.exists(args.template):
+        print("ERROR: Can't find the template file " + args.template)
+
+    shared_gmm_params = kvlReadSharedGMMParameters(args.shared_gmm_params)
+    prepareAtlasDirectory(args.atlasdir,
+                          args.mesh,
+                          args.compression_lut,
+                          shared_gmm_params,
+                          args.template,
+                          showFigures=args.show_figs)
+
+if __name__ == "__main__":
+    sys.exit(main(sys.argv))

python/requirements-extra.txt

@@ -16,3 +16,7 @@ pandas
 matplotlib
 transforms3d
 scikit-image==0.16.2
+
+## Samseg vis
+pyqtgraph
+PyQt5