From 58ecd66545a43ef436a89ebdebadabcb8e8a4a61 Mon Sep 17 00:00:00 2001 From: Xiao Furen Date: Sun, 2 Feb 2025 12:29:19 +0800 Subject: [PATCH] working --- .gitignore | 2 + mri_synthmorph/CMakeLists.txt | 22 + mri_synthmorph/Dockerfile | 46 + mri_synthmorph/README.md | 106 +++ mri_synthmorph/container-script | 96 ++ mri_synthmorph/fs-synthmorph-reg | 1001 +++++++++++++++++++++ mri_synthmorph/mri_synthmorph | 419 +++++++++ mri_synthmorph/mri_synthmorph_apply | 237 +++++ mri_synthmorph/synthmorph/registration.py | 313 +++++++ mri_synthmorph/synthmorph/utils.py | 93 ++ mri_synthmorph/test_apply.sh | 69 ++ mri_synthmorph/test_register.sh | 83 ++ src/conda-create.sh | 4 + src/g4synthmorph.py | 376 ++++++++ src/mri_synthmorph | 1 + src/requirements.txt | 6 + src/synthmorph | 1 + src/synthmorph-ct.py | 59 ++ 18 files changed, 2934 insertions(+) create mode 100644 mri_synthmorph/CMakeLists.txt create mode 100644 mri_synthmorph/Dockerfile create mode 100644 mri_synthmorph/README.md create mode 100644 mri_synthmorph/container-script create mode 100644 mri_synthmorph/fs-synthmorph-reg create mode 100644 mri_synthmorph/mri_synthmorph create mode 100644 mri_synthmorph/mri_synthmorph_apply create mode 100644 mri_synthmorph/synthmorph/registration.py create mode 100644 mri_synthmorph/synthmorph/utils.py create mode 100644 mri_synthmorph/test_apply.sh create mode 100644 mri_synthmorph/test_register.sh create mode 100644 src/conda-create.sh create mode 100644 src/g4synthmorph.py create mode 120000 src/mri_synthmorph create mode 100644 src/requirements.txt create mode 120000 src/synthmorph create mode 100644 src/synthmorph-ct.py diff --git a/.gitignore b/.gitignore index ab3e8ce..550b8ce 100644 --- a/.gitignore +++ b/.gitignore @@ -162,3 +162,5 @@ cython_debug/ # option (not recommended) you can uncomment the following to ignore the entire idea folder. #.idea/ +*.h5 +*.gz diff --git a/mri_synthmorph/CMakeLists.txt b/mri_synthmorph/CMakeLists.txt new file mode 100644 index 0000000..a1f55ee --- /dev/null +++ b/mri_synthmorph/CMakeLists.txt @@ -0,0 +1,22 @@ +project(mri_synthmorph) + +file(GLOB WEIGHTS "synthmorph.*.h5") + +if (FSPYTHON_INSTALL_TREE) + install_pyscript_fspython_tree(mri_synthmorph) + install_symlinks_fspython_tree(TYPE files DESTINATION models ${WEIGHTS}) + install_directories_fspython_tree(synthmorph DESTINATION python/packages) +else() + install_pyscript(mri_synthmorph) + install_symlinks(TYPE files DESTINATION models ${WEIGHTS}) + install_directories(synthmorph DESTINATION python/packages) +endif() + +install_configured(fs-synthmorph-reg DESTINATION bin) + +## 08/2024 - currently failing on Intel Mac + +if(NOT APPLE) + add_test_script(NAME mri_synthmorph_test_register SCRIPT test_register.sh) + add_test_script(NAME mri_synthmorph_test_apply SCRIPT test_apply.sh) +endif() diff --git a/mri_synthmorph/Dockerfile b/mri_synthmorph/Dockerfile new file mode 100644 index 0000000..5961100 --- /dev/null +++ b/mri_synthmorph/Dockerfile @@ -0,0 +1,46 @@ +# Define base image. Set HOME to avoid Matplotlib warning about non-writable +# MPLCONFIGDIR on Neurite import when running as non-root user. +FROM tensorflow/tensorflow:2.17.0-gpu AS base +ENV FREESURFER_HOME=/freesurfer +ENV PYTHONUSERBASE="$FREESURFER_HOME/env" +ENV PATH="$FREESURFER_HOME:$PATH" +ENV HOME=/tmp + + +# Intermediate build stage. Install Python packages to user base for easy COPY. +FROM base AS copy + +COPY --chmod=0775 mri_synthmorph $FREESURFER_HOME/ +COPY --chmod=0664 synthmorph/*.py $FREESURFER_HOME/synthmorph/ +COPY --chmod=0664 synthmorph.*.h5 $FREESURFER_HOME/models/ + +RUN apt-get update && apt-get install -y --no-install-recommends git +RUN python3 -m pip install -U pip +RUN python3 -m pip install --user \ + 'numpy<2.0' \ + git+https://github.com/adalca/pystrum.git@ba35d4b357f54e5ed577cbd413076a07ef810a21 \ + git+https://github.com/adalca/neurite.git@9ae2f5cec2201eedbcc6929cecf852193cef7646 \ + git+https://github.com/freesurfer/surfa.git@041905fca717447780e0cc211197669e3218de2f \ + git+https://github.com/voxelmorph/voxelmorph.git@53d1b95fa734648c92fd8af4f3807b09cb56c342 + +WORKDIR /artifacts +RUN python3 -V >python.txt +RUN python3 -m pip freeze >requirements.txt +RUN mri_synthmorph -h >help.general.txt +RUN mri_synthmorph register -h >help.register.txt +RUN mri_synthmorph apply -h >help.apply.txt + + +# Export Python requirements for reference. Build artifacts will only exist in +# in the target stage `export`. +FROM scratch AS export +COPY --from=copy /artifacts/*.txt / + + +# Exclude Git and caches from final image to save space. Copy only once to +# avoid unnecessary container layers. Set working directory to /mnt for data +# exchange with the host without having to specify the full path. +FROM base +COPY --from=copy $FREESURFER_HOME $FREESURFER_HOME +WORKDIR /mnt +ENTRYPOINT ["mri_synthmorph"] diff --git a/mri_synthmorph/README.md b/mri_synthmorph/README.md new file mode 100644 index 0000000..e88a279 --- /dev/null +++ b/mri_synthmorph/README.md @@ -0,0 +1,106 @@ +# SynthMorph + +This guide explains how to build SynthMorph container images. +It assumes you execute commands in the `mri_synthmorph` directory. +For general information about SynthMorph, visit [synthmorph.io](https://synthmorph.io). + + +## Managing weight files with git-annex + +Weight files are large and therefore managed with `git-annex`. +Instructions with examples are available elsewhere: + +* https://surfer.nmr.mgh.harvard.edu/fswiki/GitAnnex +* https://git-annex.branchable.com/walkthrough + + +## Building SynthMorph images with Docker + +FreeSurfer automatically ships the most recent `mri_synthmorph` and weight files. +Building a standalone container image requires fetching and unlocking the model files with `git-annex`, replacing the symbolic links with the actual files. + +```sh +git fetch datasrc +git annex get . +git annex unlock synthmorph.*.h5 +``` + +Build a new image with the appropriate version tag: + +```sh +tag=X +docker build -t freesurfer/synthmorph:$tag . +``` + + +## Testing the local Docker image + +Update the version reference in the wrapper script and run it to test the local image with Docker. + +```sh +sed -i "s/^\(version = \).*/\1$tag/" synthmorph +./synthmorph -h +``` + + +## Testing with Apptainer + +Testing the image with Apptainer (Singularity) before making it public requires conversion. +If your home directory has a low quota, set up a cache elsewhere: + +```sh +d=$(mktemp -d) +export APPTAINER_CACHEDIR="$d" +export APPTAINER_TMPDIR="$d" +``` + +On the machine running Docker, convert the image with: + +```sh +apptainer build -f synthmorph_$tag.sif docker-daemon://freesurfer/synthmorph:$tag +``` + +If you want to test the image on another machine, save it first. +After transfer to the target machine, build a SIF file as a non-root user using the fakeroot feature. +This relies on namespace mappings set up in /etc/subuid and /etc/subgid (likely by Help). + +```sh +docker save synthmorph:$tag | gzip >synthmorph_$tag.tar.gz +apptainer build -f synthmorph_$tag.sif docker-archive://synthmorph_$tag.tar.gz +``` + +Finally, run the image. + +```sh +apptainer run --nv -e -B /autofs synthmorph_$tag.sif +``` + + +## Pushing to the Docker Hub + +Push the new image to the Docker Hub to make it public. +Update the default "latest" tag, so that `docker pull freesurfer/synthmorph` without a tag will fetch the most recent image. + +```sh +docker push freesurfer/synthmorph:$tag +docker tag freesurfer/synthmorph:$tag freesurfer/synthmorph:latest +docker push freesurfer/synthmorph:latest +``` + + +## Exporting Python requirements + +Export build artifacts for users who wish to create a custom Python environment. + +```sh +docker build --target export --output env . +``` + + +## Final steps + +Lock the annexed weight files again to prevent modification. + +```sh +git annex lock synthmorph.*.h5 +``` diff --git a/mri_synthmorph/container-script b/mri_synthmorph/container-script new file mode 100644 index 0000000..58a7dc2 --- /dev/null +++ b/mri_synthmorph/container-script @@ -0,0 +1,96 @@ +#!/usr/bin/env python3 + +# This wrapper script facilitates setup and use of SynthMorph containers by +# pulling them from the Docker Hub and mounting the host directory defined by +# environment variable SUBJECTS_DIR to /mnt in the container. Invoke the script +# just like `mri_synthmorph` in FreeSurfer, with one exception: you can only +# read and write data under SUBJECTS_DIR, which will be the working directory +# in the container. If unset, SUBJECTS_DIR defaults to your current directory. +# This means you can access relative paths under your working directory without +# setting SUBJECTS_DIR. In other words, SUBJECTS_DIR sets the working directory +# for SynthMorph, and you can specify paths relative to it. + +# Update the version to pull a different image, unless you already have it. +version = 4 + +# Local image location for Apptainer/Singularity. Set an absolute path to avoid +# pulling new images when you change the folder. Ignored for Docker and Podman. +sif_file = f'synthmorph_{version}.sif' + +# We will use the first of the below container systems found in your PATH, from +# left to right. You may wish to reorder them, If you have several installed. +tools = ('docker', 'apptainer', 'singularity', 'podman') + + +import os +import sys +import signal +import shutil +import subprocess + + +# Report version. Avoid errors when piping, for example to `head`. +signal.signal(signal.SIGPIPE, handler=signal.SIG_DFL) +hub = 'https://hub.docker.com/u/freesurfer' +print(f'Running SynthMorph version {version} from {hub}') + + +# Find a container system. +for tool in tools: + path = shutil.which(tool) + if path: + print(f'Using {path} to manage containers') + break + +if not path: + print(f'Cannot find container tools {tools} in PATH', file=sys.stderr) + exit(1) + + +# Prepare bind path and URL. Mount SUBJECTS_DIR as /mnt inside the container, +# which we made the working directory when building the image. While Docker +# and Podman will respect it, they require absolute paths for bind mounts. +host = os.environ.get('SUBJECTS_DIR', os.getcwd()) +host = os.path.abspath(host) +print(f'Will bind /mnt in image to SUBJECTS_DIR="{host}"') + +image = f'freesurfer/synthmorph:{version}' +if tool != 'docker': + image = f'docker://{image}' + + +# Run Docker containers with the UID and GID of the host user. This user will +# own bind mounts inside the container, preventing output files owned by root. +# Root inside a rootless Podman container maps to the non-root host user, which +# is what we want. If we set UID and GID inside the container to the non-root +# host user as we do for Docker, then these would get remapped according to +# /etc/subuid outside, causing problems with read and write permissions. +if tool in ('docker', 'podman'): + arg = ('run', '--rm', '-v', f'{host}:/mnt') + + # Pretty-print help text. + if sys.stdout.isatty(): + arg = (*arg, '-t') + if tool == 'docker': + arg = (*arg, '-u', f'{os.getuid()}:{os.getgid()}') + + arg = (*arg, image) + + +# For Apptainer/Singularity, the user inside and outside the container is the +# same. The working directory is also the same, unless we explicitly set it. +if tool in ('apptainer', 'singularity'): + arg = ('run', '--nv', '--pwd', '/mnt', '-e', '-B', f'{host}:/mnt', sif_file) + + if not os.path.isfile(sif_file): + print(f'Cannot find image {sif_file}, pulling it', file=sys.stderr) + proc = subprocess.run((tool, 'pull', sif_file, image)) + if proc.returncode: + exit(proc.returncode) + + +# Summarize and launch container. +print('Command:', ' '.join((tool, *arg))) +print('SynthMorph arguments:', *sys.argv[1:]) +proc = subprocess.run((tool, *arg, *sys.argv[1:])) +exit(proc.returncode) diff --git a/mri_synthmorph/fs-synthmorph-reg b/mri_synthmorph/fs-synthmorph-reg new file mode 100644 index 0000000..ab9a533 --- /dev/null +++ b/mri_synthmorph/fs-synthmorph-reg @@ -0,0 +1,1001 @@ +#!/bin/tcsh -f +# fs-synthmorph-reg - sources +if(-e $FREESURFER_HOME/sources.csh) then + source $FREESURFER_HOME/sources.csh +endif + +if($?FSMNI152DIR == 0) setenv FSMNI152DIR $FREESURFER/average/mni_icbm152_nlin_asym_09c +unsetenv FS_LOCAL_PYTHONPATH + +set VERSION = '$Id$'; +set scriptname = `basename $0` + +set outdir = (); +set subject = (); +set m3z = (); +set m3zinv = () +set invol = () +set targvol = () +set lta2 = () +set vgthresh = 1e-5 +set StripInput = 0 +set StripTarget = 0 +set ReRun = 0 + +set threads = 1 +set ForceUpdate = 0 +set antsreg = antsRegistrationSyNQuick.sh +set UseQuick = 0 +set CropInput = 1 +set CropTarget = 1 +set MNITarget = 1 +# Synthmorph uses 1mm internally, so useing target res=1mm does not +# slow things down or take more mem and it is probably gives a little +# bit better registration. ANTs speed and performance will be +# affected. +set MNITargetRes = 1.0mm +set MNIOutputRes = 1.0mm +set DoCBIG = 0 +set DoTest = 0 +set dim = 3 +set UseAnts = 0 # 0 means use synthmorph +set ComputeInverse = 1 +set PitStr = "" +set AffineOnly = 0 + +set tmpdir = (); +set cleanup = 1; +set LF = (); + +set inputargs = ($argv); +set PrintHelp = 0; +if($#argv == 0) goto usage_exit; +set n = `echo $argv | grep -e -help | wc -l` +if($n != 0) then + set PrintHelp = 1; + goto usage_exit; +endif +set n = `echo $argv | grep -e -version | wc -l` +if($n != 0) then + echo $VERSION + exit 0; +endif +goto parse_args; +parse_args_return: +goto check_params; +check_params_return: + +set StartTime = `date`; +set tSecStart = `date '+%s'`; +set year = `date +%Y` +set month = `date +%m` +set day = `date +%d` +set hour = `date +%H` +set min = `date +%M` + +if($#outdir) then + mkdir -p $outdir/log +else + set outdir = `dirname $m3z` + mkdir -p $outdir +endif +pushd $outdir > /dev/null +set outdir = `pwd`; +popd > /dev/null + +if($#tmpdir == 0) set tmpdir = $outdir/tmp +mkdir -p $tmpdir + +# Set up log file +if($#LF == 0) set LF = $outdir/log/fs-synthmorph-reg.Y$year.M$month.D$day.H$hour.M$min.log +if($LF != /dev/null) rm -f $LF +echo "Log file for fs-synthmorph-reg" >> $LF +date | tee -a $LF +echo "" | tee -a $LF +echo "setenv SUBJECTS_DIR $SUBJECTS_DIR" | tee -a $LF +echo "cd `pwd`" | tee -a $LF +echo $0 $inputargs | tee -a $LF +ls -l $0 | tee -a $LF +echo "" | tee -a $LF +cat $FREESURFER_HOME/build-stamp.txt | tee -a $LF +echo $VERSION | tee -a $LF +uname -a | tee -a $LF +echo "pid $$" | tee -a $LF +if($?PBS_JOBID) then + echo "pbsjob $PBS_JOBID" >> $LF +endif +if($?SLURM_JOB_ID) then + echo SLURM_JOB_ID $SLURM_JOB_ID >> $LF +endif + +#======================================================== +# Note: Ants might not be deterministc with threads +setenv OMP_NUM_THREADS $threads +setenv ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS $threads + +if($StripInput) then + #set involstripped = $tmpdir/invol.stripped.nii.gz # breaks header + set involstripped = $tmpdir/invol.stripped.mgz + set ud = `UpdateNeeded $involstripped $invol` + if($ud || $ForceUpdate) then + set cmd = (mri_synthstrip -i $invol -o $involstripped -t $threads) + echo $cmd | tee -a $LF + $cmd | tee -a $LF + if($status) exit 1 + endif + set invol = $involstripped +endif +if($StripTarget) then + #set targvolstripped = $tmpdir/targ.stripped.nii.gz # breaks header + set targvolstripped = $tmpdir/targ.stripped.mgz + set ud = `UpdateNeeded $targvolstripped $targvol` + if($ud || $ForceUpdate) then + set cmd = (mri_synthstrip -i $targvol -o $targvolstripped -t $threads) + echo $cmd | tee -a $LF + $cmd | tee -a $LF + if($status) exit 1 + endif + set targvol = $targvolstripped +endif + +if($CropInput) then # Crop for speed + # Input volume + set involcrop = $tmpdir/invol.crop.nii.gz + set ud = `UpdateNeeded $involcrop $invol` + if($ud || $ForceUpdate) then + set cmd = (mri_mask -bb 3 $invol $invol $involcrop) + echo $cmd | tee -a $LF + $cmd | tee -a $LF + if($status) exit 1 + endif + # To get back to the full FoV + set regcroptoinvol = $tmpdir/reg.crop-to-invol.lta + set ud = `UpdateNeeded $regcroptoinvol $involcrop $invol` + if($ud || $ForceUpdate) then + set cmd = (lta_convert --inlta identity.nofile --src $involcrop \ + --trg $invol --outlta $regcroptoinvol) + echo $cmd | tee -a $LF + $cmd |& tee -a $LF + if($status) exit 1 + endif +else + set involcrop = $invol + set regcroptoinvol = () +endif + +if($CropTarget) then + if(! $MNITarget) then + # Target vol + set targvolcrop = $tmpdir/targvol.crop.nii.gz + set ud = `UpdateNeeded $targvolcrop $targvol` + if($ud || $ForceUpdate) then + set cmd = (mri_mask -crop 3 $targvol $targvol $targvolcrop) + echo $cmd | tee -a $LF + $cmd | tee -a $LF + if($status) exit 1 + endif + # To get back to the full FoV + set regcroptotarg = $tmpdir/reg.crop-to-targ.lta + set ud = `UpdateNeeded $regcroptotarg $targvolcrop $targvol` + if($ud || $ForceUpdate) then + set cmd = (lta_convert --inlta identity.nofile --src $targvolcrop \ + --trg $targvol --outlta $regcroptotarg) + echo $cmd | tee -a $LF + $cmd |& tee -a $LF + if($status) exit 1 + endif + else + set targvolcrop = $FSMNI152DIR/reg-targets/mni152.${MNITargetRes}$PitStr.cropped.nii.gz + set regcroptotarg = $FSMNI152DIR/reg-targets/reg.${MNITargetRes}.cropped.to.${MNIOutputRes}.lta + # Note that lta does not have PitStr + endif +else + set targvolcrop = $targvol + set regcroptotarg = () +endif + +if($UseAnts) then +# Compute both the affine and the warp with ANTs. The Affine goes from +# the input croped space to the target cropped space. The warp goes +# from target cropped to target cropped. +set warp = $tmpdir/reg.1Warp.nii.gz +set affinemat = $tmpdir/reg.0GenericAffine.mat +set ud = `UpdateNeeded $warp $involcrop $targvol` +if($ud || $ForceUpdate) then + date | tee -a $LF + pushd $tmpdir # this program can crash if stuff in the current dir + if(! $DoCBIG) then + set cmd = ($antsreg -d $dim -m $involcrop -f $targvolcrop -o reg. -n $threads) + else + set cmd = (antsRegistration --dimensionality $dim \ + --output [reg.,reg.Warped.nii.gz,reg.InverseWarped.nii.gz] \ + --initial-moving-transform [ $involcrop, $targvolcrop, 1 ] \ + --collapse-output-transforms 1 \ + --use-histogram-matching 1 \ + --use-estimate-learning-rate-once 1 \ + --metric mattes[ $fixed, $moving, 1, 32, regular, 0.3] \ + --transform affine[ 0.1 ] \ + --convergence [ 100x100x200, 1.e-8, 20 ] \ + --smoothing-sigmas 4x2x1vox \ + --shrink-factors 3x2x1 -l 1 \ + -metric cc[ $fixed, $moving, 1, 4] \ + --transform SyN[ .20, 3, 0] \ + --convergence [ 100x100x50, 0, 5 ] \ + --smoothing-sigmas 1x0.5x0vox \ + --shrink-factors 4x2x1) + endif + echo "\n\n"| tee -a $LF + echo $cmd | tee -a $LF + fs_time $cmd |& tee -a $LF + if($status) goto error_exit + echo "\n\n"| tee -a $LF + if($CropInput || $CropTarget) then + # Somehow ANTs will change the geom slightly, so map it back to the + # the target space. This is a bit of a hack that seems to work, but + # it is not principled (ie, it was derived by trial and error, so + # there is a danger that it will not work all the time). + set cmd = (mri_vol2vol --regheader --targ $targvolcrop --mov $warp --o $warp) + echo "\n\n"| tee -a $LF + echo $cmd | tee -a $LF + fs_time $cmd |& tee -a $LF + if($status) goto error_exit + echo "\n\n"| tee -a $LF + endif + popd +endif + +# Convert the ANTs affine to txt file (maps involcrop to the input of the targvolcrop space) +set affinetxt = $tmpdir/reg.0GenericAffine.txt +set ud = `UpdateNeeded $affinetxt $affinemat` +if($ud || $ForceUpdate) then + date | tee -a $LF + set cmd = (ConvertTransformFile $dim $affinemat $affinetxt --hm --ras) + echo "\n\n"| tee -a $LF + echo $cmd | tee -a $LF + fs_time $cmd |& tee -a $LF + if($status) goto error_exit + echo "\n\n"| tee -a $LF +endif + +# Convert the ANTs affine to an LTA +set affinelta = $tmpdir/reg.0GenericAffine.lta +set ud = `UpdateNeeded $affinelta $affinetxt` +if($ud || $ForceUpdate) then + date | tee -a $LF + set cmd = (lta_convert --src $involcrop --trg $targvolcrop --outlta $affinelta) + if($dim == 2) set cmd = ($cmd --inniftyreg2d $affinetxt); + if($dim == 3) set cmd = ($cmd --inniftyreg $affinetxt); + echo "\n\n"| tee -a $LF + echo $cmd | tee -a $LF + fs_time $cmd |& tee -a $LF + if($status) goto error_exit + echo "\n\n"| tee -a $LF +endif + +# Create one lta that maps from the uncropped input vol space to the taret space +if($CropInput || $CropTarget) then + set regcroptargtocropinvol = $tmpdir/reg.croptarg_to_cropinvol.lta + set ud = `UpdateNeeded $regcroptargtocropinvol $affinelta $regcroptoinvol` + if($ud || $ForceUpdate) then + date | tee -a $LF + set cmd = (mri_concatenate_lta -invert1 -invertout $affinelta $regcroptoinvol $regcroptargtocropinvol) + echo "\n\n"| tee -a $LF + echo $cmd | tee -a $LF + fs_time $cmd |& tee -a $LF + if($status) goto error_exit + echo "\n\n"| tee -a $LF + endif +else + set regcroptargtocropinvol = $affinelta +endif + +# Concatenate the affines and the warp (ANTS) +set ud = `UpdateNeeded $m3z $warp $targvolcrop $regcroptargtocropinvol $regcroptotarg ` +if($ud || $ForceUpdate) then + date | tee -a $LF + set cmd = (mri_warp_convert --initk $warp --insrcgeom $targvolcrop \ + --outm3z $m3z --vg-thresh $vgthresh) + if($CropInput) set cmd = ($cmd --lta1-inv $regcroptargtocropinvol ) + if($#regcroptotarg) set cmd = ($cmd --lta2 $regcroptotarg) + echo "\n\n"| tee -a $LF + echo $cmd | tee -a $LF + fs_time $cmd |& tee -a $LF + if($status) goto error_exit + echo "\n\n"| tee -a $LF +endif + +else # End ANTS ================= else Use synthmorph + +set gpuopt = "" +# Compute both the affine, maps from involcrop to targvolcrop +set affinelta = $outdir/aff.lta +set ud = `UpdateNeeded $affinelta $involcrop $targvolcrop` +if($ud || $ForceUpdate) then + date | tee -a $LF + set cmd = (mri_synthmorph -m affine -t $affinelta $involcrop $targvolcrop -j $threads $gpuopt) + echo "\n\n"| tee -a $LF + echo $cmd | tee -a $LF + fs_time $cmd |& tee -a $LF + if($status) goto error_exit +endif + +# Create one lta that maps from the uncropped input vol space to the (cropped) target space +set regtargtoinvol = $outdir/reg.targ_to_invol.lta +set reginvoltotarg = $outdir/reg.invol_to_targ.lta +if($CropInput || $CropTarget) then + set reginvoltocroptarg = $tmpdir/reg.invol_to_croptarg.lta + set ud = `UpdateNeeded $reginvoltocroptarg $affinelta $regcroptoinvol` + if($ud || $ForceUpdate) then + date | tee -a $LF + set cmd = (mri_concatenate_lta -invert1 -invertout $affinelta $regcroptoinvol $reginvoltocroptarg) + echo "\n\n"| tee -a $LF + echo $cmd | tee -a $LF + fs_time $cmd |& tee -a $LF + if($status) goto error_exit + echo "\n\n"| tee -a $LF + endif + if($MNITarget) then + # Create one lta that maps from the uncropped input vol space to the uncropped target space + # This produces a linear tranform that can be used to map from target to input volume just + # like the nonlinear transform + set ud = `UpdateNeeded $regtargtoinvol $reginvoltocroptarg` + if($ud || $ForceUpdate) then + set regtargtocropped = $FSMNI152DIR/reg-targets/reg.$MNIOutputRes.to.$MNITargetRes.cropped.lta + set cmd = (mri_concatenate_lta -invert2 $regtargtocropped $reginvoltocroptarg $regtargtoinvol) + echo $cmd | tee -a $LF + fs_time $cmd |& tee -a $LF + if($status) goto error_exit + echo "\n\n"| tee -a $LF + endif + endif + # Compute the inverse: invol-to-targ + set ud = `UpdateNeeded $reginvoltotarg $regtargtoinvol` + if($ud || $ForceUpdate) then + set cmd = (lta_convert --invert --inlta $regtargtoinvol --outlta $reginvoltotarg) + echo $cmd | tee -a $LF + fs_time $cmd |& tee -a $LF + if($status) goto error_exit + echo "\n\n"| tee -a $LF + endif +else + pushd $tmpdir + ln -fs aff.lta reg.invol_to_targ.lta + popd + set reginvoltocroptarg = $affinelta + set reginvoltotarg = $affinelta + # Compute target to invol + set ud = `UpdateNeeded $regtargtoinvol $reginvoltotarg` + if($ud || $ForceUpdate) then + set cmd = (lta_convert --invert --inlta $reginvoltotarg --outlta $regtargtoinvol) + echo $cmd | tee -a $LF + fs_time $cmd |& tee -a $LF + if($status) goto error_exit + echo "\n\n"| tee -a $LF + endif + +endif + +if($MNITarget) then + # Create an xfm file which can be used as the taliarch.xfm + set xfm = $outdir/reg.invol_to_targ.xfm + set ud = `UpdateNeeded $xfm $reginvoltotarg` + if($ud) then + set cmd = (lta_convert --inlta $reginvoltotarg --outmni $xfm) + echo $cmd | tee -a $LF + $cmd |& tee -a $LF + if($status) goto error_exit + echo "\n\n"| tee -a $LF + endif +endif + +echo "" | tee -a $LF +echo "To check affine registration" | tee -a $LF +echo "tkregisterfv --mov $invol --targ $targvol --reg $reginvoltotarg" | tee -a $LF +echo "" | tee -a $LF + +if($AffineOnly) then + echo "AffineOnly specified, so exiting now" | tee -a $LF + goto done +endif + +# Now compute the nonlinear part +set deform = $tmpdir/deform.mgz +set smvol = $tmpdir/synthmorph.out.mgz +set ud = `UpdateNeeded $deform $affinelta $involcrop $targvolcrop` +if($ud) then + set cmd = (mri_synthmorph -m deform -t $deform -i $affinelta $involcrop $targvolcrop -j $threads $gpuopt) + if($DoTest) set cmd = ($cmd -o $smvol); # can be a pain when you want to rerun to create test output + echo "\n\n" | tee -a $LF + echo $cmd | tee -a $LF + fs_time $cmd |& tee -a $LF + if($status) goto error_exit +endif + +# Compute warp to the full FoV of the input (synthmorph) +set ud = `UpdateNeeded $m3z $deform $targvolcrop $reginvoltocroptarg $regcroptotarg` +if($ud || $ForceUpdate) then + # This call is not intuitive to me. It appears that the insrcgeom + # will set the vox2ras for the image-side of the warp. The warp + # maps from target voxel to cropped input RAS. The insrcgeom + # indicates how to go from that RAS to a CRS in the cropped input, + # then the lta1 indicates how to go from the cropped CRS to the + # uncropped CRS. + set cmd = (mri_warp_convert --inras $deform --insrcgeom $involcrop\ + --outm3z $m3z --vg-thresh $vgthresh) + if($CropInput) set cmd = ($cmd --lta1-inv $regcroptoinvol) + if($#regcroptotarg) set cmd = ($cmd --lta2 $regcroptotarg) + echo "\n\n$cmd" | tee -a $LF + fs_time $cmd |& tee -a $LF + if($status) goto error_exit +endif + +endif # Synthmorph + +if($ComputeInverse) then + set ud = `UpdateNeeded $m3zinv $m3z` + if($ud || $ForceUpdate) then + set cmd = (mri_ca_register -invert-and-save $m3z $m3zinv) + echo "\n\n$cmd" | tee -a $LF + fs_time $cmd |& tee -a $LF + if($status) goto error_exit + endif +endif + +# Test against output generated by the registration program +if($DoTest) then + if($UseAnts) then + set mvol = $tmpdir/reg.Warped.nii.gz + else + set mvol = $smvol + endif + if($CropInput || $CropTarget) then + # Ants will output to the cropped target space, need to resample + # to the full uncropped space. + set mvol2 = $tmpdir/morph.out.nii.gz + set ud = `UpdateNeeded $mvol2 $mvol $targvol` + if($ud || $ForceUpdate) then + set cmd = (mri_vol2vol --regheader --mov $mvol --targ $targvol --o $mvol2) + echo $cmd | tee -a $LF + fs_time $cmd |& tee -a $LF + if($status) exit 1 + endif + set mvol = $mvol2 + endif + + set ud = `UpdateNeeded $testvol $invol $m3z $mvol` + if($ud || $ForceUpdate) then + set cmd = (mri_convert -rt nearest $invol -at $m3z $testvol) + echo $cmd | tee -a $LF + fs_time $cmd |& tee -a $LF + if($status) exit 1 + set cmd = (mri_diff --po $testvol $mvol) + echo $cmd | tee -a $LF + $cmd | tee $tmpdir/test.diff.dat |& tee -a $LF + endif + echo "tkmeditfv -f $targvol -aux $testvol" |& tee -a $LF +endif + +#======================================================== + +done: + +# Cleanup +if($cleanup) rm -rf $tmpdir + +# Done +echo " " |& tee -a $LF +set tSecEnd = `date '+%s'`; +@ tSecRun = $tSecEnd - $tSecStart; +set tRunMin = `echo $tSecRun/60|bc -l` +set tRunMin = `printf %5.2f $tRunMin` +set tRunHours = `echo $tSecRun/3600|bc -l` +set tRunHours = `printf %5.2f $tRunHours` +echo "Started at $StartTime " |& tee -a $LF +echo "Ended at `date`" |& tee -a $LF +echo "Fs-Synthmorph-Reg-Run-Time-Sec $tSecRun" |& tee -a $LF +echo "Fs-Synthmorph-Reg-Run-Time-Min $tRunMin" |& tee -a $LF +echo "Fs-Synthmorph-Reg-Run-Time-Hours $tRunHours" |& tee -a $LF +echo " " |& tee -a $LF +echo "fs-synthmorph-reg Done" |& tee -a $LF +exit 0 + +############################################### + +############--------------################## +error_exit: +echo "ERROR:" + +exit 1; +############################################### + +############--------------################## +parse_args: +set cmdline = ($argv); +while( $#argv != 0 ) + + set flag = $argv[1]; shift; + + switch($flag) + + case "--i": + if($#argv < 1) goto arg1err; + set invol = $argv[1]; shift; + if(! -e $invol) then + echo "ERROR: cannot find $invol" + exit 1 + endif + set invol = `getfullpath $invol` + breaksw + + case "--t": + if($#argv < 1) goto arg1err; + set targvol = $argv[1]; shift; + if(! -e $targvol) then + echo "ERROR: cannot find $targvol" + exit 1 + endif + set targvol = `getfullpath $targvol` + set MNITarget = 0 + set Crop = 1 + breaksw + + case "--warp": + case "--m3z": + if($#argv < 1) goto arg1err; + set m3z = $argv[1]; shift; + breaksw + + case "--affine-only": + set AffineOnly = 1 + breaksw + + case "--o": + if($#argv < 1) goto arg1err; + set outdir = $argv[1]; shift; + breaksw + + case "--s": + if($#argv < 1) goto arg1err; + set subject = $argv[1]; shift; + breaksw + + case "--sd": + if($#argv < 1) goto arg1err; + setenv SUBJECTS_DIR $argv[1]; shift; + breaksw + + case "--threads": + if($#argv < 1) goto arg1err; + set threads = $argv[1]; shift; + breaksw + + case "--vg-thresh": + if($#argv < 1) goto arg1err; + set vgthresh = $argv[1]; shift; + breaksw + + case "--synthmorph": + set UseAnts = 0 + breaksw + case "--ants": + set UseAnts = 1 + breaksw + + case "--test": + set DoTest = 1 + breaksw + case "--no-test": + set DoTest = 0 + breaksw + + case "--2d": + set dim = 2 + set DoTest = 0 + breaksw + + case "--quick": + set antsreg = antsRegistrationSyNQuick.sh + set UseQuick = 1 + breaksw + case "--no-quick": + case "--syn": + set antsreg = antsRegistrationSyN.sh + set UseQuick = 0 + breaksw + + case "--crop": + set CropInput = 1 + set CropTarget = 1 + breaksw + case "--no-crop": + set CropInput = 0 + set CropTarget = 0 + breaksw + + case "--strip": + set StripInput = 1 + set StripTarget = 1 + breaksw + case "--no-strip": + set StripInput = 0 + set StripTarget = 0 + breaksw + case "--strip-input": + set StripInput = 1 + breaksw + case "--no-strip-input": + set StripInput = 0 + breaksw + case "--strip-target": + set StripTarget = 1 + breaksw + case "--no-strip-target": + set StripTarget = 0 + breaksw + + case "--inv": + set ComputeInverse = 1 + breaksw + case "--no-inv": + set ComputeInverse = 0 + breaksw + + case "--mni": + set MNITarget = 1 + breaksw + case "--no-nmni": + set MNITarget = 0 + breaksw + + case "--mni-res": + case "--mni-int-res": + case "--mni-targ-res": + if($#argv < 1) goto arg1err; + set MNITargetRes = $argv[1]; shift; + if($MNITargetRes != 1.0mm && $MNITargetRes != 1.5mm && $MNITargetRes != 2.0mm) then + echo "ERROR: --mni-res must be 1.0mm, 1.5mm, or 2.0mm" + exit 1 + endif + set MNITarget = 1 + breaksw + + case "--mni-output-res": + case "--mni-out-res": + if($#argv < 1) goto arg1err; + set MNIOutputRes = $argv[1]; shift; + if($MNIOutputRes != 1.0mm && $MNIOutputRes != 1.5mm && $MNIOutputRes != 2.0mm && $MNIOutputRes != conformed) then + echo "ERROR: --mni-out-res must be 1.0mm, 1.5mm, 2.0mm, or conformed" + exit 1 + endif + set MNITarget = 1 + breaksw + + case "--mni-1": + set MNIOutputRes = 1.0mm + set MNITarget = 1 + breaksw + + case "--mni-1.5": + set MNIOutputRes = 1.5mm + set MNITarget = 1 + breaksw + + case "--mni-2": + set MNIOutputRes = 2.0mm + set MNITarget = 1 + breaksw + + case "--mni-conformed": + set MNIOutputRes = conformed + set MNITarget = 1 + breaksw + + case "--pituitary": + set PitStr = ".pit" + breaksw + + case "--cbig": + set DoCBig = 1 + breaksw + case "--no-cbig": + set DoCBig = 0 + breaksw + + case "--force": + set ForceUpdate = 1 + breaksw + case "--rerun": + # Prevents it from bailing out on the first check. This can be + # handy when debugging and you want it to get into the interior + # of the code. + set ReRun = 1 + breaksw + + case "--log": + if($#argv < 1) goto arg1err; + set LF = $argv[1]; shift; + breaksw + + case "--nolog": + case "--no-log": + set LF = /dev/null + breaksw + + case "--tmp": + case "--tmpdir": + if($#argv < 1) goto arg1err; + set tmpdir = $argv[1]; shift; + set cleanup = 0; + breaksw + + case "--nocleanup": + set cleanup = 0; + breaksw + + case "--cleanup": + set cleanup = 1; + breaksw + + case "--debug": + set verbose = 1; + set echo = 1; + breaksw + + default: + echo ERROR: Flag $flag unrecognized. + echo $cmdline + exit 1 + breaksw + endsw + +end + +goto parse_args_return; +############--------------################## + +############--------------################## +check_params: + +if($UseAnts == -1) then + echo "ERROR: must spec either --ants or --synthmorph" + exit 1 +endif + +# This will keep the temp stuff from being deleted +if($#outdir && $#tmpdir == 0) then + set tmpdir = $outdir + set cleanup = 0 +endif + +if($#subject) then + set sd = $SUBJECTS_DIR/$subject + if(! -e $sd) then + echo "ERROR: cannot find $subject" + exit 1; + endif + if($#invol == 0) set invol = $sd/mri/norm.mgz + if($#outdir == 0 && $MNITarget) then + if($UseAnts) then + set outdir = $sd/mri/transforms/ants.warp.$MNITargetRes.$MNIOutputRes$PitStr + else + set outdir = $sd/mri/transforms/synthmorph.$MNITargetRes.$MNIOutputRes$PitStr + endif + #don't create particular files in the transform dir + #set m3z = $outdir.nii.gz + #set m3zinv = $outdir.inv.nii.gz # Just in case the invert is requested + endif +endif + +# This is too complicated now as it is not just setting the target volume to +# the cropped; have to set the reg-crop-to-target too +#if($MNITarget) set CropTarget = 0; + +if(! $MNITarget) then + # Could get this to work, but not really needed + set CropTarget = 0; + set CropInput = 0; +endif + +if($#outdir == 0) then + echo "ERROR: must spec output dir" + exit 1 +endif +if($#invol == 0) then + echo "ERROR: must spec input volume" + exit 1 +endif +foreach f ($invol $targvol) + if(! -e $f) then + echo "ERROR: cannot find $f" + exit 1 + endif +end + +mkdir -p $outdir +set outdir = `getfullpath $outdir` +if($#m3z == 0) then + if($MNITarget) then + # Lose ants vs synthseg distinction in the output + set m3z = $outdir/warp.to.mni152.${MNITargetRes}.${MNIOutputRes}.nii.gz + set m3zinv = $outdir/warp.to.mni152.${MNITargetRes}.${MNIOutputRes}.inv.nii.gz + else + set m3z = $outdir/warp.nii.gz + set m3zinv = $outdir/warp.inv.nii.gz + endif +endif + +if($ComputeInverse && $#m3zinv == 0) then + set stem = `fname2stem $m3z` + set ext = `fname2ext $m3z` + set m3zinv = $stem.inv.$ext +endif + +if($#targvol && $MNITarget) then + echo "ERROR: cannot specify both target volume and mni target" + exit 1 +endif + +# MNI is already skull stripped +if($MNITarget) set targvol = $FSMNI152DIR/reg-targets/mni152.${MNIOutputRes}${PitStr}.nii.gz + +if($UseQuick) then + set antsreg = antsRegistrationSyNQuick.sh +else + set antsreg = antsRegistrationSyN.sh +endif + +set testvol = () +if($DoTest) set testvol = $outdir/test.nii.gz + +set ud1 = `UpdateNeeded $m3z $targvol $invol` +set ud2 = 0 +if($ComputeInverse) set ud2 = `UpdateNeeded $m3zinv $targvol $invol` +if($ReRun == 0 && $ud1 == 0 && $ud2 == 0 && $ForceUpdate == 0) then + echo "fs-synthmorph-reg: update not needed" + exit 0 +endif + +goto check_params_return; +############--------------################## + +############--------------################## +arg1err: + echo "ERROR: flag $flag requires one argument" + exit 1 +############--------------################## +arg2err: + echo "ERROR: flag $flag requires two arguments" + exit 1 +############--------------################## + +############--------------################## +usage_exit: + echo "" + echo "fs-synthmorph-reg -- frontend for running mri_synthmorph (24GB)" + echo " --i invol" + echo " --t targetvol" + echo " --warp warp.{m3z,mgz,nii.gz} : output warp file (or save in outdir)" + echo " --o outdir" + echo " --s subject (invol=norm.mgz, warp=mri/transforms/synthmorph.1.0mm.1.0mm.nii.gz targ=MNI152.1mm)" + echo " --threads threads" + echo " --mni-targ-res 1.0mm, 1.5mm, 2.0mm (default is 1.0mm)" + echo " --mni-out-res 1.0mm, 1.5mm, 2.0mm, conformed (default is 1.0mm)" + echo " --no-inv : do not compute warp inverse (computed by default)" + echo " --affine-only : only perform the affine registration part" + echo " --strip : skull strip input and target" + echo " --strip-input : skull strip input" + echo " --strip-target : skull strip target" + echo " --crop/--no-crop (default is $CropInput, but no-crop if not registering to mni)" + echo " --pituitary : select the MNI target volume that does not mask out the pituitary" + echo " --test : resample the input volume to the output space and store as test.nii.gz" + echo " --vg-thresh vgthresh : threshold for testing diffs in volume geom" + echo " --ants : use ANTs registration intstead of synthmorph" + echo " --force : regenerate all output " + echo " --rerun : for debugging only, when you want it to get into the inteiror of the code" + echo " --tmp tmpdir" + echo " " + + + if(! $PrintHelp) exit 1; + echo $VERSION + cat $0 | awk 'BEGIN{prt=0}{if(prt) print $0; if($1 == "BEGINHELP") prt = 1 }' +exit 1; + +#---- Everything below here is printed out as part of help -----# +BEGINHELP + +This is a frontend for running mri_synthmorph, especially for +registering to mni152 space. For typical usage, it will consume less +than 15GB of memory. For affine only it requires less than 7GB. + +Simple usage: + +fs-synthmorph-reg --i inputvol.nii.gz --o synthmorphdir + +This will register to the mni152 saving output files in synthmorphdir. The nonlinear +warp will be outdir/warp.to.mni152.1.0mm.1.0mm.nii.gz which can be applied like: + +mri_convert inputvol.nii.gz -at warp.to.mni152.1.0mm.1.0mm.nii.gz inputvol.mni152.nii.gz + +which can be checked with + +freeview inputvol.mni152.nii.gz $FREESURFER/average/mni_icbm152_nlin_asym_09c/reg-targets/mni152.1.0mm.nii.gz + +If you have created the warp from a FreeSurfer subject space (either with --s or passing a +conformed volume), then you can apply the warp to the surfaces, eg, + +mris_apply_reg --warp lh.white warp.to.mni152.1.0mm.1.0mm.inv.nii.gz lh.warp.mni152 + +which you can check with + +freeview inputvol.mni152.nii.gz $FREESURFER/average/mni_icbm152_nlin_asym_09c/reg-targets/mni152.1.0mm.nii.gz -f lh.warp.mni152 + +Skull stripping: mri_synthmorph is quite robust to the presence or +absence of skull stripping, but there will be some differences, so, to +be on the safe side, we recommend that images be skull stripped prior +to registration. When registering to the mni152, it will automatically +use a skull stripped volume. If you want fs-synthmorph-reg to skull +strip for you, then add --strip (uses mri_synthstrip). Note that if +you pass a FS subject (--s), it will use the norm.mgz volume, which is +already skull stripped. Note that ANTs performance and speed will be +very sensitive to stripping. + +Cropping: internally, fs-synthmorph-reg will crop the input to a +minimal window around non-zero voxels in the image. In the end, this +probably does not have much of an effect for synthmorph as it will +sample to 256^3 internally. For ANTs, this can speed up the program +dramatically and reduce its memory footprint. fs-synthmorph-reg will +manage all of the transforms so that they all reference the uncropped +space, ie, the user is completely insulated from the cropping process +and its implications. To turn off cropping, use --no-crop. If you are +not stripping, it probably does not make sense to crop. + +Target and Output MNI Resolution: the first 1.0mm in the warp file +name means that the registration is done to a version of the mni152 +with an isotropic resolution of 1.0mm (this will be how finely the +warp field is sampled). The second 1.0mm means that the warp file will +actually warp to the the mni152 with an isotropic resolution of 1.0mm. +You can change the resolution of the warp field with the +--mni-targ-res flag (options are 1.0mm, 1.5mm, 2.0mm), but, +internally, synthmorph uses 1mm, so this makes sense. You can change +the resolution of the output in several ways. From this program, you +can specify the output resolution with --mni-out-res (options are +1.0mm, 1.5mm, 2.0mm). You can also create an LTA file to the desired +output resolution and then use mri_warp_convert to change the warp. +Or you can create an LTA file to the desired resolution, and then use +mri_vol2vol --gcam to apply the warp and the LTA (instead of +mri_convert, which will only the the warp). + +Affine and affine-only: by default, this script will perform a +non-linear registration to the target space. The first step of this +will be to perform an affine registration. Sometimes, this is useful +too. This affine registration is stored in reg.targ_to_invol.lta. If +you ONLY want the affine registration, then you can specify +--affine-only. + +Arbitrary Targets: you can use fs-synthmorph-reg to register to target +volumes other than the mni152. To do this, simply specify --t +targetvol. You can have fs-synthmorph-reg skull strip the target by +adding --strip-target. By default, the target will be cropped unless +--no-crop (see Cropping above). + +Please cite SynthMorph: +Anatomy-specific acquisition-agnostic affine registration learned from +fictitious images +Hoffmann M, Hoopes A, Fischl B*, Dalca AV* (*equal contribution) +SPIE Medical Imaging: Image Processing, 12464, p 1246402, 2023 +https://doi.org/10.1117/12.2653251 +https://synthmorph.io/#papers (PDF) + +SynthMorph: learning contrast-invariant registration without acquired images +Hoffmann M, Billot B, Greve DN, Iglesias JE, Fischl B, Dalca AV +IEEE Transactions on Medical Imaging, 41 (3), 543-558, 2022 +https://doi.org/10.1109/TMI.2021.3116879 + +If you use SynthStrip in your analysis, please cite: +SynthStrip: Skull-Stripping for Any Brain Image +A Hoopes, JS Mora, AV Dalca, B Fischl, M Hoffmann +NeuroImage 206 (2022), 119474 +https://doi.org/10.1016/j.neuroimage.2022.119474 +Website: https://synthstrip.io + diff --git a/mri_synthmorph/mri_synthmorph b/mri_synthmorph/mri_synthmorph new file mode 100644 index 0000000..bfee4e5 --- /dev/null +++ b/mri_synthmorph/mri_synthmorph @@ -0,0 +1,419 @@ +#!/usr/bin/env python3 + +import os +import sys +import pathlib +import argparse +import textwrap +import surfa as sf +from synthmorph import utils + + +# Argument settings. +default = { + 'model': 'joint', + 'hyper': 0.5, + 'extent': 256, + 'steps': 7, + 'method': 'linear', + 'type': 'float32', + 'fill': 0, +} +choices = { + 'model': ('joint', 'deform', 'affine', 'rigid'), + 'extent': (192, 256), + 'method': ('linear', 'nearest'), + 'type': ('uint8', 'uint16', 'int16', 'int32', 'float32'), +} +limits = { + 'steps': 5, +} +resolve = ('model', 'method') + + +# Documentation. +n = '\033[0m' if sys.stdout.isatty() else '' +b = '\033[1m' if sys.stdout.isatty() else '' +u = '\033[4m' if sys.stdout.isatty() else '' +prog = os.path.basename(sys.argv[0]) + + +# References. +ref = f''' +SynthMorph: learning contrast-invariant registration without acquired images\t +Hoffmann M, Billot B, Greve DN, Iglesias JE, Fischl B, Dalca AV\t +IEEE Transactions on Medical Imaging, 41 (3), 543-558, 2022\t +https://doi.org/10.1109/TMI.2021.3116879 + +Anatomy-specific acquisition-agnostic {u}affine{n} registration learned from fictitious images\t +Hoffmann M, Hoopes A, Fischl B*, Dalca AV* (*equal contribution)\t +SPIE Medical Imaging: Image Processing, 12464, 1246402, 2023\t +https://doi.org/10.1117/12.2653251\t +https://synthmorph.io/#papers (PDF) + +Anatomy-aware and acquisition-agnostic {u}joint{n} registration with SynthMorph\t +Hoffmann M, Hoopes A, Greve DN, Fischl B*, Dalca AV* (*equal contribution)\t +Imaging Neuroscience, 2, 1-33, 2024\t +https://doi.org/10.1162/imag_a_00197 + +Website: https://synthmorph.io +''' + +help_general = f'''{prog} + +{b}NAME{n} + {b}{prog}{n} - register 3D brain images without preprocessing + +{b}SYNOPSIS{n} + {b}{prog}{n} [-h] {u}command{n} [options] + +{b}DESCRIPTION{n} + SynthMorph is a deep-learning tool for symmetric, acquisition-agnostic + registration of single-frame brain MRI of any geometry. The + registration is anatomy-aware, removing the need for skull-stripping, + and you can control the warp smoothness. + + Pass an option or {u}command{n} from the following list. You can omit + trailing characters, as long as there is no ambiguity. + + {b}register{n} + Register 3D brain images without preprocessing. + + {b}apply{n} + Apply an existing transform to another 3D image or label map. + + {b}-h{n} + Print this help text and exit. + +{b}IMAGE FORMAT{n} + The registration supports single-frame image volumes of any size, + resolution, and orientation. The moving and the fixed image geometries + can differ. The accepted image file formats are: MGH (.mgz) and NIfTI + (.nii.gz, .nii). + + Internally, the registration converts image buffers to: isotropic 1-mm + voxels, intensities min-max normalized into the interval [0, 1], and + left-inferior-anterior (LIA) axes. This conversion requires intact + image-to-world matrices. That is, the head must have the correct + anatomical orientation in a viewer like {b}freeview{n}. + +{b}TRANSFORM FORMAT{n} + SynthMorph transformations operate in physical RAS space. We save + matrix transforms as text in LTA format (.lta) and displacement fields + as images with three frames indicating shifts in RAS direction. + +{b}ENVIRONMENT{n} + The following environment variables affect {b}{prog}{n}: + + SUBJECTS_DIR + Ignored unless {b}{prog}{n} runs inside a container. Mounts the + host directory SUBJECTS_DIR to {u}/mnt{n} inside the container. + Defaults to the current working directory. + +{b}SEE ALSO{n} + For converting, composing, and applying transforms, consider FreeSurfer + tools {b}lta_convert{n}, {b}mri_warp_convert{n}, + {b}mri_concatenate_lta{n}, {b}mri_concatenate_gcam{n}, and + {b}mri_convert{n}. + +{b}CONTACT{n} + Reach out to freesurfer@nmr.mgh.harvard.edu or at + https://voxelmorph.net. + +{b}REFERENCES{n} + If you use SynthMorph in a publication, please cite us! +''' + textwrap.indent(ref, prefix=' ' * 8) + +help_register = f'''{prog}-register + +{b}NAME{n} + {b}{prog}-register{n} - register 3D brain images without preprocessing + +{b}SYNOPSIS{n} + {b}{prog} register{n} [options] {u}moving{n} {u}fixed{n} + +{b}DESCRIPTION{n} + SynthMorph is a deep-learning tool for symmetric, acquisition-agnostic + registration of brain MRI with any volume size, resolution, and + orientation. The registration is anatomy-aware, removing the need for + skull-stripping, and you can control the warp smoothness. + + SynthMorph registers a {u}moving{n} (source) image to a {u}fixed{n} + (target) image. Their geometries can differ. The options are as + follows: + + {b}-m{n} {u}model{n} + Transformation model ({', '.join(choices['model'])}). Defaults + to {default['model']}. Joint includes affine and deformable but + differs from running both in sequence in that it applies the + deformable step in an affine mid-space to guarantee symmetric + joint transforms. Deformable assumes prior affine alignment or + initialization with {b}-i{n}. + + {b}-o{n} {u}image{n} + Save {u}moving{n} registered to {u}fixed{n}. + + {b}-O{n} {u}image{n} + Save {u}fixed{n} registered to {u}moving{n}. + + {b}-H{n} + Update the voxel-to-world matrix instead of resampling when + saving images with {b}-o{n} and {b}-O{n}. For matrix transforms + only. Not all software supports headers with shear from affine + registration. + + {b}-t{n} {u}trans{n} + Save the transform from {u}moving{n} to {u}fixed{n}, including + any initialization. + + {b}-T{n} {u}trans{n} + Save the transform from {u}fixed{n} to {u}moving{n}, including + any initialization. + + {b}-i{n} {u}trans{n} + Apply an initial matrix transform to {u}moving{n} before the + registration. + + {b}-M{n} + Apply half the initial matrix transform to {u}moving{n} and + (the inverse of) the other half to {u}fixed{n}, for symmetry. + This will make running the deformable after an affine step + equivalent to joint registration. Requires {b}-i{n}. + + {b}-j{n} {u}threads{n} + Number of TensorFlow threads. System default if unspecified. + + {b}-g{n} + Use the GPU in environment variable CUDA_VISIBLE_DEVICES or GPU + 0 if the variable is unset or empty. + + {b}-r{n} {u}lambda{n} + Regularization parameter in the open interval (0, 1) for + deformable registration. Higher values lead to smoother warps. + Defaults to {default['hyper']}. + + {b}-n{n} {u}steps{n} + Integration steps for deformable registration. Lower numbers + improve speed and memory use but can lead to inaccuracies and + folding voxels. Defaults to {default['steps']}. Should not be + less than {limits['steps']}. + + {b}-e{n} {u}extent{n} + Isotropic extent of the registration space in unit voxels + {choices['extent']}. Lower values improve speed and memory use + but may crop the anatomy of interest. Defaults to + {default['extent']}. + + {b}-w{n} {u}weights{n} + Use alternative model weights, exclusively. Repeat the flag + to set affine and deformable weights for joint registration, + or the result will disappoint. + + {b}-h{n} + Print this help text and exit. + +{b}ENVIRONMENT{n} + The following environment variables affect {b}{prog}-register{n}: + + CUDA_VISIBLE_DEVICES + Use a specific GPU. If unset or empty, passing {b}-g{n} will + select GPU 0. Ignored without {b}-g{n}. + + FREESURFER_HOME + Load model weights from directory {u}FREESURFER_HOME/models{n}. + Ignored when specifying weights with {b}-w{n}. + +{b}EXAMPLES{n} + Joint affine-deformable registration, saving the moved image: + # {prog} register -o out.nii mov.nii fix.nii + + Joint registration at 25% warp smoothness: + # {prog} register -r 0.25 -o out.nii mov.nii fix.nii + + Affine registration saving the transform: + # {prog} register -m affine -t aff.lta mov.nii.gz fix.nii.gz + + Deformable registration only, assuming prior affine alignment: + # {prog} register -m deform -t def.mgz mov.mgz fix.mgz + + Deformable step initialized with an affine transform: + # {prog} reg -m def -i aff.lta -o out.mgz mov.mgz fix.mgz + + Rigid registration, setting the output image header (no resampling): + # {prog} register -m rigid -Ho out.mgz mov.mgz fix.mgz +''' + +help_apply = f'''{prog}-apply + +{b}NAME{n} + {b}{prog}-apply{n} - apply an existing SynthMorph transform + +{b}SYNOPSIS{n} + {b}{prog} apply{n} [options] {u}trans{n} {u}image{n} {u}output{n} + [{u}image{n} {u}output{n} ...] + +{b}DESCRIPTION{n} + Apply a spatial transform {u}trans{n} estimated by SynthMorph to a 3D + {u}image{n} and write the result to {u}output{n}. You can pass any + number of image-output pairs to be processed in the same way. + + The following options identically affect all input-output pairs. + + {b}-H{n} + Update the voxel-to-world matrix of {u}output{n} instead of + resampling. For matrix transforms only. Not all software and + file formats support headers with shear from affine + registration. + + {b}-m{n} {u}method{n} + Interpolation method ({', '.join(choices['method'])}). Defaults + to {default['method']}. Choose linear for images and nearest + for label (segmentation) maps. + + {b}-t{n} {u}type{n} + Output data type ({', '.join(choices['type'])}). Defaults to + {default['type']}. Casting to a type other than + {default['type']} after linear interpolation may result in + information loss. + + {b}-f{n} {u}fill{n} + Extrapolation fill value for areas outside the field-of-view of + {u}image{n}. Defaults to {default['fill']}. + + {b}-h{n} + Print this help text and exit. + +{b}EXAMPLES{n} + Apply an affine transform to an image: + # {prog} apply affine.lta image.nii out.nii.gz + + Apply a warp to an image, saving the output in floating-point format: + # {prog} apply -t float32 warp.nii image.nii out.nii + + Apply the same transform to two images: + # {prog} app warp.mgz im_1.mgz out_1.mgz im_2.mgz out_2.mgz + + Transform a label map: + # {prog} apply -m nearest warp.nii labels.nii out.nii +''' + + +# Command-line parsing. +p = argparse.ArgumentParser() +p.format_help = lambda: utils.rewrap_text(help_general, end='\n\n') + +sub = p.add_subparsers(dest='command') +commands = {f: sub.add_parser(f) for f in ('register', 'apply')} + + +def add_flags(f): + out = {} + + if f in choices: + out.update(choices=choices[f]) + if f in resolve: + out.update(type=lambda x: utils.resolve_abbrev(x, strings=choices[f])) + + if f in default: + out.update(default=default[f]) + + return out + + +# Registration arguments. +r = commands['register'] +r.add_argument('moving') +r.add_argument('fixed') +r.add_argument('-m', dest='model', **add_flags('model')) +r.add_argument('-o', dest='out_moving', metavar='image') +r.add_argument('-O', dest='out_fixed', metavar='image') +r.add_argument('-H', dest='header_only', action='store_true') +r.add_argument('-t', dest='trans', metavar='trans') +r.add_argument('-T', dest='inverse', metavar='trans') +r.add_argument('-i', dest='init', metavar='trans') +r.add_argument('-M', dest='mid_space', action='store_true') +r.add_argument('-j', dest='threads', metavar='threads', type=int) +r.add_argument('-g', dest='gpu', action='store_true') +r.add_argument('-r', dest='hyper', metavar='lambda', type=float, **add_flags('hyper')) +r.add_argument('-n', dest='steps', metavar='steps', type=int, **add_flags('steps')) +r.add_argument('-e', dest='extent', type=int, **add_flags('extent')) +r.add_argument('-w', dest='weights', metavar='weights', action='append') +r.add_argument('-v', dest='verbose', action='store_true') +r.add_argument('-d', dest='out_dir', metavar='dir', type=pathlib.Path) +r.format_help = lambda: utils.rewrap_text(help_register, end='\n\n') + +# Transformation arguments. +a = commands['apply'] +a.add_argument('trans') +a.add_argument('pairs', metavar='image output', nargs='+') +a.add_argument('-H', dest='header_only', action='store_true') +a.add_argument('-m', dest='method', **add_flags('method')) +a.add_argument('-t', dest='type', **add_flags('type')) +a.add_argument('-f', dest='fill', metavar='fill', type=float, **add_flags('fill')) +a.format_help = lambda: utils.rewrap_text(help_apply, end='\n\n') + + +# Parse arguments. +if len(sys.argv) == 1: + p.print_usage() + exit(0) + +# Command resolution. +c = sys.argv[1] = utils.resolve_abbrev(sys.argv[1], commands) +if len(sys.argv) == 2 and c in commands: + commands[c].print_usage() + exit(0) + +# Default command. +if c not in (*commands, '-h'): + sys.argv.insert(1, 'register') + +arg = p.parse_args() + + +# Command. +if arg.command == 'register': + + # Argument checking. + if arg.header_only and not arg.model in ('affine', 'rigid'): + sf.system.fatal('-H is not compatible with deformable registration') + + if arg.mid_space and not arg.init: + sf.system.fatal('-M requires matrix initialization') + + if not 0 < arg.hyper < 1: + sf.system.fatal('regularization strength not in open interval (0, 1)') + + if arg.steps < limits['steps']: + sf.system.fatal('too few integration steps') + + # TensorFlow setup. + gpu = os.environ.get('CUDA_VISIBLE_DEVICES', '0') + os.environ['CUDA_VISIBLE_DEVICES'] = gpu if arg.gpu else '' + os.environ['TF_CPP_MIN_LOG_LEVEL'] = '0' if arg.verbose else '3' + os.environ['NEURITE_BACKEND'] = 'tensorflow' + os.environ['VXM_BACKEND'] = 'tensorflow' + from synthmorph import registration + + registration.register(arg) + + +if arg.command == 'apply': + + # Argument checking. + which = 'affine' if arg.trans.endswith('.lta') else 'warp' + if arg.header_only and which == 'warp': + sf.system.fatal('-H is not compatible with deformable transforms') + + if len(arg.pairs) % 2: + sf.system.fatal('list of input-output pairs not of even length') + + # Transform. + trans = getattr(sf, f'load_{which}')(arg.trans) + prop = dict(method=arg.method, resample=not arg.header_only, fill=arg.fill) + for inp, out in zip(arg.pairs[::2], arg.pairs[1::2]): + sf.load_volume(inp).transform(trans, **prop).astype(arg.type).save(out) + + +print('Thank you for choosing SynthMorph. Please cite us!') +print(utils.rewrap_text(ref)) diff --git a/mri_synthmorph/mri_synthmorph_apply b/mri_synthmorph/mri_synthmorph_apply new file mode 100644 index 0000000..e6b1d09 --- /dev/null +++ b/mri_synthmorph/mri_synthmorph_apply @@ -0,0 +1,237 @@ +#!/usr/bin/env python3 + +import os +import sys +import shutil +import textwrap +import argparse +import surfa as sf + + +# Settings. +default = { + 'method': 'linear', + 'fill': 0, + 'type': 'float32', +} +choices = { + 'method': ('linear', 'nearest'), + 'type': ('uint8', 'uint16', 'int16', 'int32', 'float32'), +} + + +def rewrap(text, width=None, hard='\t\n', hard_indent=0): + """Rewrap text such that lines fill the available horizontal space. + + Reformats individual paragraphs of a text body, considering subsequent + lines with identical indentation as paragraphs. For unspecified width, the + function will attempt to determine the extent of the terminal. + + Parameters + ---------- + text : str + Text to rewrap. + width : int, optional + Maximum line width. None means the width of the terminal as determined + by `textwrap`, defaulting to 80 characters for background processes. + hard : str, optional + String interpreted as a hard break when terminating a line. Useful for + inserting a line break without changing the indentation level. Must end + with a line break and will be removed from the output text. + hard_indent : int, optional + Number of additional whitespace characters by which to indent the lines + following a hard break. See `hard`. + + Returns + ------- + out : str + Reformatted text. + + """ + # Inputs. + if width is None: + width = shutil.get_terminal_size().columns + lines = text.splitlines(keepends=True) + + # Merge lines to paragraphs. + pad = [] + pad_hard = [] + par = [] + for i, line in enumerate(lines): + ind = len(line) - len(line.lstrip()) + if i == 0 or ind != pad[-1] or lines[i - 1].endswith(hard): + par.append('') + pad.append(ind) + pad_hard.append(ind) + + if line.endswith(hard): + line = line.replace(hard, '\n') + pad_hard[-1] += hard_indent + par[-1] += line[ind:] + + # Reformat paragraphs. + for i, _ in enumerate(par): + par[i] = textwrap.fill( + par[i], width, + initial_indent=' ' * pad[i], subsequent_indent=' ' * pad_hard[i], + ) + + return '\n'.join(par) + + +# Documentation. +n = '\033[0m' if sys.stdout.isatty() else '' +b = '\033[1m' if sys.stdout.isatty() else '' +u = '\033[4m' if sys.stdout.isatty() else '' +prog = os.path.basename(sys.argv[0]) +doc = f'''{prog} + +{b}NAME{n} + {b}{prog}{n} - apply a SynthMorph transform to 3D images + +{b}SYNOPSIS{n} + {b}{prog}{n} [options] {u}trans{n} {u}image{n} {u}output{n} + [{u}image{n} {u}output{n} ...] + +{b}DESCRIPTION{n} + Apply a spatial transform {u}trans{n} estimated by SynthMorph to a 3D + {u}image{n} and write the result to {u}output{n}. You can pass any + number of image-output pairs to be processed in the same way. + + The following options identically affect all image-output pairs. + + {b}-H{n} + Update the voxel-to-world matrix of the output image instead of + resampling. For matrix transforms only. Not all software and + file formats support headers with shear from affine + registration. + + {b}-i{n} {u}method{n} + Interpolation method ({', '.join(choices['method'])}). Defaults + to {default['method']}. Choose linear for images and nearest + for label (segmentation) maps. + + {b}-t{n} {u}type{n} + Output data type ({', '.join(choices['type'])}). Defaults to + {default['type']}. Casting to a narrower type can result in + information loss. + + {b}-f{n} {u}fill{n} + Extrapolation fill value for areas outside the field-of-view of + the input image. Defaults to {default['fill']}. + + {b}-h{n} + Print this help text and exit. + +{b}IMAGE FORMAT{n} + Accepted file formats include: MGH (.mgz) and NIfTI (.nii.gz, .nii). + +{b}TRANSFORMS{n} + Refer to the help text of the registration utility for information on + transform file formats. + + For converting, composing, and applying transforms, consider the + FreeSurfer tools lta_convert, mri_warp_convert, mri_concatenate_lta, + mri_concatenate_gcam, mri_convert, mri_info. + +{b}ENVIRONMENT{n} + The following environment variables affect {b}{prog}{n}: + + SUBJECTS_DIR + Ignored unless {b}{prog}{n} runs inside a container. Mount the + host directory SUBJECTS_DIR to {u}/mnt{n} inside the container. + Defaults to the current working directory. + +{b}EXAMPLES{n} + Apply an affine transform to an image: + # {prog} affine.lta image.nii out.nii.gz + + Apply a warp to an image, saving the output in floating-point format: + # {prog} -t float32 warp.mgz image.mgz out.mgz + + Apply a transform to each of two images: + # {prog} warp.mgz image_1.mgz out_1.mgz image_2.mgz out_2.mgz + + Transform a label map: + # {prog} -i nearest warp.mgz labels.mgz out.mgz + +{b}CONTACT{n} + Reach out to freesurfer@nmr.mgh.harvard.edu or at + https://github.com/voxelmorph/voxelmorph. + +{b}REFERENCES{n} + If you use SynthMorph in a publication, please cite us! +''' + + +# References. +ref = ''' +SynthMorph: learning contrast-invariant registration without acquired images\t +Hoffmann M, Billot B, Greve DN, Iglesias JE, Fischl B, Dalca AV\t +IEEE Transactions on Medical Imaging, 41 (3), 543-558, 2022\t +https://doi.org/10.1109/TMI.2021.3116879 + +Anatomy-specific acquisition-agnostic affine registration learned from fictitious images\t +Hoffmann M, Hoopes A, Fischl B*, Dalca AV* (*equal contribution)\t +SPIE Medical Imaging: Image Processing, 12464, 1246402, 2023\t +https://doi.org/10.1117/12.2653251\t +https://synthmorph.io/#papers (PDF) + +Anatomy-aware and acquisition-agnostic joint registration with SynthMorph\t +Hoffmann M, Hoopes A, Greve DN, Fischl B*, Dalca AV* (*equal contribution)\t +Imaging Neuroscience, 2, 1-33, 2024\t +https://doi.org/10.1162/imag_a_00197 + +Website: https://synthmorph.io +''' +doc += textwrap.indent(ref, prefix=' ' * 8) + + +print(rewrap(( + f'Warning: {prog} is deprecated in favor of `mri_synthmorph apply` and ' + 'will be removed in the future.' +))) + + +# Arguments. +p = argparse.ArgumentParser(add_help=False) +p.add_argument('trans') +p.add_argument('pairs', metavar='image output', nargs='+') +p.add_argument('-H', dest='header_only', action='store_true') +p.add_argument('-i', dest='method', choices=choices['method'], default=default['method']) +p.add_argument('-t', dest='type', choices=choices['type'], default=default['type']) +p.add_argument('-f', dest='fill', metavar='fill', type=float, default=default['fill']) +p.add_argument('-h', action='store_true') + + +# Help. +if len(sys.argv) == 1: + p.print_usage() + exit(0) + +if any(f[0] == '-' and 'h' in f for f in sys.argv): + print(rewrap(doc), end='\n\n') + exit(0) + + +# Parsing. +arg = p.parse_args() + +if len(arg.pairs) % 2: + sf.system.fatal('did not receive even-length list of input-output pairs') + + +# Transform. +f = 'affine' if arg.trans.endswith('.lta') else 'warp' +trans = getattr(sf, f'load_{f}')(arg.trans) + + +# Application. +pairs = zip(arg.pairs[::2], arg.pairs[1::2]) +prop = dict(method=arg.method, resample=not arg.header_only, fill=arg.fill) +for inp, out in pairs: + sf.load_volume(inp).transform(trans, **prop).astype(arg.type).save(out) + + +print('Thank you for choosing SynthMorph. Please cite us!') +print(rewrap(ref)) diff --git a/mri_synthmorph/synthmorph/registration.py b/mri_synthmorph/synthmorph/registration.py new file mode 100644 index 0000000..7351fa1 --- /dev/null +++ b/mri_synthmorph/synthmorph/registration.py @@ -0,0 +1,313 @@ +import os +import h5py +import numpy as np +import surfa as sf +import tensorflow as tf +import voxelmorph as vxm + + +# Settings. +weights = { + 'joint': ('synthmorph.affine.2.h5', 'synthmorph.deform.3.h5',), + 'deform': ('synthmorph.deform.3.h5',), + 'affine': ('synthmorph.affine.2.h5',), + 'rigid': ('synthmorph.rigid.1.h5',), +} + + +def network_space(im, shape, center=None): + """Construct transform from network space to the voxel space of an image. + + Constructs a coordinate transform from the space the network will operate + in to the zero-based image index space. The network space has isotropic + 1-mm voxels, left-inferior-anterior (LIA) orientation, and no shear. It is + centered on the field of view, or that of a reference image. This space is + an indexed voxel space, not world space. + + Parameters + ---------- + im : surfa.Volume + Input image to construct the transform for. + shape : (3,) array-like + Spatial shape of the network space. + center : surfa.Volume, optional + Center the network space on the center of a reference image. + + Returns + ------- + out : tuple of (3, 4) NumPy arrays + Transform from network to input-image space and its inverse, thinking + coordinates. + + """ + old = im.geom + new = sf.ImageGeometry( + shape=shape, + voxsize=1, + rotation='LIA', + center=old.center if center is None else center.geom.center, + shear=None, + ) + + net_to_vox = old.world2vox @ new.vox2world + vox_to_net = new.world2vox @ old.vox2world + return net_to_vox.matrix, vox_to_net.matrix + + +def transform(im, trans, shape=None, normalize=False, batch=False): + """Apply a spatial transform to 3D image voxel data in dimensions. + + Applies a transformation matrix operating in zero-based index space or a + displacement field to an image buffer. + + Parameters + ---------- + im : surfa.Volume or NumPy array or TensorFlow tensor + Input image to transform, without batch dimension. + trans : array-like + Transform to apply to the image. A matrix of shape (3, 4), a matrix + of shape (4, 4), or a displacement field of shape (*space, 3), + without batch dimension. + shape : (3,) array-like, optional + Output shape used for converting matrices to dense transforms. None + means the shape of the input image will be used. + normalize : bool, optional + Min-max normalize the image intensities into the interval [0, 1]. + batch : bool, optional + Prepend a singleton batch dimension to the output tensor. + + Returns + ------- + out : float TensorFlow tensor + Transformed image with a trailing feature dimension. + + """ + # Add singleton feature dimension if needed. + if tf.rank(im) == 3: + im = im[..., tf.newaxis] + + out = vxm.utils.transform( + im, trans, fill_value=0, shift_center=False, shape=shape, + ) + + if normalize: + out -= tf.reduce_min(out) + out /= tf.reduce_max(out) + + if batch: + out = out[tf.newaxis, ...] + + return out + + +def load_weights(model, weights): + """Load weights into model or submodel. + + Attempts to load (all) weights into a model or one of its submodels. If + that fails, `model` may be a submodel of what we got weights for, and we + attempt to load the weights of a submodel (layer) into `model`. + + Parameters + ---------- + model : TensorFlow model + Model to initialize. + weights : str or pathlib.Path + Path to weights file. + + Raises + ------ + ValueError + If unsuccessful at loading any weights. + + """ + # Extract submodels. + models = [model] + i = 0 + while i < len(models): + layers = [f for f in models[i].layers if isinstance(f, tf.keras.Model)] + models.extend(layers) + i += 1 + + # Add models wrapping a single model in case this was done in training. + # Requires list expansion or Python will get stuck. + models.extend([tf.keras.Model(m.inputs, m(m.inputs)) for m in models]) + + # Attempt to load all weights into one of the models. + for mod in models: + try: + mod.load_weights(weights) + return + except ValueError as e: + pass + + # Assume `model` is a submodel of what we got weights for. + with h5py.File(weights, mode='r') as h5: + layers = h5.attrs['layer_names'] + weights = [list(h5[lay].attrs['weight_names']) for lay in layers] + + # Layers with weights. Attempt loading. + layers, weights = zip(*filter(lambda f: f[1], zip(layers, weights))) + for lay, wei in zip(layers, weights): + try: + model.set_weights([h5[lay][w] for w in wei]) + return + except ValueError as e: + if lay is layers[-1]: + raise e + + +def register(arg): + + # Parse arguments. + in_shape = (arg.extent,) * 3 + is_mat = arg.model in ('affine', 'rigid') + + # Threading. + if arg.threads: + tf.config.threading.set_inter_op_parallelism_threads(arg.threads) + tf.config.threading.set_intra_op_parallelism_threads(arg.threads) + + # Input data. + mov = sf.load_volume(arg.moving) + fix = sf.load_volume(arg.fixed) + if not len(mov.shape) == len(fix.shape) == 3: + sf.system.fatal('input images are not single-frame volumes') + + # Transforms between native voxel and network coordinates. Voxel and + # network spaces differ for each image. The networks expect isotropic 1-mm + # LIA spaces. Center these on the original images, except in the deformable + # case: it assumes prior affine registration, so we center the moving + # network space on the fixed image, to take into account affine transforms + # via resampling, updating the header, or passed on the command line alike. + center = fix if arg.model == 'deform' else None + net_to_mov, mov_to_net = network_space(mov, shape=in_shape, center=center) + net_to_fix, fix_to_net = network_space(fix, shape=in_shape) + + # Coordinate transforms from and to world space. There is only one world. + mov_to_ras = mov.geom.vox2world.matrix + fix_to_ras = fix.geom.vox2world.matrix + ras_to_mov = mov.geom.world2vox.matrix + ras_to_fix = fix.geom.world2vox.matrix + + # Incorporate an initial matrix transform from moving to fixed coordinates, + # as LTAs store the inverse. For mid-space initialization, compute the + # square root of the transform between fixed and moving network space. + if arg.init: + init = sf.load_affine(arg.init).convert(space='voxel') + if init.ndim != 3 \ + or not sf.transform.image_geometry_equal(mov.geom, init.source, tol=1e-3) \ + or not sf.transform.image_geometry_equal(fix.geom, init.target, tol=1e-3): + sf.system.fatal('initial transform geometry does not match images') + + init = fix_to_net @ init @ net_to_mov + if arg.mid_space: + init = tf.linalg.sqrtm(init) + if np.any(np.isnan(init)): + sf.system.fatal(f'cannot compute matrix square root of {arg.init}') + net_to_fix = net_to_fix @ init + fix_to_net = np.linalg.inv(net_to_fix) + + net_to_mov = net_to_mov @ tf.linalg.inv(init) + mov_to_net = np.linalg.inv(net_to_mov) + + # Take the input images to network space. When saving the moving image with + # the correct voxel-to-RAS matrix after incorporating an initial transform, + # an image viewer taking this matrix into account will show an unchanged + # image. The networks only see the voxel data, which have been moved. + inputs = ( + transform(mov, net_to_mov, shape=in_shape, normalize=True, batch=True), + transform(fix, net_to_fix, shape=in_shape, normalize=True, batch=True), + ) + + # Network. For deformable-only registration, `HyperVxmJoint` ignores the + # `mid_space` argument, and the initialization will determine the space. + prop = dict(in_shape=in_shape, bidir=True) + if is_mat: + prop.update(make_dense=False, rigid=arg.model == 'rigid') + model = vxm.networks.VxmAffineFeatureDetector(**prop) + + else: + prop.update(mid_space=True, int_steps=arg.steps, skip_affine=arg.model == 'deform') + model = vxm.networks.HyperVxmJoint(**prop) + inputs = (tf.constant([arg.hyper]), *inputs) + + # Weights. + if not arg.weights: + fs = os.environ.get('FREESURFER_HOME') + if not fs: + sf.system.fatal('set environment variable FREESURFER_HOME or weights') + arg.weights = [os.path.join(fs, 'models', f) for f in weights[arg.model]] + + for f in arg.weights: + load_weights(model, weights=f) + + # Inference. The first transform maps from the moving to the fixed image, + # or equivalently, from fixed to moving coordinates. The second is the + # inverse. Convert transforms between moving and fixed network spaces to + # transforms between the original voxel spaces. + pred = tuple(map(tf.squeeze, model(inputs))) + fw, bw = pred + fw = vxm.utils.compose((net_to_mov, fw, fix_to_net), shift_center=False, shape=fix.shape) + bw = vxm.utils.compose((net_to_fix, bw, mov_to_net), shift_center=False, shape=mov.shape) + + # Associate image geometries with the transforms. LTAs store the inverse. + if is_mat: + fw, bw = bw, fw + fw = sf.Affine(fw, source=mov, target=fix, space='voxel') + bw = sf.Affine(bw, source=fix, target=mov, space='voxel') + format = dict(space='world') + + else: + fw = sf.Warp(fw, source=mov, target=fix, format=sf.Warp.Format.disp_crs) + bw = sf.Warp(bw, source=fix, target=mov, format=sf.Warp.Format.disp_crs) + format = dict(format=sf.Warp.Format.disp_ras) + + # Output transforms. + if arg.trans: + fw.convert(**format).save(arg.trans) + + if arg.inverse: + bw.convert(**format).save(arg.inverse) + + # Moved images. + if arg.out_moving: + mov.transform(fw, resample=not arg.header_only).save(arg.out_moving) + + if arg.out_fixed: + fix.transform(bw, resample=not arg.header_only).save(arg.out_fixed) + + # Outputs in network space. + if arg.out_dir: + arg.out_dir.mkdir(exist_ok=True) + + # Input images. + mov = sf.ImageGeometry(in_shape, vox2world=mov_to_ras @ net_to_mov) + fix = sf.ImageGeometry(in_shape, vox2world=fix_to_ras @ net_to_fix) + mov = sf.Volume(inputs[-2][0], geometry=fix if arg.init else mov) + fix = sf.Volume(inputs[-1][0], geometry=fix) + mov.save(filename=arg.out_dir / 'inp_1.nii.gz') + fix.save(filename=arg.out_dir / 'inp_2.nii.gz') + + fw, bw = pred + if is_mat: + fw, bw = bw, fw + fw = sf.Affine(fw, source=mov, target=fix, space='voxel') + bw = sf.Affine(bw, source=fix, target=mov, space='voxel') + ext = 'lta' + + else: + fw = sf.Warp(fw, source=mov, target=fix, format=sf.Warp.Format.disp_crs) + bw = sf.Warp(bw, source=fix, target=mov, format=sf.Warp.Format.disp_crs) + ext = 'nii.gz' + + # Transforms. + fw.convert(**format).save(filename=arg.out_dir / f'tra_1.{ext}') + bw.convert(**format).save(filename=arg.out_dir / f'tra_2.{ext}') + + # Moved images. + mov.transform(fw).save(filename=arg.out_dir / 'out_1.nii.gz') + fix.transform(bw).save(filename=arg.out_dir / 'out_2.nii.gz') + + vmpeak = sf.system.vmpeak() + if vmpeak is not None: + print(f'#@# mri_synthmorph: {arg.model}, threads: {arg.threads}, VmPeak: {vmpeak}') diff --git a/mri_synthmorph/synthmorph/utils.py b/mri_synthmorph/synthmorph/utils.py new file mode 100644 index 0000000..6ca415b --- /dev/null +++ b/mri_synthmorph/synthmorph/utils.py @@ -0,0 +1,93 @@ +import shutil +import textwrap + + +def resolve_abbrev(needle, strings, lower=False): + """Return a full-length string matching a substring from the beginning. + + Parameters + ---------- + needle : str + Substring of one of several `strings`. + strings : str or iterable of str + Full-length strings, one of which should begin with `needle`. + lower : bool, optional + Convert needle to lowercase before matching. + + Returns + ------- + str + String in `strings` that begins with `needle` if there is no ambiguity. + If there is not exactly one match, the function will return `needle`. + + """ + if isinstance(strings, str): + strings = [strings] + strings = tuple(strings) + + if lower: + needle = needle.lower() + + matches = [f for f in strings if f.startswith(needle)] + return matches[0] if len(matches) == 1 else needle + + +def rewrap_text(text, width=None, hard='\t\n', hard_indent=0, end=''): + """Rewrap text such that lines fill the available horizontal space. + + Reformats individual paragraphs of a text body, considering subsequent + lines with identical indentation as paragraphs. For unspecified width, the + function will attempt to determine the extent of the terminal. + + Parameters + ---------- + text : str + Text to rewrap. + width : int, optional + Maximum line width. None means the width of the terminal as determined + by `textwrap`, defaulting to 80 characters for background processes. + hard : str, optional + String interpreted as a hard break when terminating a line. Useful for + inserting a line break without changing the indentation level. Must end + with a line break and will be removed from the output text. + hard_indent : int, optional + Number of additional whitespace characters by which to indent the lines + following a hard break. See `hard`. + end : str, optional + Append to the reformatted text. + + Returns + ------- + out : str + Reformatted text. + + """ + # Inputs. + if width is None: + width = shutil.get_terminal_size().columns + lines = text.splitlines(keepends=True) + + # Merge lines to paragraphs. + pad = [] + pad_hard = [] + par = [] + for i, line in enumerate(lines): + ind = len(line) - len(line.lstrip()) + if i == 0 or ind != pad[-1] or lines[i - 1].endswith(hard): + par.append('') + pad.append(ind) + pad_hard.append(ind) + + if line.endswith(hard): + line = line.replace(hard, '\n') + pad_hard[-1] += hard_indent + par[-1] += line[ind:] + + # Reformat paragraphs. + for i, _ in enumerate(par): + par[i] = textwrap.fill( + par[i], width, + initial_indent=' ' * pad[i], subsequent_indent=' ' * pad_hard[i], + ) + + return '\n'.join(par) + end diff --git a/mri_synthmorph/test_apply.sh b/mri_synthmorph/test_apply.sh new file mode 100644 index 0000000..f77e6de --- /dev/null +++ b/mri_synthmorph/test_apply.sh @@ -0,0 +1,69 @@ +#!/usr/bin/env bash +. "$(dirname "$0")/../test.sh" + +t() { test_command mri_synthmorph "$@" ; } + +# image interpolation +t apply affine.lta moving.mgz out.mgz +compare_vol out.mgz affine.mgz --thresh 0.02 --res-thresh 1e-3 --geo-thresh 1e-3 + +# NIfTI format +t apply affine.lta moving.mgz out.nii.gz +compare_vol out.nii.gz affine.mgz --thresh 0.02 --res-thresh 1e-3 --geo-thresh 1e-3 + +# dense transform +t apply identity.mgz moving.mgz out.mgz -t uint8 +compare_vol out.mgz moving.mgz --res-thresh 1e-3 --geo-thresh 1e-3 + +# matrix update +t apply -H rigid.lta -t uint8 moving.mgz out.mgz +compare_vol out.mgz header.mgz --res-thresh 1e-3 --geo-thresh 1e-3 + +# label interpolation +t apply -m nearest -t uint8 affine.lta labels.moving.mgz out.mgz +compare_vol out.mgz labels.affine.mgz --res-thresh 1e-3 --geo-thresh 1e-3 + +# multiple input pairs +t apply rigid.lta moving.mgz out_1.mgz moving.mgz out_2.mgz +compare_vol out_1.mgz rigid.mgz --thresh 0.02 --res-thresh 1e-3 --geo-thresh 1e-3 +compare_vol out_2.mgz rigid.mgz --thresh 0.02 --res-thresh 1e-3 --geo-thresh 1e-3 + +# data types, command abbreviation +t a -t uint8 rigid.lta moving.mgz out.mgz +run_comparison mri_info --type out.mgz | grep -Fx uchar + +t ap -t uint16 rigid.lta moving.mgz out.mgz +run_comparison mri_info --type out.mgz | grep -Fx ushrt + +t app -t int16 rigid.lta moving.mgz out.mgz +run_comparison mri_info --type out.mgz | grep -Fx short + +t appl -t int32 rigid.lta moving.mgz out.mgz +run_comparison mri_info --type out.mgz | grep -Fx int + +t apply -t float32 rigid.lta moving.mgz out.mgz +run_comparison mri_info --type out.mgz | grep -Fx float + +# method abbreviation +FSTEST_NO_DATA_RESET=1 +for method in l li lin line linea linear n ne nea near neare neares nearest; do + t apply -m "$method" rigid.lta moving.mgz out.mgz +done + +# usage, help +t +t -h +t apply +t apply -h + +# NIfTI warp +FSTEST_NO_DATA_RESET=1 +mri_convert identity.mgz identity.nii.gz +t apply identity.nii.gz moving.mgz out.mgz -t uint8 +compare_vol out.mgz moving.mgz --res-thresh 1e-3 --geo-thresh 1e-3 + +# illegal arguments +EXPECT_FAILURE=1 +t slice +t apply -H identity.mgz moving.mgz fixed.mgz +t apply affine.lta moving.mgz out.mgz odd-number-of-io-pairs.mgz diff --git a/mri_synthmorph/test_register.sh b/mri_synthmorph/test_register.sh new file mode 100644 index 0000000..fb0a4e7 --- /dev/null +++ b/mri_synthmorph/test_register.sh @@ -0,0 +1,83 @@ +#!/usr/bin/env bash +. "$(dirname "$0")/../test.sh" + +t() { test_command mri_synthmorph "$@" ; } + +# affine registration +t -m affine -o out.mgz moving.mgz fixed.mgz +compare_vol out.mgz affine.mgz --thresh 0.02 + +# affine symmetry +t -m aff -O out.mgz fixed.mgz moving.mgz +compare_vol out.mgz affine.mgz --thresh 0.02 + +# affine inverse consistency +t -m a -t out.lta -T inv.lta moving.mgz fixed.mgz +lta_diff out.lta inv.lta --invert2 | awk 'END {print $0; exit !($0<1e-3)}' + +# output directory creation +t -ma -d outputs fixed.mgz moving.mgz +[ -d outputs ] + +# rigid inverse consistency +t -m rigid -t out.lta -T inv.lta moving.mgz fixed.mgz +lta_diff out.lta inv.lta --invert2 | awk 'END {print $0; exit !($0<1e-3)}' + +# rigid registration +t -m rig -o out.mgz moving.mgz fixed.mgz +compare_vol out.mgz rigid.mgz --thresh 0.02 + +# geometry update +t -m r -Ho out.mgz moving.mgz fixed.mgz +compare_vol out.mgz header.mgz --thresh 0.02 --res-thresh 1e-3 --geo-thresh 1e-3 + +# deformable registration with initialization +t -m deform -i affine.lta -o out_1.mgz -O out_2.mgz moving.mgz fixed.mgz +compare_vol out_1.mgz deform_1.mgz --thresh 0.02 +compare_vol out_2.mgz deform_2.mgz --thresh 0.02 + +# deformable registration with mid-space initialization +t -m def -Mi affine.lta -o out.nii.gz moving.mgz fixed.mgz +compare_vol out.nii.gz deform_mid.nii.gz --thresh 0.02 --geo-thresh 1e-4 + +# joint registration +t -m joint -o out.mgz moving.mgz fixed.mgz +compare_vol out.mgz joint.mgz --thresh 0.02 + +# default model +t -o out.mgz moving.mgz fixed.mgz +compare_vol out.mgz joint.mgz --thresh 0.02 + +# help, usage, command abbreviation +FSTEST_NO_DATA_RESET=1 +t register -h +for cmd in r re reg regi regis registe register; do + t "$cmd" +done + +# deformable flags, explicit command +t register moving.mgz fixed.mgz -md -j16 -e256 -n7 -r0.5 +t register moving.mgz fixed.mgz -m j -j 16 -e 192 -n 5 -r 0.7 + +# NIfTI warps +FSTEST_NO_DATA_RESET=1 +mri_convert=$(find_path $FSTEST_CWD mri_convert/mri_convert) +t -t out.nii.gz moving.mgz fixed.mgz +test_command $mri_convert -odt float -at out.nii.gz moving.mgz out.mgz +compare_vol out.mgz joint.mgz --thresh 1 + +# displacements in RAS space +mri_warp_convert=$(find_path $FSTEST_CWD mri_warp_convert/mri_warp_convert) +test_command $mri_warp_convert -g moving.mgz --inras out.nii.gz --outmgzwarp out.mgz +test_command $mri_convert -odt float -at out.mgz moving.mgz moved.mgz +compare_vol moved.mgz joint.mgz --thresh 1 --geo-thresh 1e-4 + +# illegal arguments +EXPECT_FAILURE=1 +t moving.mgz fixed.mgz -m banana +t moving.mgz fixed.mgz -e 1 +t moving.mgz fixed.mgz -n 4 +t moving.mgz fixed.mgz -r 0 +t moving.mgz fixed.mgz -r 1 +t moving.mgz fixed.mgz -Hm deform +t moving.mgz fixed.mgz -Hm joint diff --git a/src/conda-create.sh b/src/conda-create.sh new file mode 100644 index 0000000..8ba1433 --- /dev/null +++ b/src/conda-create.sh @@ -0,0 +1,4 @@ +conda deactivate +conda create -y -n 25reg -c conda-forge simpleitk tensorflow-gpu +conda activate 25reg +pip install -r requirements.txt diff --git a/src/g4synthmorph.py b/src/g4synthmorph.py new file mode 100644 index 0000000..9b2ee9e --- /dev/null +++ b/src/g4synthmorph.py @@ -0,0 +1,376 @@ + +''' +Use SynthMorph to register G4 images + +https://download-directory.github.io/ +https://github.com/freesurfer/freesurfer/tree/dev/mri_synthmorph + + + +XLA_FLAGS=--xla_gpu_cuda_data_dir=/home/xfr/.conda/envs/25reg time ./mri_synthmorph -m affine -o ../test.nii.gz -g '/mnt/1218/Public/dataset2/M6/ZYRGTRKJ/20230728/MR/3D_SAG_T1_MPRAGE_+C_MPR_Tra_20230728143005_14.nii.gz' '/mnt/1218/Public/dataset2/M6/ZYRGTRKJ/20230728/CT/1.1_CyberKnife_head(MAR)_20230728111920_3.nii.gz' + +XLA_FLAGS=--xla_gpu_cuda_data_dir=/home/xfr/.conda/envs/25reg time mri_synthmorph/mri_synthmorph -m affine -o affine.nii.gz -g moving.nii.gz clipped.nii.gz + + +''' + +from pathlib import Path + +import argparse +import logging +import json +import os +# import pathlib +import shelve +import shutil +import time + +from skimage.metrics import normalized_mutual_information + +import filelock +import matplotlib.pyplot as plt +import numpy as np +import SimpleITK as sitk + +from mri_synthmorph.synthmorph import registration +# from synthmorph import registration + +import surfa as sf + +PATIENTS_ROOT = '/mnt/1220/Public/dataset2/G4' +OUT_ROOT = '/mnt/1220/Public/dataset2/G4-synthmorph' +SHELVE = os.path.join(OUT_ROOT, '0shelve') + +MAX_Y = 256 + +SIZE_X = 249 +SIZE_Y = 249 +SIZE_Z = 192 +# SIZE_Z = 256 + +MIN_OVERLAP = 0.50 +MIN_METRIC = -0.50 + + +logging.basicConfig( + level=logging.INFO, + format='%(asctime)s - %(levelname)s - %(message)s', + handlers=[ + logging.StreamHandler(), + logging.FileHandler('g4synthmorph.log') + ] +) +logger = logging.getLogger(__name__) + + +# def resize_with_crop_or_pad(image, tx = SIZE_X, ty = SIZE_Y, tz = SIZE_Z): +def resize_with_pad(image, tx = SIZE_X, ty = SIZE_Y, tz = SIZE_Z): + sx, sy, sz = image.GetSize() + l = [(tx-sx)//2, + (ty-sy)//2, + (tz-sz)//2,] + u = [tx-sx-l[0], + ty-sy-l[1], + tz-sz-l[2], + ] + # print (l, u) + return sitk.ConstantPad(image, l, u) + + +def draw_sitk(image, d, post): + a = sitk.GetArrayFromImage(image) + s = a.shape + + fig, axs = plt.subplots(1, 3) + # fig.suptitle('%dx%dx%d'%(s[2], s[1], s[0])) + axs.flat[0].imshow(a[s[0]//2,:,:], cmap='gray') + axs.flat[1].imshow(a[:,s[1]//2,:], cmap='gray') + axs.flat[2].imshow(a[:,:,s[2]//2], cmap='gray') + axs.flat[0].axis('off') + axs.flat[1].axis('off') + axs.flat[2].axis('off') + axs.flat[1].invert_yaxis() + axs.flat[2].invert_yaxis() + plt.tight_layout() + os.makedirs(d, exist_ok=True) + plt.savefig(os.path.join(d, '%dx%dx%d-%s'%(s[2],s[1],s[0],post))) + plt.close() + # exit() + +def bbox2_3D(img): + + r = np.any(img, axis=(1, 2)) + c = np.any(img, axis=(0, 2)) + z = np.any(img, axis=(0, 1)) + + if not np.any(r): + return -1, -1, -1, -1, -1, -1 + + rmin, rmax = np.where(r)[0][[0, -1]] + cmin, cmax = np.where(c)[0][[0, -1]] + zmin, zmax = np.where(z)[0][[0, -1]] + + return rmin, rmax, cmin, cmax, zmin, zmax + +''' +Namespace(command='register', moving='/nn/7295866/20250127/nii/7_3D_SAG_T1_MPRAGE_+C_20250127132612_100.nii.gz', fixed='/123/onlylian/0/tmpgp96622o/clipped.nii.gz', +model='joint', out_moving='/123/onlylian/0/tmpgp96622o/joint.nii.gz', out_fixed='/123/onlylian/0/tmpgp96622o/out_fixed-joint.nii.gz', +header_only=False, trans='/123/onlylian/0/tmpgp96622o/moving_to_fixed-joint.nii.gz', inverse='/123/onlylian/0/tmpgp96622o/fixed_to_moving-joint.nii.gz', +init=None, mid_space=False, threads=None, gpu=True, hyper=0.5, steps=7, extent=256, weights=None, verbose=False, out_dir=None) +''' +def register(ct0, ct1, moving, out_root): + FREESURFER_HOME = '/mnt/1218/Public/packages/freesurfer-8.0.0-beta/' + # out_root = Path(ct0).resolve().parent/os.path.basename(mr).replace('.nii.gz','') + + # print(out_root) + modality = os.path.basename(out_root) + # exit() + + out_root = Path(out_root)/os.path.basename(moving).replace('.nii.gz','') + out_root.mkdir(exist_ok=True) + + logger.info(' '.join((ct0, ct1, moving, str(out_root)))) + + orig = sf.load_volume(moving) + if modality == 'CT': + clipped = out_root/'clipped.nii.gz' + + cl = orig.clip(0, 80) + cl.save(clipped) + + MODELS = [ + 'rigid', + # 'affine', + # 'joint', + ] + + else: + clipped = moving + MODELS = [ + 'rigid', + 'affine', + 'joint', + ] + + + + # exit() + + + + default = { + + 'command': 'register', + 'header_only': False, + 'init': None, + 'mid_space': False, + 'threads': None, + + # 'gpu': False, + 'gpu': True, + + 'verbose': False, + # 'verbose': True, + + 'hyper': 0.5, + 'steps': 7, + 'extent': 256, + + 'weights': None, + + # 'model': 'affine', + + # 'out_dir': None, + # 'out_fixed': 'out_fixed.nii.gz', + # 'out_moving': 'out_moving.nii.gz', + # 'trans': None, + # 'inverse': None, + + 'out_fixed': None, + 'out_moving': None, + 'trans': None, + 'inverse': None, + + 'moving' : clipped, + 'fixed' : ct1, + +# 'weights': str(Path(__file__).resolve().parent/'mri_synthmorph/models/synthmorph.affine.2.h5'), + + } + + os.environ["FREESURFER_HOME"] = FREESURFER_HOME + os.environ["XLA_FLAGS"] = '--xla_gpu_cuda_data_dir=%s'% os.environ["CONDA_PREFIX"] + + for m in MODELS: + default['model'] = m + default['out_dir'] = out_root/m + + # if m in ('affine', 'rigid'): + # default['trans'] = 'trans.lta' + # default['inverse'] = 'inverse.lta' + # else: + # default['trans'] = 'trans.nii.gz' + # default['inverse'] = 'inverse.nii.gz' + + arg=argparse.Namespace(**default) + +# CONDA_PREFIX=/home/xfr/.conda/envs/25reg +# XLA_FLAGS=--xla_gpu_cuda_data_dir=/path/to/cuda + + logger.info('registering %s'%m) + registration.register(arg) + logger.info('registered %s'%m) + + if m in ( + 'rigid', + 'affine', + 'joint', + ): + # which = 'affine' if arg.trans.endswith('.lta') else 'warp' + + out = out_root/('%s.nii.gz'%m) + if m in ['affine', 'rigid']: + trans = sf.load_affine(default['out_dir']/'tra_1.lta') + prop = dict(method='linear', resample=True, fill=0) + orig.transform(trans, **prop).save(out) + logger.info('transformed %s'%out) + else: + # need to resample before transform in warp, too complicated, just copy it + # trans1 = default['out_dir']/'tra_1.nii.gz' + # trans = sf.load_warp(trans1) + + shutil.copy(default['out_dir']/'out_1.nii.gz', out) + logger.info('copied %s'% out) + + with open(out_root/'metric.txt', 'w') as f_metrics: + for m in MODELS: + out1 = sf.load_volume(out_root/m/'out_1.nii.gz').data + inp2 = sf.load_volume(out_root/m/'inp_2.nii.gz').data + met = normalized_mutual_information(out1, inp2) + f_metrics.write('%s\t%f\n'%(m, met)) + + return out_root + +def check(epath): + registered = 0 + for root, dirs, files in os.walk(epath): + dirs.sort() + + RT_DIR = os.path.join(root, 'RT') + + ORGAN_DIR = os.path.join(RT_DIR, 'ORGAN') + if not os.path.isdir(ORGAN_DIR): + continue + + # if there is no eye, it's no a brain image + eye = None + organs = sorted(os.scandir(ORGAN_DIR), key=lambda e: e.name) + for o in organs: + if 'eye' in o.name.lower(): + eye = o + if eye is None: + logger.info('no eye... skip ' + root) + # exit() + return None + + ct_image = os.path.join(RT_DIR, 'ct_image.nii.gz') + + + + outdir = os.path.join(OUT_ROOT, os.path.relpath(root, PATIENTS_ROOT)) + logger.info(outdir) + + os.makedirs(outdir, exist_ok=True) + # ct0_nii = os.path.join(outdir, 'ct0.nii.gz') + ct1_nii = os.path.join(outdir, 'clipped.nii.gz') + # shutil.copy(ct_image, ct0_nii) + + ct = sf.load_volume(ct_image) + clipped = ct.clip(0, 80) + clipped.save(ct1_nii) + + for root2, dirs2, files2 in os.walk(root): + dirs2.sort() + outdir = os.path.join(OUT_ROOT, os.path.relpath(root2, PATIENTS_ROOT)) + if root2.endswith('RT'): + modality = 'RT' + logger.info('copying %s %s' %(root2, outdir)) + shutil.copytree(root2, outdir) + # exit() + continue + skip = (root2==root) or ('RT' in root2.split('/')) + if skip: + continue + if root2.endswith('CT'): + modality = 'CT' + else: + modality = 'other' + logger.info(' '.join([str(skip), root2, modality])) + outdir = os.path.join(OUT_ROOT, os.path.relpath(root2, PATIENTS_ROOT)) + os.makedirs(outdir, exist_ok=True) + for e in sorted(os.scandir(root2), key=lambda e: e.name): + if not e.name.endswith('.nii.gz'): + continue + if '_RTDOSE_' in e.name: + continue + if '_DTI_' in e.name: + continue + if '_ROI1.' in e.name: + continue + + OUT_IMG = os.path.join(outdir, e.name) + if os.path.isfile(OUT_IMG): + logger.info('skip '+ OUT_IMG) + continue + + logger.info(' '.join([e.name, e.path])) + + moving = e.path + register(ct_image, ct1_nii, moving, outdir) + registered += 1 + # exit() + + + + # exit() + return registered + + + +def main(): + + # check('/mnt/1220/Public/dataset2/G4/3L6LOEER') # bad registration + # exit() + + EXCLUDE = ( + # 'LLUQJUY4', #cervical + ) + + os.makedirs(OUT_ROOT, exist_ok=True) + + LOCK_DIR = os.path.join(OUT_ROOT, '0lock') + os.makedirs(LOCK_DIR, exist_ok=True) + for e in sorted(os.scandir(PATIENTS_ROOT), key=lambda e: e.name): + if e.is_dir(): + d = shelve.open(SHELVE) + if e.name in d or e.name in EXCLUDE: + logger.info('skip '+ e.name) + d.close() + continue + d.close() + lock_path = os.path.join(LOCK_DIR, '%s.lock'%e.name) + lock = filelock.FileLock(lock_path, timeout=1) + try: + lock.acquire() + except: + logger.info(lock_path + ' locked') + continue + ret = check(e.path) + lock.release() + # exit() + d = shelve.open(SHELVE) + d[e.name] = ret + d.close() + +if __name__ == '__main__': + main() diff --git a/src/mri_synthmorph b/src/mri_synthmorph new file mode 120000 index 0000000..1c96e4c --- /dev/null +++ b/src/mri_synthmorph @@ -0,0 +1 @@ +../mri_synthmorph/ \ No newline at end of file diff --git a/src/requirements.txt b/src/requirements.txt new file mode 100644 index 0000000..94c5307 --- /dev/null +++ b/src/requirements.txt @@ -0,0 +1,6 @@ + +filelock + +git+https://github.com/adalca/neurite.git +git+https://github.com/freesurfer/surfa.git +git+https://github.com/voxelmorph/voxelmorph.git diff --git a/src/synthmorph b/src/synthmorph new file mode 120000 index 0000000..6c67576 --- /dev/null +++ b/src/synthmorph @@ -0,0 +1 @@ +../mri_synthmorph/synthmorph \ No newline at end of file diff --git a/src/synthmorph-ct.py b/src/synthmorph-ct.py new file mode 100644 index 0000000..16be785 --- /dev/null +++ b/src/synthmorph-ct.py @@ -0,0 +1,59 @@ +''' +conda create -n voxelmorph -c conda-forge python simpleitk tensorflow-gpu +conda create -n voxelmorph -c conda-forge "python<3.11" simpleitk "tensorflow-gpu<2.16" +conda activate voxelmorph +pip install git+https://github.com/adalca/neurite.git git+https://github.com/freesurfer/surfa.git git+https://github.com/voxelmorph/voxelmorph.git + +wget https://raw.githubusercontent.com/freesurfer/freesurfer/dev/mri_synthmorph/mri_synthmorph +wget https://surfer.nmr.mgh.harvard.edu/docs/synthmorph/synthmorph.affine.2.h5 + +export FREESURFER_HOME=/home/xfr/git9/Taipei-1/trials +time ./mri_synthmorph -m affine -t trans.lta '/mnt/1218/Public/dataset2/M6/ZYRGTRKJ/20230728/MR/3D_SAG_T1_MPRAGE_+C_MPR_Tra_20230728143005_14.nii.gz' clipped.nii.gz + + +#### https://github.com/freesurfer/freesurfer/tree/dev/mri_synthmorph + +conda create -n synthmorph -c conda-forge -c nvidia python=3.11 simpleitk tensorflow-gpu=2.17 cuda-nvcc +conda create -y -n synthmorph -c conda-forge -c nvidia python simpleitk tensorflow-gpu cuda-nvcc +conda activate synthmorph +pip install git+https://github.com/adalca/neurite.git git+https://github.com/freesurfer/surfa.git git+https://github.com/voxelmorph/voxelmorph.git +XLA_FLAGS=--xla_gpu_cuda_data_dir=/home/xfr/.conda/envs/synthmorph time ./mri_synthmorph -m affine -t trans.lta -o out-aff.nii.gz moving.nii.gz clipped.nii.gz -g +XLA_FLAGS=--xla_gpu_cuda_data_dir=/home/xfr/.conda/envs/synthmorph time ./mri_synthmorph -o out.nii.gz moving.nii.gz clipped.nii.gz -g + +export FREESURFER_HOME=/mnt/1218/Public/packages/freesurfer-7.4.1 +export FREESURFER_HOME=/mnt/1218/Public/packages/freesurfer-8.0.0-beta +source $FREESURFER_HOME/SetUpFreeSurfer.sh +time mri_synthmorph -m affine -t trans.lta -o out-aff.nii.gz moving.nii.gz clipped.nii.gz -g + +time mri_easyreg --ref clipped.nii.gz --flo moving.nii.gz + + + + ./cuda_sdk_lib + /home/conda/feedstock_root/build_artifacts/tensorflow-split_1729095706337/_build_env/targets/x86_64-linux + /usr/local/cuda + /home/xfr/.conda/envs/synthmorph/lib/python3.11/site-packages/tensorflow/python/platform/../../../nvidia/cuda_nvcc + /home/xfr/.conda/envs/synthmorph/lib/python3.11/site-packages/tensorflow/python/platform/../../../../nvidia/cuda_nvcc + + +Less than 10 seconds!!! +''' + +fi = '/nn/7295866/20250127/nii/a_1.1_CyberKnife_head(MAR)_20250127111447_5.nii.gz' +mv = '/nn/7295866/20250127/nii/7_3D_SAG_T1_MPRAGE_+C_20250127132612_100.nii.gz' + +fi = '/mnt/1218/Public/dataset2/M6/ZYRGTRKJ/20230728/CT/1.1_CyberKnife_head(MAR)_20230728111920_3.nii.gz' +mv = '/mnt/1218/Public/dataset2/M6/ZYRGTRKJ/20230728/MR/3D_SAG_T1_MPRAGE_+C_MPR_Tra_20230728143005_14.nii.gz' + +import shutil + +import surfa as sf + + +ct = sf.load_volume(fi) +clipped = ct.clip(0, 80) +clipped.save('../clipped.nii.gz') + +shutil.copy(mv, '../moving.nii.gz') + +