mridefacer

#!/bin/bash
# emacs: -*- mode: shell-script; indent-tabs-mode: nil -*-
# vi: set ft=sh sts=2 ts=2 sw=2 et:

# play save
set -e
PS4='+ $LINENO: '

# version of this script
mridefacer_version=0.2.2

############
# Defaults #
############

# where is the de-face mask and template
md_datadir=${MRIDEFACER_DATA_DIR:-}
# produce fancy colored output
md_color_output=1
# enables additional status message if set to true
md_verbose=0
# output directory, use same as input when empty
md_outdir=''
# TODO: clarify those params... confusing!
# flag whether to output directory, use same as input when empty
md_apply2orig=0
# flag whether to keep original file if --apply
md_keep_orig=0
# flag whether to output directory, and place aside the original
md_apply=1
# flag either to do 'annex' treatment of the original and produced files
md_annex=0
# init next xfm with the result of the preceding one.
# not advised to guarantee reproducibility across independent runs
# (i.e. per file vs per group of files)
md_use_prev_xfm=0


print_version()
{
cat << EOT
mridefacer $mridefacer_version

Copyright (C) 2013-2016 Michael Hanke <michael.hanke@gmail.com>
              2017      Yaroslav Halchenko <debian@onerussian.com>

Release under the MIT License.

Written by Michael Hanke

EOT
}


print_short_description()
{
cat << EOT
Helper to aid de-identification of MRI images (3D or 4D).

A de-face mask image will be created for each input image, by aligning
template mask to the input. Optionally, de-face masks can be applied
to the input images to either replace the original image, or create
modified copies in a specified output directory.

This tool can process single or series of images. In the latter case,
the computed transformation between template image and input image will
be updated incrementally for the next image in the series. This feature
is most suitable for processing images that have been recorded in
temporal succession.

EOT
}


print_help()
{
cat << EOT

Usage:  mridefacer [OPTIONS] <imagefile> [<imagefile> [...]]


Options:

--apply
  If this flag is set, the defacing is applied to the input images.
  By default, defacing mask images are created only. If --outdir is
  not provided, the actual input images will be replaced by defaced
  versions. In this case, the resulting images will be in FSL's
  standard radiological orientation.

--apply-only
  If this flag is set, the defacing is applied to the input images,
  and stored into a new file (with _defaced suffix) instead of replacing
  original one.

--keep-orig
  If this flag is set (in effect only if --apply), original image will be
  stored with _orig suffix.

--annex
  If this flag is set, it will add all produced niftis to git annex, and
  annotate original anatomical with metadata
    distribution-restrictions=sensitive (datalad convention)
  to make it easy to automate non-redistribution of these materials.

--use-prev-xfm
  If multiple files are provided, alignmet for the next file will
  be initialized with the result for the previous file.  Not advised
  if you want to achieve reproducible results independent of the order
  of the files

-h, --help
  Print short description, usage summary and option list.

--no-colored
  If set, mridefacer won't colorize its status and error messages.

--outdir
  If set, output files will be put into this directory. By default,
  output files are placed in the same directory as the respective input
  files.

--verbose
  Enable additional status messages.

--verbose-help
  Print all available help.

--version
  Print version information and exit.

EOT
}


print_additional_description()
{
cat << EOT

Examples:

mridefacer ADMME
  Deface a single image

Environment:

mridefacer requires FSL and the environment variable FSLDIR needs to be set
accordingly. mridefacer acknowledges MRIDEFACER_DATA_DIR to determine the
location of the required MRI template images.


Report bugs to <michael.hanke@gmail.com>.

EOT
}


################################
# Commandline options handling #
################################

# Parse commandline options (taken from the getopt examples from the Debian util-linux package)
# Note that we use `"$@"' to let each command-line parameter expand to a
# separate word. The quotes around `$@' are essential!
# We need CLOPTS as the `eval set --' would nuke the return value of getopt.
CLOPTS=`getopt -o h --long help,verbose-help,version,no-color,verbose,outdir:,apply,apply-only,keep-orig,annex,use-prev-xfm -n 'mridefacer' -- "$@"`

