def _processParameterNode(self, parameterNode):
    if not self.registrationResult:
      self.registrationResult = RegistrationResult("01: RegistrationResult")
    result = self.registrationResult
    result.fixedVolume = slicer.mrmlScene.GetNodeByID(parameterNode.GetAttribute('FixedImageNodeID'))
    result.fixedLabel = slicer.mrmlScene.GetNodeByID(parameterNode.GetAttribute('FixedLabelNodeID'))
    result.movingLabel = slicer.mrmlScene.GetNodeByID(parameterNode.GetAttribute('MovingLabelNodeID'))
    movingVolume = slicer.mrmlScene.GetNodeByID(parameterNode.GetAttribute('MovingImageNodeID'))
    result.movingVolume = self.volumesLogic.CloneVolume(slicer.mrmlScene, movingVolume, "temp-movingVolume_" + str(self.counter))
    self.counter += 1

    logging.debug("Fixed Image Name: %s" % result.fixedVolume.GetName())
    logging.debug("Fixed Label Name: %s" % result.fixedLabel.GetName())
    logging.debug("Moving Image Name: %s" % movingVolume.GetName())
    logging.debug("Moving Label Name: %s" % result.movingLabel.GetName())
    initialTransform = parameterNode.GetAttribute('InitialTransformNodeID')
    if initialTransform:
      initialTransform = slicer.mrmlScene.GetNodeByID(initialTransform)
      logging.debug("Initial Registration Name: %s" % initialTransform.GetName())
    return result
  def runReRegistration(self, parameterNode, progressCallback=None):
    self.progressCallback = progressCallback
    registrationTypes = ['rigid', 'bSpline']

    if not self.registrationResult:
      self.registrationResult = RegistrationResult("01: RegistrationResult")
    result = self.registrationResult
    result.fixedVolume = slicer.mrmlScene.GetNodeByID(parameterNode.GetAttribute('FixedImageNodeID'))
    fixedLabel = slicer.mrmlScene.GetNodeByID(parameterNode.GetAttribute('FixedLabelNodeID'))
    result.fixedLabel = self.volumesLogic.CreateAndAddLabelVolume(slicer.mrmlScene, result.fixedVolume,
                                                                  result.fixedVolume.GetName() + '-label')
    result.movingLabel = slicer.mrmlScene.GetNodeByID(parameterNode.GetAttribute('MovingLabelNodeID'))
    movingVolume = slicer.mrmlScene.GetNodeByID(parameterNode.GetAttribute('MovingImageNodeID'))
    result.movingVolume = self.volumesLogic.CloneVolume(slicer.mrmlScene, movingVolume, "temp-movingVolume")

    self.createVolumeAndTransformNodes(registrationTypes, prefix=str(result.seriesNumber), suffix=result.suffix)

    initialTransform = parameterNode.GetAttribute('InitialTransformNodeID')
    if initialTransform:
      initialTransform = slicer.mrmlScene.GetNodeByID(initialTransform)

    self.runBRAINSResample(inputVolume=fixedLabel, referenceVolume=result.fixedVolume, outputVolume=result.fixedLabel,
                           warpTransform=initialTransform)
    if initialTransform:
      self.doRigidRegistration(movingBinaryVolume=result.movingLabel, initialTransform=initialTransform)
    else:
      self.doRigidRegistration(movingBinaryVolume=result.movingLabel)

    self.dilateMask(result.fixedLabel)
    self.doBSplineRegistration(initialTransform=result.rigidTransform, useScaleVersor3D=True,
                             useScaleSkewVersor3D=True, useAffine=True)

    targetsNodeID = parameterNode.GetAttribute('TargetsNodeID')
    if targetsNodeID:
      result.originalTargets = slicer.mrmlScene.GetNodeByID(targetsNodeID)
      self.transformTargets(registrationTypes, result.originalTargets, str(result.seriesNumber), suffix=result.suffix)
    result.movingVolume = movingVolume
  def run(self, parameterNode, progressCallback=None):
    self.progressCallback = progressCallback
    registrationTypes = ['rigid', 'affine', 'bSpline']

    if not self.registrationResult:
      self.registrationResult = RegistrationResult("01: RegistrationResult")
    result = self.registrationResult
    result.fixedVolume = slicer.mrmlScene.GetNodeByID(parameterNode.GetAttribute('FixedImageNodeID'))
    result.fixedLabel = slicer.mrmlScene.GetNodeByID(parameterNode.GetAttribute('FixedLabelNodeID'))
    result.movingLabel = slicer.mrmlScene.GetNodeByID(parameterNode.GetAttribute('MovingLabelNodeID'))
    movingVolume = slicer.mrmlScene.GetNodeByID(parameterNode.GetAttribute('MovingImageNodeID'))
    result.movingVolume = self.volumesLogic.CloneVolume(slicer.mrmlScene, movingVolume, "temp-movingVolume")

    self.createVolumeAndTransformNodes(registrationTypes, prefix=str(result.seriesNumber), suffix=result.suffix)

    self.doRigidRegistration(movingBinaryVolume=result.movingLabel, initializeTransformMode="useCenterOfROIAlign")
    self.doAffineRegistration()
    self.doBSplineRegistration(initialTransform=result.affineTransform)

    targetsNodeID = parameterNode.GetAttribute('TargetsNodeID')
    if targetsNodeID:
      result.originalTargets = slicer.mrmlScene.GetNodeByID(targetsNodeID)
      self.transformTargets(registrationTypes, result.originalTargets, str(result.seriesNumber), suffix=result.suffix)
    result.movingVolume = movingVolume
