code: # emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
  # vi: set ft=python sts=4 ts=4 sw=4 et:
  #
  # Copyright 2021 The NiPreps Developers <nipreps@gmail.com>
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  #
  # We support and encourage derived works from this project, please read
  # about our expectations at
  #
  #     https://www.nipreps.org/community/licensing/
  #
  """Nipype interfaces to support anatomical workflow."""
  import os.path as op
  from builtins import zip
  
  import nibabel as nb
  import numpy as np
  import scipy.ndimage as nd
  from mriqc.qc.anatomical import (
      art_qi1,
      art_qi2,
      cjv,
      cnr,
      efc,
      fber,
      rpve,
      snr,
      snr_dietrich,
      summary_stats,
      volume_fraction,
      wm2max,
  )
  from mriqc.utils.misc import _flatten_dict
  from nipype.interfaces.base import (
      BaseInterfaceInputSpec,
      File,
      InputMultiPath,
      SimpleInterface,
      TraitedSpec,
      isdefined,
      traits,
  )
  from nipype.utils.filemanip import fname_presuffix
  
  
  class StructuralQCInputSpec(BaseInterfaceInputSpec):
      in_file = File(exists=True, mandatory=True, desc="file to be plotted")
      in_noinu = File(exists=True, mandatory=True, desc="image after INU correction")
      in_segm = File(exists=True, mandatory=True, desc="segmentation file from FSL FAST")
      in_bias = File(exists=True, mandatory=True, desc="bias file")
      head_msk = File(exists=True, mandatory=True, desc="head mask")
      air_msk = File(exists=True, mandatory=True, desc="air mask")
      rot_msk = File(exists=True, mandatory=True, desc="rotation mask")
      artifact_msk = File(exists=True, mandatory=True, desc="air mask")
      in_pvms = InputMultiPath(
          File(exists=True),
          mandatory=True,
          desc="partial volume maps from FSL FAST",
      )
      in_tpms = InputMultiPath(File(), desc="tissue probability maps from FSL FAST")
      mni_tpms = InputMultiPath(File(), desc="tissue probability maps from FSL FAST")
      in_fwhm = traits.List(
          traits.Float, mandatory=True, desc="smoothness estimated with AFNI"
      )
      human = traits.Bool(True, usedefault=True, desc="human workflow")
  
  
  class StructuralQCOutputSpec(TraitedSpec):
      summary = traits.Dict(desc="summary statistics per tissue")
      icvs = traits.Dict(desc="intracranial volume (ICV) fractions")
      rpve = traits.Dict(desc="partial volume fractions")
      size = traits.Dict(desc="image sizes")
      spacing = traits.Dict(desc="image sizes")
      fwhm = traits.Dict(desc="full width half-maximum measure")
      inu = traits.Dict(desc="summary statistics of the bias field")
      snr = traits.Dict
      snrd = traits.Dict
      cnr = traits.Float
      fber = traits.Float
      efc = traits.Float
      qi_1 = traits.Float
      wm2max = traits.Float
      cjv = traits.Float
      out_qc = traits.Dict(desc="output flattened dictionary with all measures")
      out_noisefit = File(exists=True, desc="plot of background noise and chi fitting")
      tpm_overlap = traits.Dict
  
  
  class StructuralQC(SimpleInterface):
      """
      Computes anatomical :abbr:`QC (Quality Control)` measures on the
      structural image given as input
  
      """
  
      input_spec = StructuralQCInputSpec
      output_spec = StructuralQCOutputSpec
  
      def _run_interface(self, runtime):  # pylint: disable=R0914,E1101
          imnii = nb.load(self.inputs.in_noinu)
          erode = (
              np.all(np.array(imnii.header.get_zooms()[:3], dtype=np.float32) < 1.9)
              if self.inputs.human
              else False
          )
  
          # Load image corrected for INU
          inudata = np.nan_to_num(imnii.get_fdata())
          inudata[inudata < 0] = 0
  
          if np.all(inudata < 1e-5):
              raise RuntimeError(
                  "Input inhomogeneity-corrected data seem empty. "
                  "MRIQC failed to process this dataset."
              )
  
          # Load binary segmentation from FSL FAST
          segnii = nb.load(self.inputs.in_segm)
          segdata = np.asanyarray(segnii.dataobj).astype(np.uint8)
  
          if np.sum(segdata > 0) < 1e3:
              raise RuntimeError(
                  "Input segmentation data is likely corrupt. "
                  "MRIQC failed to process this dataset."
              )
  
          # Load air, artifacts and head masks
          airdata = np.asanyarray(nb.load(self.inputs.air_msk).dataobj).astype(np.uint8)
          artdata = np.asanyarray(nb.load(self.inputs.artifact_msk).dataobj).astype(
              np.uint8
          )
  
          headdata = np.asanyarray(nb.load(self.inputs.head_msk).dataobj).astype(np.uint8)
          if np.sum(headdata > 0) < 100:
              raise RuntimeError(
                  "Detected less than 100 voxels belonging to the head mask. "
                  "MRIQC failed to process this dataset."
              )
  
          rotdata = np.asanyarray(nb.load(self.inputs.rot_msk).dataobj).astype(np.uint8)
  
          # Load Partial Volume Maps (pvms) from FSL FAST
          pvmdata = []
          for fname in self.inputs.in_pvms:
              pvmdata.append(nb.load(fname).get_fdata(dtype="float32"))
              if np.sum(pvmdata[-1] > 1e-4) < 10:
                  raise RuntimeError(
                      "Detected less than 10 voxels belonging to one tissue prob. map. "
                      "MRIQC failed to process this dataset."
                  )
  
          # Summary stats
          stats = summary_stats(inudata, pvmdata, airdata, erode=erode)
          self._results["summary"] = stats
  
          # SNR
          snrvals = []
          self._results["snr"] = {}
          for tlabel in ["csf", "wm", "gm"]:
              snrvals.append(
                  snr(
                      stats[tlabel]["median"],
                      stats[tlabel]["stdv"],
                      stats[tlabel]["n"],
                  )
              )
              self._results["snr"][tlabel] = snrvals[-1]
          self._results["snr"]["total"] = float(np.mean(snrvals))
  
          snrvals = []
          self._results["snrd"] = {
              tlabel: snr_dietrich(
                  stats[tlabel]["median"],
                  mad_air=stats["bg"]["mad"],
                  sigma_air=stats["bg"]["stdv"],
              )
              for tlabel in ["csf", "wm", "gm"]
          }
          self._results["snrd"]["total"] = float(
              np.mean([val for _, val in list(self._results["snrd"].items())])
          )
  
          # CNR
          self._results["cnr"] = cnr(
              stats["wm"]["median"],
              stats["gm"]["median"],
              stats["bg"]["stdv"],
              stats["wm"]["stdv"],
              stats["gm"]["stdv"],
          )
  
          # FBER
          self._results["fber"] = fber(inudata, headdata, rotdata)
  
          # EFC
          self._results["efc"] = efc(inudata, rotdata)
  
          # M2WM
          self._results["wm2max"] = wm2max(inudata, stats["wm"]["median"])
  
          # Artifacts
          self._results["qi_1"] = art_qi1(airdata, artdata)
  
          # CJV
          self._results["cjv"] = cjv(
              # mu_wm, mu_gm, sigma_wm, sigma_gm
              stats["wm"]["median"],
              stats["gm"]["median"],
              stats["wm"]["mad"],
              stats["gm"]["mad"],
          )
  
          # FWHM
          fwhm = np.array(self.inputs.in_fwhm[:3]) / np.array(
              imnii.header.get_zooms()[:3]
          )
          self._results["fwhm"] = {
              "x": float(fwhm[0]),
              "y": float(fwhm[1]),
              "z": float(fwhm[2]),
              "avg": float(np.average(fwhm)),
          }
  
          # ICVs
          self._results["icvs"] = volume_fraction(pvmdata)
  
          # RPVE
          self._results["rpve"] = rpve(pvmdata, segdata)
  
          # Image specs
          self._results["size"] = {
              "x": int(inudata.shape[0]),
              "y": int(inudata.shape[1]),
              "z": int(inudata.shape[2]),
          }
          self._results["spacing"] = {
              i: float(v) for i, v in zip(["x", "y", "z"], imnii.header.get_zooms()[:3])
          }
  
          try:
              self._results["size"]["t"] = int(inudata.shape[3])
          except IndexError:
              pass
  
          try:
              self._results["spacing"]["tr"] = float(imnii.header.get_zooms()[3])
          except IndexError:
              pass
  
          # Bias
          bias = nb.load(self.inputs.in_bias).get_fdata()[segdata > 0]
          self._results["inu"] = {
              "range": float(
                  np.abs(np.percentile(bias, 95.0) - np.percentile(bias, 5.0))
              ),
              "med": float(np.median(bias)),
          }  # pylint: disable=E1101
  
          mni_tpms = [nb.load(tpm).get_fdata() for tpm in self.inputs.mni_tpms]
          in_tpms = [nb.load(tpm).get_fdata() for tpm in self.inputs.in_pvms]
          overlap = fuzzy_jaccard(in_tpms, mni_tpms)
          self._results["tpm_overlap"] = {
              "csf": overlap[0],
              "gm": overlap[1],
              "wm": overlap[2],
          }
  
          # Flatten the dictionary
          self._results["out_qc"] = _flatten_dict(self._results)
          return runtime
  
  
  class ArtifactMaskInputSpec(BaseInterfaceInputSpec):
      in_file = File(exists=True, mandatory=True, desc="File to be plotted")
      head_mask = File(exists=True, mandatory=True, desc="head mask")
      rot_mask = File(exists=True, desc="a rotation mask")
      nasion_post_mask = File(
          exists=True,
          mandatory=True,
          desc="nasion to posterior of cerebellum mask",
      )
  
  
  class ArtifactMaskOutputSpec(TraitedSpec):
      out_hat_msk = File(exists=True, desc='output "hat" mask')
      out_art_msk = File(exists=True, desc="output artifacts mask")
      out_air_msk = File(exists=True, desc='output "hat" mask, without artifacts')
  
  
  class ArtifactMask(SimpleInterface):
      """
      Computes the artifact mask using the method described in [Mortamet2009]_.
      """
  
      input_spec = ArtifactMaskInputSpec
      output_spec = ArtifactMaskOutputSpec
  
      def _run_interface(self, runtime):
          imnii = nb.load(self.inputs.in_file)
          imdata = np.nan_to_num(imnii.get_fdata().astype(np.float32))
  
          # Remove negative values
          imdata[imdata < 0] = 0
  
          hmdata = np.asanyarray(nb.load(self.inputs.head_mask).dataobj)
          npdata = np.asanyarray(nb.load(self.inputs.nasion_post_mask).dataobj)
  
          # Invert head mask
          airdata = np.ones_like(hmdata, dtype=np.uint8)
          airdata[hmdata == 1] = 0
  
          # Calculate distance to border
          dist = nd.morphology.distance_transform_edt(airdata)
  
          # Apply nasion-to-posterior mask
          airdata[npdata == 1] = 0
          dist[npdata == 1] = 0
          dist /= dist.max()
  
          # Apply rotation mask (if supplied)
          if isdefined(self.inputs.rot_mask):
              rotmskdata = np.asanyarray(nb.load(self.inputs.rot_mask).dataobj)
              airdata[rotmskdata == 1] = 0
  
          # Run the artifact detection
          qi1_img = artifact_mask(imdata, airdata, dist)
  
          fname, ext = op.splitext(op.basename(self.inputs.in_file))
          if ext == ".gz":
              fname, ext2 = op.splitext(fname)
              ext = ext2 + ext
  
          self._results["out_hat_msk"] = op.abspath("{}_hat{}".format(fname, ext))
          self._results["out_art_msk"] = op.abspath("{}_art{}".format(fname, ext))
          self._results["out_air_msk"] = op.abspath("{}_air{}".format(fname, ext))
  
          hdr = imnii.header.copy()
          hdr.set_data_dtype(np.uint8)
          nb.Nifti1Image(qi1_img, imnii.affine, hdr).to_filename(
              self._results["out_art_msk"]
          )
  
          nb.Nifti1Image(airdata, imnii.affine, hdr).to_filename(
              self._results["out_hat_msk"]
          )
  
          airdata[qi1_img > 0] = 0
          nb.Nifti1Image(airdata, imnii.affine, hdr).to_filename(
              self._results["out_air_msk"]
          )
          return runtime
  
  
  class ComputeQI2InputSpec(BaseInterfaceInputSpec):
      in_file = File(exists=True, mandatory=True, desc="File to be plotted")
      air_msk = File(exists=True, mandatory=True, desc="air (without artifacts) mask")
  
  
  class ComputeQI2OutputSpec(TraitedSpec):
      qi2 = traits.Float(desc="computed QI2 value")
      out_file = File(desc="output plot: noise fit")
  
  
  class ComputeQI2(SimpleInterface):
      """
      Computes the artifact mask using the method described in [Mortamet2009]_.
      """
  
      input_spec = ComputeQI2InputSpec
      output_spec = ComputeQI2OutputSpec
  
      def _run_interface(self, runtime):
          imdata = nb.load(self.inputs.in_file).get_fdata()
          airdata = nb.load(self.inputs.air_msk).get_fdata()
          qi2, out_file = art_qi2(imdata, airdata)
          self._results["qi2"] = qi2
          self._results["out_file"] = out_file
          return runtime
  
  
  class HarmonizeInputSpec(BaseInterfaceInputSpec):
      in_file = File(
          exists=True, mandatory=True, desc="input data (after bias correction)"
      )
      wm_mask = File(exists=True, mandatory=True, desc="white-matter mask")
      erodemsk = traits.Bool(True, usedefault=True, desc="erode mask")
      thresh = traits.Float(0.9, usedefault=True, desc="WM probability threshold")
  
  
  class HarmonizeOutputSpec(TraitedSpec):
      out_file = File(exists=True, desc="input data (after intensity harmonization)")
  
  
  class Harmonize(SimpleInterface):
      """
      Computes the artifact mask using the method described in [Mortamet2009]_.
      """
  
      input_spec = HarmonizeInputSpec
      output_spec = HarmonizeOutputSpec
  
      def _run_interface(self, runtime):
  
          in_file = nb.load(self.inputs.in_file)
          wm_mask = nb.load(self.inputs.wm_mask).get_fdata()
          wm_mask[wm_mask < 0.9] = 0
          wm_mask[wm_mask > 0] = 1
          wm_mask = wm_mask.astype(np.uint8)
  
          if self.inputs.erodemsk:
              # Create a structural element to be used in an opening operation.
              struc = nd.generate_binary_structure(3, 2)
              # Perform an opening operation on the background data.
              wm_mask = nd.binary_erosion(wm_mask, structure=struc).astype(np.uint8)
  
          data = in_file.get_fdata()
          data *= 1000.0 / np.median(data[wm_mask > 0])
  
          out_file = fname_presuffix(
              self.inputs.in_file, suffix="_harmonized", newpath="."
          )
          in_file.__class__(data, in_file.affine, in_file.header).to_filename(out_file)
  
          self._results["out_file"] = out_file
  
          return runtime
  
  
  class RotationMaskInputSpec(BaseInterfaceInputSpec):
      in_file = File(exists=True, mandatory=True, desc="input data")
  
  
  class RotationMaskOutputSpec(TraitedSpec):
      out_file = File(exists=True, desc="rotation mask (if any)")
  
  
  class RotationMask(SimpleInterface):
      """
      Computes the artifact mask using the method described in [Mortamet2009]_.
      """
  
      input_spec = RotationMaskInputSpec
      output_spec = RotationMaskOutputSpec
  
      def _run_interface(self, runtime):
          in_file = nb.load(self.inputs.in_file)
          data = in_file.get_fdata()
          mask = data <= 0
  
          # Pad one pixel to control behavior on borders of binary_opening
          mask = np.pad(mask, pad_width=(1,), mode="constant", constant_values=1)
  
          # Remove noise
          struc = nd.generate_binary_structure(3, 2)
          mask = nd.binary_opening(mask, structure=struc).astype(np.uint8)
  
          # Remove small objects
          label_im, nb_labels = nd.label(mask)
          if nb_labels > 2:
              sizes = nd.sum(mask, label_im, list(range(nb_labels + 1)))
              ordered = list(reversed(sorted(zip(sizes, list(range(nb_labels + 1))))))
              for _, label in ordered[2:]:
                  mask[label_im == label] = 0
  
          # Un-pad
          mask = mask[1:-1, 1:-1, 1:-1]
  
          # If mask is small, clean-up
          if mask.sum() < 500:
              mask = np.zeros_like(mask, dtype=np.uint8)
  
          out_img = in_file.__class__(mask, in_file.affine, in_file.header)
          out_img.header.set_data_dtype(np.uint8)
  
          out_file = fname_presuffix(self.inputs.in_file, suffix="_rotmask", newpath=".")
          out_img.to_filename(out_file)
          self._results["out_file"] = out_file
          return runtime
  
  
  def artifact_mask(imdata, airdata, distance, zscore=10.0):
      """Computes a mask of artifacts found in the air region"""
      from statsmodels.robust.scale import mad
  
      if not np.issubdtype(airdata.dtype, np.integer):
          airdata[airdata < 0.95] = 0
          airdata[airdata > 0.0] = 1
  
      bg_img = imdata * airdata
      if np.sum((bg_img > 0).astype(np.uint8)) < 100:
          return np.zeros_like(airdata)
  
      # Find the background threshold (the most frequently occurring value
      # excluding 0)
      bg_location = np.median(bg_img[bg_img > 0])
      bg_spread = mad(bg_img[bg_img > 0])
      bg_img[bg_img > 0] -= bg_location
      bg_img[bg_img > 0] /= bg_spread
  
      # Apply this threshold to the background voxels to identify voxels
      # contributing artifacts.
      qi1_img = np.zeros_like(bg_img)
      qi1_img[bg_img > zscore] = 1
      qi1_img[distance < 0.10] = 0
  
      # Create a structural element to be used in an opening operation.
      struc = nd.generate_binary_structure(3, 1)
      qi1_img = nd.binary_opening(qi1_img, struc).astype(np.uint8)
      qi1_img[airdata <= 0] = 0
  
      return qi1_img
  
  
  def fuzzy_jaccard(in_tpms, in_mni_tpms):
      overlaps = []
      for tpm, mni_tpm in zip(in_tpms, in_mni_tpms):
          tpm = tpm.reshape(-1)
          mni_tpm = mni_tpm.reshape(-1)
  
          num = np.min([tpm, mni_tpm], axis=0).sum()
          den = np.max([tpm, mni_tpm], axis=0).sum()
          overlaps.append(float(num / den))
      return overlaps
  
apis: numpy.sum numpy.issubdtype numpy.ones_like scipy.ndimage.morphology.distance_transform_edt numpy.average numpy.mean scipy.ndimage.generate_binary_structure scipy.ndimage.label scipy.ndimage.binary_opening numpy.median numpy.asanyarray scipy.ndimage.binary_erosion numpy.all numpy.max numpy.min numpy.pad numpy.percentile numpy.zeros_like numpy.array