if [ $? != 0 ] ; then
  echo "Terminating..." >&2
  exit 1
fi

# Note the quotes around `$CLOPTS': they are essential!
eval set -- "$CLOPTS"

while true ; do
  case "$1" in
	  --no-color) md_color_output=0; shift;;
	  --verbose) md_verbose=1; shift;;
	  --annex) md_annex=1; shift;;
	  --apply) md_apply2orig=1; shift;;
	  --apply-only) md_apply=1; shift;;
	  --keep-orig) md_keep_orig=1; shift;;
	  --use-prev-xfm) md_use_prev_xfm=1; shift;;
	  -h|--help) print_short_description; print_help; exit 0;;
	  --verbose-help) print_short_description; print_help; print_additional_description; exit 0;;
	  --version) print_version; exit 0;;
	  --outdir) md_outdir="$2"; shift; shift;;
	  --) shift ; break ;;
	  *) echo "Internal error! ($1)"; exit 1;;
  esac
done

# colorful output requested?
if [ "$md_color_output" = 1 ]; then
  black='\e[0;30m'; Black='\e[1;30m'
  red='\e[0;31m'; Red='\e[1;31m'
  green='\e[0;32m'; Green='\e[1;32m'
  yellow='\e[0;33m'; Yellow='\e[1;33m'
  blue='\e[0;34m'; Blue='\e[1;34m'
  cyan='\e[0;36m'; Cyan='\e[1;36m'
  white='\e[0;37m'; White='\e[1;37m'
  NC='\e[0m' #no color
else
  black=''; Black=''
  red=''; Red=''
  green=''; Green=''
  yellow=''; Yellow=''
  blue=''; Blue=''
  cyan=''; Cyan=''
  white=''; White=''
  NC='' #no color
fi

_fsl_env="/etc/fsl/fsl.sh"

if [ -z "${FSLDIR:-}" ]; then
  if [ -e "$_fsl_env" ]; then
    printf "${Green}FSLDIR environment variable not set. Sourcing present $_fsl_env${NC}\n"
    source "$_fsl_env"
  else
    printf "${Red}FSLDIR environment variable not set${NC}\n" && exit 1
  fi
fi

# play safer
set -u

# Try both in-source and "as installed" settings
if [ -z "$md_datadir" ]; then
  exe_dir="$(dirname $0)"
  tried=
  for cand in data ../share/mridefacer; do
    cand_dir="$exe_dir/$cand"
	tried+=" '$cand_dir'"
    if [ -e "$cand_dir/face_teeth_ear_mask.nii.gz" ] && [ -e "$cand_dir/head_tmpl_brain.nii.gz" ]; then
      md_datadir="$cand_dir"
      break
	fi
  done
  if [ -z "$md_datadir" ]; then
    printf "${Red}Could not determined directory with mridefacer data. Tried: $tried${NC}\n" && exit 1
  fi
fi