class SliceTrackerRegistrationLogic(ScriptedLoadableModuleLogic, ModuleLogicMixin):

  counter = 1

  def __init__(self):
    ScriptedLoadableModuleLogic.__init__(self)
    self.volumesLogic = slicer.modules.volumes.logic()
    self.markupsLogic = slicer.modules.markups.logic()
    self.registrationResult = None

  def _processParameterNode(self, parameterNode):
    if not self.registrationResult:
      self.registrationResult = RegistrationResult("01: RegistrationResult")
    result = self.registrationResult
    result.fixedVolume = slicer.mrmlScene.GetNodeByID(parameterNode.GetAttribute('FixedImageNodeID'))
    result.fixedLabel = slicer.mrmlScene.GetNodeByID(parameterNode.GetAttribute('FixedLabelNodeID'))
    result.movingLabel = slicer.mrmlScene.GetNodeByID(parameterNode.GetAttribute('MovingLabelNodeID'))
    movingVolume = slicer.mrmlScene.GetNodeByID(parameterNode.GetAttribute('MovingImageNodeID'))
    result.movingVolume = self.volumesLogic.CloneVolume(slicer.mrmlScene, movingVolume, "temp-movingVolume_" + str(self.counter))
    self.counter += 1

    logging.debug("Fixed Image Name: %s" % result.fixedVolume.GetName())
    logging.debug("Fixed Label Name: %s" % result.fixedLabel.GetName())
    logging.debug("Moving Image Name: %s" % movingVolume.GetName())
    logging.debug("Moving Label Name: %s" % result.movingLabel.GetName())
    initialTransform = parameterNode.GetAttribute('InitialTransformNodeID')
    if initialTransform:
      initialTransform = slicer.mrmlScene.GetNodeByID(initialTransform)
      logging.debug("Initial Registration Name: %s" % initialTransform.GetName())
    return result

  def run(self, parameterNode, progressCallback=None):
    self.progressCallback = progressCallback
    result = self._processParameterNode(parameterNode)

    registrationTypes = ['rigid', 'affine', 'bSpline']
    self.createVolumeAndTransformNodes(registrationTypes, prefix=str(result.seriesNumber), suffix=result.suffix)

    self.doRigidRegistration(movingBinaryVolume=result.movingLabel, initializeTransformMode="useCenterOfROIAlign")
    self.doAffineRegistration()
    self.doBSplineRegistration(initialTransform=result.affineTransform)

    targetsNodeID = parameterNode.GetAttribute('TargetsNodeID')
    if targetsNodeID:
      result.originalTargets = slicer.mrmlScene.GetNodeByID(targetsNodeID)
      self.transformTargets(registrationTypes, result.originalTargets, str(result.seriesNumber), suffix=result.suffix)
    result.movingVolume = slicer.mrmlScene.GetNodeByID(parameterNode.GetAttribute('MovingImageNodeID'))

  def runReRegistration(self, parameterNode, progressCallback=None):
    logging.debug("Starting Re-Registration")

    self.progressCallback = progressCallback

    self._processParameterNode(parameterNode)
    result = self.registrationResult

    registrationTypes = ['rigid', 'bSpline']
    self.createVolumeAndTransformNodes(registrationTypes, prefix=str(result.seriesNumber), suffix=result.suffix)
    initialTransform = parameterNode.GetAttribute('InitialTransformNodeID')

    if initialTransform:
      initialTransform = slicer.mrmlScene.GetNodeByID(initialTransform)

    # TODO: label value should be delivered by parameterNode
    self.dilateMask(result.fixedLabel, dilateValue=8)
    self.doRigidRegistration(movingBinaryVolume=result.movingLabel,
                             initialTransform=initialTransform if initialTransform else None)
    self.doBSplineRegistration(initialTransform=result.rigidTransform, useScaleVersor3D=True, useScaleSkewVersor3D=True,
                               useAffine=True)

    targetsNodeID = parameterNode.GetAttribute('TargetsNodeID')
    if targetsNodeID:
      result.originalTargets = slicer.mrmlScene.GetNodeByID(targetsNodeID)
      self.transformTargets(registrationTypes, result.originalTargets, str(result.seriesNumber), suffix=result.suffix)
    result.movingVolume = slicer.mrmlScene.GetNodeByID(parameterNode.GetAttribute('MovingImageNodeID'))

  def createVolumeAndTransformNodes(self, registrationTypes, prefix, suffix=""):
    for regType in registrationTypes:
      self.registrationResult.setVolume(regType, self.createScalarVolumeNode(prefix + '-VOLUME-' + regType + suffix))
      transformName = prefix + '-TRANSFORM-' + regType + suffix
      transform = self.createBSplineTransformNode(transformName) if regType == 'bSpline' \
        else self.createLinearTransformNode(transformName)
      self.registrationResult.setTransform(regType, transform)

  def transformTargets(self, registrations, targets, prefix, suffix=""):
    if targets:
      for registration in registrations:
        name = prefix + '-TARGETS-' + registration + suffix
        clone = self.cloneFiducialAndTransform(name, targets, self.registrationResult.getTransform(registration))
        self.markupsLogic.SetAllMarkupsLocked(clone, True)
        self.registrationResult.setTargets(registration, clone)

  def cloneFiducialAndTransform(self, cloneName, originalTargets, transformNode):
    tfmLogic = slicer.modules.transforms.logic()
    clonedTargets = self.cloneFiducials(originalTargets, cloneName)
    clonedTargets.SetAndObserveTransformNodeID(transformNode.GetID())
    tfmLogic.hardenTransform(clonedTargets)
    return clonedTargets

  def doRigidRegistration(self, **kwargs):
    self.updateProgress(labelText='\nRigid registration', value=2)
    paramsRigid = {'fixedVolume': self.registrationResult.fixedVolume,
                   'movingVolume': self.registrationResult.movingVolume,
                   'fixedBinaryVolume': self.registrationResult.fixedLabel,
                   'outputTransform': self.registrationResult.rigidTransform.GetID(),
                   'outputVolume': self.registrationResult.rigidVolume.GetID(),
                   'maskProcessingMode': "ROI",
                   'useRigid': True}
    for key, value in kwargs.iteritems():
      paramsRigid[key] = value
    slicer.cli.run(slicer.modules.brainsfit, None, paramsRigid, wait_for_completion=True)
    self.registrationResult.cmdArguments += "Rigid Registration Parameters: %s" % str(paramsRigid) + "\n\n"

  def doAffineRegistration(self):
    self.updateProgress(labelText='\nAffine registration', value=2)
    paramsAffine = {'fixedVolume': self.registrationResult.fixedVolume,
                    'movingVolume': self.registrationResult.movingVolume,
                    'fixedBinaryVolume': self.registrationResult.fixedLabel,
                    'movingBinaryVolume': self.registrationResult.movingLabel,
                    'outputTransform': self.registrationResult.affineTransform.GetID(),
                    'outputVolume': self.registrationResult.affineVolume.GetID(),
                    'maskProcessingMode': "ROI",
                    'useAffine': True,
                    'initialTransform': self.registrationResult.rigidTransform}
    slicer.cli.run(slicer.modules.brainsfit, None, paramsAffine, wait_for_completion=True)
    self.registrationResult.cmdArguments += "Affine Registration Parameters: %s" % str(paramsAffine) + "\n\n"

  def doBSplineRegistration(self, initialTransform, **kwargs):
    self.updateProgress(labelText='\nBSpline registration', value=3)
    paramsBSpline = {'fixedVolume': self.registrationResult.fixedVolume,
                     'movingVolume': self.registrationResult.movingVolume,
                     'outputVolume': self.registrationResult.bSplineVolume.GetID(),
                     'bsplineTransform': self.registrationResult.bSplineTransform.GetID(),
                     'fixedBinaryVolume': self.registrationResult.fixedLabel,
                     'movingBinaryVolume': self.registrationResult.movingLabel,
                     'useROIBSpline': True,
                     'useBSpline': True,
                     'splineGridSize': "3,3,3",
                     'maskProcessing': "ROI",
                     'minimumStepLength': "0.005",
                     'maximumStepLength': "0.2",
                     'costFunctionConvergenceFactor': "1.00E+09",
                     'maskProcessingMode': "ROI",
                     'initialTransform': initialTransform}
    for key, value in kwargs.iteritems():
      paramsBSpline[key] = value

    slicer.cli.run(slicer.modules.brainsfit, None, paramsBSpline, wait_for_completion=True)
    self.registrationResult.cmdArguments += "BSpline Registration Parameters: %s" % str(paramsBSpline) + "\n\n"

    self.updateProgress(labelText='\nCompleted registration', value=4)

  def updateProgress(self, **kwargs):
    if self.progressCallback:
      self.progressCallback(**kwargs)