cmtkImagePairNonrigidRegistrationFunctionalTemplate.h

Go to the documentation of this file.

/*
//
//  Copyright 2004-2011 SRI International
//
//  Copyright 1997-2009 Torsten Rohlfing
//
//  This file is part of the Computational Morphometry Toolkit.
//
//  http://www.nitrc.org/projects/cmtk/
//
//  The Computational Morphometry Toolkit is free software: you can
//  redistribute it and/or modify it under the terms of the GNU General Public
//  License as published by the Free Software Foundation, either version 3 of
//  the License, or (at your option) any later version.
//
//  The Computational Morphometry Toolkit is distributed in the hope that it
//  will be useful, but WITHOUT ANY WARRANTY; without even the implied
//  warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//  GNU General Public License for more details.
//
//  You should have received a copy of the GNU General Public License along
//  with the Computational Morphometry Toolkit.  If not, see
//  <http://www.gnu.org/licenses/>.
//
//  $Revision: 2752 $
//
//  $LastChangedDate: 2011-01-17 11:33:31 -0800 (Mon, 17 Jan 2011) $
//
//  $LastChangedBy: torstenrohlfing $
//
*/

#ifndef __cmtkImagePairNonrigidRegistrationFunctionalTemplate_h_included_
#define __cmtkImagePairNonrigidRegistrationFunctionalTemplate_h_included_

#include <cmtkconfig.h>

#include <Registration/cmtkImagePairNonrigidRegistrationFunctional.h>

#include <Base/cmtkSplineWarpXform.h>
#include <Base/cmtkDataTypeTraits.h>

#ifdef CMTK_BUILD_DEMO
#  include <IO/cmtkXformIO.h>
#endif // #ifdef CMTK_BUILD_DEMO