if [ $# -lt 1 ]; then
  printf "${Red}Error: This command needs at least one positional argument.\n${NC}"
  exit 1
fi

#
# End of boilerplate
#

getdeface () {
  # image to deface
  target=$1
  # seed the alignment with an initial transformation
  # make empty string to disable
  init_xfm=$2
  # outfile basename
  out_base=$3
  # flag whether to realign image (1) or use initial transform as such (0)
  update_xfm=$4
  # search radius in degrees
  sr="${5:-90}"
  # bet frac
  bf="${6:-0.5}"

  $FSLDIR/bin/fslreorient2std "${target}" ${out_base}_instd
  # do again to get the matrix
  $FSLDIR/bin/fslreorient2std "${target}" > ${out_base}_instd.mat
  $FSLDIR/bin/convert_xfm -omat ${out_base}_fromstd.mat -inverse ${out_base}_instd.mat

  if [[ $($FSLDIR/bin/fslinfo ${out_base}_instd | grep '^dim4' | sed -e 's/.* //g') -gt 1 ]]; then
    #echo "Use mean volume as reference for de-facing"
    #$FSLDIR/bin/fslmaths ${out_base}_instd -Tmean ${out_base}_meanvol
    if [ "$md_verbose" = 1 ]; then
      printf "${Green}Use first volume as reference for de-facing\n${NC}"
    fi
    $FSLDIR/bin/fslroi ${out_base}_instd ${out_base}_meanvol 0 1
  else
    $FSLDIR/bin/imln ${out_base}_instd ${out_base}_meanvol
  fi

  # brain extract
  $FSLDIR/bin/bet ${out_base}_meanvol ${out_base}_brain -R -f "${bf}"
  # protect again failing BET with unreliable exit code
  [ "`$FSLDIR/bin/imtest "${out_base}_brain"`" = "1" ]

  # subsample highres stuff -- doesn't gain much beyond 1mm resolution
  if [[ $(echo "$(fslinfo ${out_base}_meanvol | grep '^pixdim' | head -n3 | sed -e 's/.* //g' | numbound) > 2.1" | bc) == 1 ]]; then
    imln ${out_base}_brain ${out_base}_subsamp
  else
    if [ "$md_verbose" = 1 ]; then
      printf "${Green}Sub-sample high-resolution image for alignment\n${NC}"
    fi
    $FSLDIR/bin/fslmaths ${out_base}_brain -subsamp2 ${out_base}_subsamp
  fi

  # XXX SOME LIKE IT, SOME NOT
  #opts="-usesqform -bins 256 -cost corratio -searchrx -$sr $sr -searchry -$sr $sr -searchrz -$sr $sr -dof 12"
  opts="-bins 256 -cost corratio -searchrx -$sr $sr -searchry -$sr $sr -searchrz -$sr $sr -dof 12"

  if [ ! -z "$init_xfm" ]; then
    if [ "$md_verbose" = 1 ]; then
      printf "${green}Use provided transformation matrix for initialization\n${NC}"
    fi
    opts="-init $init_xfm $opts"
  fi

  # align template to input image
  if [ $update_xfm = 1 ]; then
    if [ "$md_verbose" = 1 ]; then
      printf "${Green}Align input for defacing\n${NC}"
    fi
    $FSLDIR/bin/flirt -in $md_datadir/head_tmpl_brain \
      -inweight $md_datadir/head_tmpl_weights -ref ${out_base}_subsamp \
      -omat ${out_base}.mat -out ${out_base}_alignedtmpl $opts
    popts="-init ${out_base}.mat"
  else
    if [ ! -z "$init_xfm" ]; then
       ln -s ${init_xfm} ${out_base}.mat
       popts="-init ${out_base}.mat"
    else
       popts=""
    fi
  fi

  # project de-face mask onto std reference
  $FSLDIR/bin/flirt -in $md_datadir/face_teeth_ear_mask -applyxfm \
     $popts -out ${out_base}_defacemask_aligned_std -interp trilinear \
      -ref ${out_base}_meanvol

  # project de-face mask onto reference
  $FSLDIR/bin/flirt -in ${out_base}_defacemask_aligned_std -applyxfm \
     -init ${out_base}_fromstd.mat -out ${out_base}_defacemask_aligned \
     -interp trilinear -ref "${target}"

  # threshold de-face mask and store as output
  $FSLDIR/bin/fslmaths ${out_base}_defacemask_aligned -thr 0.5 -bin \
     ${out_base}_defacemask -odt char
}


#
# main
#

# prep tmp dir
wdir=`mktemp -d --tmpdir mridefacer.XXXXXXXXXXX`
if [ "$md_verbose" = 1 ]; then
  printf "${green}Using temporary directory '$wdir'\n${NC}"
fi
# take care of clean-up on exit
trap "rm -rf $wdir" EXIT


count=0
last_xfm=''
for infile in "$@"; do
  [ "$md_use_prev_xfm" = 1 ] || last_xfm=''
  infile_dir=$(dirname "$infile")
  infile_name=$(basename "$infile")
  under_annex=0
  if [ "$md_annex" = 1 ]; then
    ( cd "$infile_dir" && git annex info --fast; )  >/dev/null 2>&1 && under_annex=1 || {
      printf "${red}File $infile is not under annex repository, no annex treatments${NC}\n"
    }
  fi

  [ $under_annex = 0 ] || (
    cd "$infile_dir"
    if git annex info --fast "$infile_name" >/dev/null 2>&1; then
      existing_restrictions=`git annex metadata -g distribution-restrictions "$infile_name"`
      if [ -z "$existing_restrictions" ]; then
        git annex metadata -s distribution-restrictions=sensitive "$infile_name" >/dev/null
      else
        printf "${yellow}File $infile already has restrictions set to '$existing_restrictions'. Not adding new${NC}\n"
      fi
    else
      printf "\n{red}File $infile is under annex but not annexed. Not enhancing its metadata${NC}\n"
    fi
  )

  getdeface "$infile" "$last_xfm" "$wdir/out_$count" 1 90 0.5
  last_xfm="`ls -1 "$wdir/out_${count}.mat" 2>/dev/null`" || true
  outbase="`$FSLDIR/bin/remove_ext "$infile"`"
  if [ ! -z "$md_outdir" ]; then
      outbase="$md_outdir/`basename "$outbase"`"
  fi
  # Since we are producing new defacemask, no sidecar file, even if there is one
  # should survive anyways, so first we remove anything in the target path.
  # Needed in case of git-annex'ed files and avoids unlocking and hopefully
  # prevents content mutation which we have managed to cause once
  rm -f "$outbase"_defacemask.*
  deface_fname="$outbase"_defacemask
  deface_dir=$(dirname "$deface_fname")

  $FSLDIR/bin/imcp "$wdir/out_${count}_defacemask" "$deface_fname"
  # make explicit test for success, as imcp doesn't
  [ "`$FSLDIR/bin/imtest "$outbase"_defacemask`" = "1" ]

  function add_to_annex() {
    for fname in "$@"; do
      d=$(dirname "$fname")
      f=$(basename "$fname")
      added=
      if [ ! $under_annex = 0 ]; then
       (
       cd "$d";
       if [ -e "$f" ]; then
          git annex add "$f";
          added+=1
       fi
       if /bin/ls "$f".* 2>/dev/null | grep -q . ; then
          git annex add "$f".*
          added+=1
       fi
       if [ -z "$added" ]; then
          echo "No files were found for $fname" >&2
          exit 1
       fi
       );
      fi
    done
  }

  if [ "$md_apply2orig" = 1 ] || [ "$md_apply" = 1 ]; then
    $FSLDIR/bin/fslmaths "$infile" -mul "$wdir/out_${count}_defacemask" "$wdir/out_${count}_defaced" -odt input
    if [ "$md_apply2orig" = 1 ]; then
      outsuffix=''
    else
      outsuffix='_defaced'
    fi

    if [ $md_apply = 1 ] && [ $md_keep_orig = 1 ]; then
      orig_fname=$(echo $infile | sed -e 's,\.nii.*,_orig,g')
      $FSLDIR/bin/immv "$infile" "${orig_fname}"
      add_to_annex "${orig_fname}"
    else
      add_to_annex "${infile}"
    fi

    $FSLDIR/bin/immv "$wdir/out_${count}_defaced" "${outbase}${outsuffix}";
    add_to_annex "${outbase}${outsuffix}"
  fi
  add_to_annex "${deface_fname}"
  count=$((count + 1))
done