namespace
cmtk
{


template<class VM> 
class ImagePairNonrigidRegistrationFunctionalTemplate
  : public ImagePairNonrigidRegistrationFunctional 
{
protected:
  SmartPointer<VM> m_IncrementalMetric;
  
public:
  typedef ImagePairNonrigidRegistrationFunctionalTemplate<VM> Self;

  typedef SmartPointer<Self> SmartPtr;

  typedef ImagePairNonrigidRegistrationFunctional Superclass;

  ImagePairNonrigidRegistrationFunctionalTemplate<VM>( UniformVolume::SmartPtr& reference, UniformVolume::SmartPtr& floating, const Interpolators::InterpolationEnum interpolation )
    : ImagePairNonrigidRegistrationFunctional( reference, floating )
  {
    this->m_InfoTaskGradient.resize( this->m_NumberOfTasks );
    this->m_InfoTaskComplete.resize( this->m_NumberOfTasks );
    
    this->m_Metric = ImagePairSimilarityMeasure::SmartPtr( new VM( reference, floating, interpolation ) );
    this->m_TaskMetric.resize( this->m_NumberOfThreads, dynamic_cast<const VM&>( *this->m_Metric ) );
  }

  virtual ~ImagePairNonrigidRegistrationFunctionalTemplate<VM>() {}

  virtual void SetForceOutside
  ( const bool flag = true, const Types::DataItem value = 0 )
  {
    this->m_ForceOutsideFlag = flag;
    this->m_ForceOutsideValueRescaled = this->m_Metric->GetFloatingValueScaled( value );
  }

  virtual void SetWarpXform ( SplineWarpXform::SmartPtr& warp )
  {
    Superclass::SetWarpXform( warp );
    this->WarpNeedsFixUpdate = true;
  }
  
  virtual void MatchRefFltIntensities();

  typename Self::ReturnType Evaluate() 
  {
    this->m_Metric->Reset();
    if ( ! this->m_WarpedVolume )
      {
      this->m_WarpedVolume = Memory::AllocateArray<Types::DataItem>(  this->m_DimsX * this->m_DimsY * this->m_DimsZ  );
      }
    
    const size_t numberOfTasks = std::min<size_t>( this->m_NumberOfTasks, this->m_DimsY * this->m_DimsZ );
    for ( size_t taskIdx = 0; taskIdx < numberOfTasks; ++taskIdx ) 
      {
      this->m_InfoTaskComplete[taskIdx].thisObject = this;
      }

    for ( size_t taskIdx = 0; taskIdx < this->m_NumberOfThreads; ++taskIdx ) 
      {
      this->m_TaskMetric[taskIdx].Reset();
      }
    
    ThreadPool::GetGlobalThreadPool().Run( EvaluateCompleteThread, this->m_InfoTaskComplete );
    
    for ( size_t taskIdx = 0; taskIdx < this->m_NumberOfThreads; ++taskIdx ) 
      {
      dynamic_cast<VM&>( *(this->m_Metric) ).Add( this->m_TaskMetric[taskIdx] );
      }
    
    return this->WeightedTotal( this->m_Metric->Get(), *(this->m_ThreadWarp[0]) );
  }

  typename Self::ReturnType EvaluateIncremental( const SplineWarpXform& warp, VM& localMetric, const DataGrid::RegionType& voi, Vector3D *const vectorCache ) 
  {
    Vector3D *pVec;
    int pX, pY, pZ, r;
    int fltIdx[3];
    Types::Coordinate fltFrac[3];

    int endLineIncrement = ( voi.From()[0] + ( this->m_DimsX - voi.To()[0]) );
    int endPlaneIncrement = this->m_DimsX * ( voi.From()[1] + (this->m_DimsY - voi.To()[1]) );
    
    const Types::DataItem unsetY = DataTypeTraits<Types::DataItem>::ChoosePaddingValue();
    localMetric = dynamic_cast<VM&>( *this->m_Metric );
    r = voi.From()[0] + this->m_DimsX * ( voi.From()[1] + this->m_DimsY * voi.From()[2] );
    for ( pZ = voi.From()[2]; pZ<voi.To()[2]; ++pZ ) 
      {
      for ( pY = voi.From()[1]; pY<voi.To()[1]; ++pY ) 
        {
        pVec = vectorCache;
        warp.GetTransformedGridRow( voi.To()[0]-voi.From()[0], pVec, voi.From()[0], pY, pZ );
        for ( pX = voi.From()[0]; pX<voi.To()[0]; ++pX, ++r, ++pVec ) 
          {
          // Remove this sample from incremental metric according to "ground warp" image.
          Types::DataItem sampleX;
          if ( this->m_Metric->GetSampleX( sampleX, r ) )
            {
            if ( this->m_WarpedVolume[r] != unsetY )
              localMetric.Decrement( sampleX, this->m_WarpedVolume[r] );
            
            // Tell us whether the current location is still within the floating volume and get the respective voxel.
            *pVec *= this->m_FloatingInverseDelta;
            if ( this->m_FloatingGrid->FindVoxelByIndex( *pVec, fltIdx, fltFrac ) ) 
              {
              // Continue metric computation.
              localMetric.Increment( sampleX, this->m_Metric->GetSampleY( fltIdx, fltFrac ) );
              } 
            else
              {
              if ( this->m_ForceOutsideFlag )
                {
                localMetric.Increment( sampleX, this->m_ForceOutsideValueRescaled );
                }
              }
            }
          }
        r += endLineIncrement;
        }
      r += endPlaneIncrement;
      }
    
    return localMetric.Get();
  }
  
  virtual typename Self::ReturnType EvaluateWithGradient( CoordinateVector& v, CoordinateVector& g, const typename Self::ParameterType step = 1 ) 
  {
    const typename Self::ReturnType current = this->EvaluateAt( v );

    if ( this->m_AdaptiveFixParameters && this->WarpNeedsFixUpdate ) 
      {
      this->UpdateWarpFixedParameters();
      }
    
    // Make sure we don't create more threads than we have parameters.
    // Actually, we shouldn't create more than the number of ACTIVE parameters.
    // May add this at some point. Anyway, unless we have A LOT of processors,
    // we shouldn't really ever have more threads than active parameters :))
    const size_t numberOfTasks = std::min<size_t>( this->m_NumberOfTasks, this->Dim );
    
    for ( size_t taskIdx = 0; taskIdx < numberOfTasks; ++taskIdx ) 
      {
      this->m_InfoTaskGradient[taskIdx].thisObject = this;
      this->m_InfoTaskGradient[taskIdx].Step = step;
      this->m_InfoTaskGradient[taskIdx].Gradient = g.Elements;
      this->m_InfoTaskGradient[taskIdx].BaseValue = current;
      this->m_InfoTaskGradient[taskIdx].Parameters = &v;
      }

    ThreadPool::GetGlobalThreadPool().Run( EvaluateGradientThread, this->m_InfoTaskGradient );
    
    return current;
  }

  virtual typename Self::ReturnType EvaluateAt ( CoordinateVector& v )
  {
    this->m_ThreadWarp[0]->SetParamVector( v );
    return this->Evaluate();
  }

#ifdef CMTK_BUILD_DEMO

  virtual void SnapshotAt( ParameterVectorType& v )
  {
    this->m_ThreadWarp[0]->SetParamVector( v );
    static int it = 0;
    char path[PATH_MAX];
    snprintf( path, PATH_MAX, "warp-%03d.xform", it++ );
    XformIO::Write( this->m_ThreadWarp[0], path );
  }
#endif

private:
  std::vector<VM> m_TaskMetric;
  
  class EvaluateGradientTaskInfo 
  {
  public:
    Self *thisObject;
    CoordinateVector *Parameters;
    typename Self::ParameterType Step;
    Types::Coordinate *Gradient;
    double BaseValue;
  };
  
  std::vector<typename Self::EvaluateGradientTaskInfo> m_InfoTaskGradient;
  
  static void EvaluateGradientThread( void* arg, const size_t taskIdx, const size_t taskCnt, const size_t threadIdx, const size_t ) 
  {
    typename Self::EvaluateGradientTaskInfo *info = static_cast<typename Self::EvaluateGradientTaskInfo*>( arg );
    
    Self *me = info->thisObject;

    SplineWarpXform& myWarp = *(me->m_ThreadWarp[threadIdx]);
    myWarp.SetParamVector( *info->Parameters );
    
    VM& threadMetric = me->m_TaskMetric[threadIdx];
    Vector3D *vectorCache = me->m_ThreadVectorCache[threadIdx];
    Types::Coordinate *p = myWarp.m_Parameters;
    
    Types::Coordinate pOld;
    double upper, lower;

    const DataGrid::RegionType *voi = me->VolumeOfInfluence + taskIdx;
    for ( size_t dim = taskIdx; dim < me->Dim; dim+=taskCnt, voi+=taskCnt ) 
      {
      if ( me->m_StepScaleVector[dim] <= 0 ) 
        {
        info->Gradient[dim] = 0;
        }
      else
        {
        const typename Self::ParameterType thisStep = info->Step * me->m_StepScaleVector[dim];
        
        pOld = p[dim];
        
        p[dim] += thisStep;
        upper = me->EvaluateIncremental( myWarp, threadMetric, *voi, vectorCache );
        p[dim] = pOld - thisStep;
        lower = me->EvaluateIncremental( myWarp, threadMetric, *voi, vectorCache );
        
        p[dim] = pOld;
        me->WeightedDerivative( lower, upper, myWarp, dim, thisStep );
        
        if ( (upper > info->BaseValue ) || (lower > info->BaseValue) ) 
          {
          // strictly mathematically speaking, we should divide here by step*StepScaleVector[dim], but StepScaleVector[idx] is either zero or a constant independent of idx
          info->Gradient[dim] = upper - lower;
          } 
        else
          {
          info->Gradient[dim] = 0;
          }
        }
      }
  }
  
  class EvaluateCompleteTaskInfo 
  {
  public:
    Self *thisObject;
  };
  
  std::vector<typename Self::EvaluateCompleteTaskInfo> m_InfoTaskComplete;
    
  static void EvaluateCompleteThread ( void *arg, const size_t taskIdx, const size_t taskCnt, const size_t threadIdx, const size_t ) 
  {
    typename Self::EvaluateCompleteTaskInfo *info = static_cast<typename Self::EvaluateCompleteTaskInfo*>( arg );
    
    Self *me = info->thisObject;
    const SplineWarpXform& warp = *(me->m_ThreadWarp[0]);
    VM& threadMetric = me->m_TaskMetric[threadIdx];
    Vector3D *vectorCache = me->m_ThreadVectorCache[threadIdx];
    
    Types::DataItem* warpedVolume = me->m_WarpedVolume;
    const Types::DataItem unsetY = ( me->m_ForceOutsideFlag ) ? me->m_ForceOutsideValueRescaled : DataTypeTraits<Types::DataItem>::ChoosePaddingValue();
    
    Vector3D *pVec;
    int pX, pY, pZ;
    
    int fltIdx[3];
    Types::Coordinate fltFrac[3];
    
    int rowCount = ( me->m_DimsY * me->m_DimsZ );
    int rowFrom = ( rowCount / taskCnt ) * taskIdx;
    int rowTo = ( taskIdx == (taskCnt-1) ) ? rowCount : ( rowCount / taskCnt ) * ( taskIdx + 1 );
    int rowsToDo = rowTo - rowFrom;
    
    int pYfrom = rowFrom % me->m_DimsY;
    int pZfrom = rowFrom / me->m_DimsY;
    
    int r = rowFrom * me->m_DimsX;
    for ( pZ = pZfrom; (pZ < me->m_DimsZ) && rowsToDo; ++pZ ) 
      {
      for ( pY = pYfrom; (pY < me->m_DimsY) && rowsToDo; pYfrom = 0, ++pY, --rowsToDo ) 
        {
        warp.GetTransformedGridRow( me->m_DimsX, vectorCache, 0, pY, pZ );
        pVec = vectorCache;
        for ( pX = 0; pX<me->m_DimsX; ++pX, ++r, ++pVec ) 
          {
          // Tell us whether the current location is still within the 
          // floating volume and get the respective voxel.
          *pVec *= me->m_FloatingInverseDelta;
          if ( me->m_FloatingGrid->FindVoxelByIndex( *pVec, fltIdx, fltFrac ) ) 
            {
            // Continue metric computation.
            warpedVolume[r] = me->m_Metric->GetSampleY( fltIdx, fltFrac );
            
            Types::DataItem value;
            if ( me->m_Metric->GetSampleX( value, r ) )
              {
              threadMetric.Increment( value, warpedVolume[r] );
              }
            } 
          else 
            {
            warpedVolume[r] = unsetY;
            }
          }
        }
      }
  }

private:
  bool WarpNeedsFixUpdate;

  JointHistogram<unsigned int>::SmartPtr m_ConsistencyHistogram;

  void UpdateWarpFixedParameters();
};


} // namespace cmtk

#include "cmtkImagePairNonrigidRegistrationFunctionalTemplate.txx"

#endif // __cmtkImagePairNonrigidRegistrationFunctionalTemplate_h_included_