itkWeightedVotingFusionImageFilter.hxx

/*=========================================================================
 *
 *  Copyright Insight Software Consortium
 *
 *  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.txt
 *
 *  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.
 *
 *=========================================================================*/
#ifndef itkWeightedVotingFusionImageFilter_hxx
#define itkWeightedVotingFusionImageFilter_hxx

#include "itkWeightedVotingFusionImageFilter.h"

#include "itkImageRegionIteratorWithIndex.h"
#include "itkProgressReporter.h"

#include <algorithm>
#include <numeric>

#include <vnl/algo/vnl_cholesky.h>
#include <vnl/algo/vnl_svd.h>
#include <vnl/vnl_inverse.h>

namespace itk {

template<typename TInputImage, typename TOutputImage>
WeightedVotingFusionImageFilter<TInputImage, TOutputImage>
::WeightedVotingFusionImageFilter() :
  m_IsWeightedAveragingComplete( false ),
  m_NumberOfAtlases( 0 ),
  m_NumberOfAtlasSegmentations( 0 ),
  m_NumberOfAtlasModalities( 0 ),
  m_Alpha( 0.1 ),
  m_Beta( 2.0 ),
  m_RetainLabelPosteriorProbabilityImages( false ),
  m_RetainAtlasVotingWeightImages( false ),
  m_ConstrainSolutionToNonnegativeWeights( false )
{
  this->m_MaskImage = nullptr;

  this->m_CountImage = nullptr;

  this->m_NeighborhoodSearchRadiusImage = nullptr;

  this->SetSimilarityMetric( Superclass::PEARSON_CORRELATION );
}

template <typename TInputImage, typename TOutputImage>
void
WeightedVotingFusionImageFilter<TInputImage, TOutputImage>
::UpdateInputs()
{
  // Set all the inputs

  this->SetNumberOfIndexedInputs( this->m_NumberOfAtlases * this->m_NumberOfAtlasModalities +
    this->m_NumberOfAtlasSegmentations + this->m_TargetImage.size() + this->m_LabelExclusionImages.size() );

  SizeValueType nthInput = 0;

  for( SizeValueType i = 0; i < this->m_TargetImage.size(); i++ )
    {
    this->SetNthInput( nthInput++, this->m_TargetImage[i] );
    }

  for( SizeValueType i = 0; i < this->m_NumberOfAtlases; i++ )
    {
    for( SizeValueType j = 0; j < this->m_NumberOfAtlasModalities; j++ )
      {
      this->SetNthInput( nthInput++, this->m_AtlasImages[i][j] );
      }
    }

  for( SizeValueType i = 0; i < this->m_NumberOfAtlasSegmentations; i++ )
    {
    this->SetNthInput( nthInput++, this->m_AtlasSegmentations[i] );
    }

  typename LabelExclusionMap::const_iterator it;
  for( it = m_LabelExclusionImages.begin(); it != m_LabelExclusionImages.end(); ++it )
    {
    this->SetNthInput( nthInput++, it->second );
    }

  if( this->m_MaskImage.IsNotNull() )
    {
    this->SetNthInput( nthInput++, this->m_MaskImage );
    }

  this->Modified();
}

template <typename TInputImage, typename TOutputImage>
void
WeightedVotingFusionImageFilter<TInputImage, TOutputImage>
::GenerateInputRequestedRegion()
{
  Superclass::GenerateInputRequestedRegion();

  // Get the output requested region
  RegionType outRegion = this->GetOutput()->GetRequestedRegion();

  // Pad this region by the search window and patch size
  if( this->m_NeighborhoodSearchRadiusImage.IsNull() )
    {
    outRegion.PadByRadius( this->GetNeighborhoodSearchRadius() );
    }
  else
    {
    NeighborhoodRadiusType maxNeighborhoodSearchRadius;
    maxNeighborhoodSearchRadius.Fill( 0 );

    ImageRegionConstIterator<RadiusImageType> ItR( this->m_NeighborhoodSearchRadiusImage,
      this->m_NeighborhoodSearchRadiusImage->GetRequestedRegion() );
    for( ItR.GoToBegin(); !ItR.IsAtEnd(); ++ItR )
      {
      RadiusValueType localSearchRadius = ItR.Get();
      if( localSearchRadius > maxNeighborhoodSearchRadius[0] )
        {
        maxNeighborhoodSearchRadius.Fill( localSearchRadius );
        }
      }
    outRegion.PadByRadius( maxNeighborhoodSearchRadius );
    }
  outRegion.PadByRadius( this->GetNeighborhoodPatchRadius() );

  // Iterate over all the inputs to this filter

  for( SizeValueType i = 0; i < this->m_TargetImage.size(); i++ )
    {
    InputImageType *input = this->m_TargetImage[i];
    if( i == 0 )
      {
      this->SetTargetImageRegion( input->GetRequestedRegion() );
      }
    RegionType region = outRegion;
    region.Crop( input->GetLargestPossibleRegion() );
    input->SetRequestedRegion( region );
    }

  for( SizeValueType i = 0; i < this->m_NumberOfAtlases; i++ )
    {
    for( SizeValueType j = 0; j < this->m_NumberOfAtlasModalities; j++ )
      {
      InputImageType *input = this->m_AtlasImages[i][j];
      RegionType region = outRegion;
      region.Crop( input->GetLargestPossibleRegion() );
      input->SetRequestedRegion( region );
      }
    }

  for( SizeValueType i = 0; i < this->m_NumberOfAtlasSegmentations; i++ )
    {
    LabelImageType *input = this->m_AtlasSegmentations[i];
    RegionType region = outRegion;
    region.Crop( input->GetLargestPossibleRegion() );
    input->SetRequestedRegion( region );
    }

  typename LabelExclusionMap::const_iterator it;
  for( it = m_LabelExclusionImages.begin(); it != m_LabelExclusionImages.end(); ++it )
    {
    LabelImageType *input = it->second;
    RegionType region = outRegion;
    region.Crop( input->GetLargestPossibleRegion() );
    input->SetRequestedRegion( region );
    }

  if( this->m_MaskImage.IsNotNull() )
    {
    MaskImageType *input = this->m_MaskImage;
    RegionType region = outRegion;
    region.Crop( input->GetLargestPossibleRegion() );
    input->SetRequestedRegion( region );
    }
}

template <typename TInputImage, typename TOutputImage>
void
WeightedVotingFusionImageFilter<TInputImage, TOutputImage>
::GenerateData()
{
  this->BeforeThreadedGenerateData();

  /**
   * Multithread processing for the weighted averaging
   */
  typename ImageSource<TOutputImage>::ThreadStruct str1;
  str1.Filter = this;

//  this->GetMultiThreader()->SetGlobalDefaultNumberOfThreads( this->GetNumberOfThreads() );
  this->GetMultiThreader()->SetSingleMethod( this->ThreaderCallback, &str1 );

  this->GetMultiThreader()->SingleMethodExecute();

  this->m_IsWeightedAveragingComplete = true;

  /**
   * Multithread processing for the image(s) reconstruction
   */

  typename ImageSource<TOutputImage>::ThreadStruct str2;
  str2.Filter = this;

//  this->GetMultiThreader()->SetGlobalDefaultNumberOfThreads( this->GetNumberOfThreads() );
  this->GetMultiThreader()->SetSingleMethod( this->ThreaderCallback, &str2 );

  this->GetMultiThreader()->SingleMethodExecute();

  this->AfterThreadedGenerateData();
}

template <typename TInputImage, typename TOutputImage>
void
WeightedVotingFusionImageFilter<TInputImage, TOutputImage>
::BeforeThreadedGenerateData()
{
  Superclass::BeforeThreadedGenerateData();

  if( this->m_NumberOfAtlasSegmentations != this->m_NumberOfAtlases )
    {
    // Set the number of atlas segmentations to 0 since we're just going to
    // doing joint intensity fusion
    this->m_NumberOfAtlasSegmentations = 0;
    }

  // Check to see if the number of target images is equal to 1 or equal to the number
  // of atlas modalities
  if( this->m_TargetImage.size() != 1 && this->m_TargetImage.size() != this->m_NumberOfAtlasModalities )
    {
    itkExceptionMacro( "The number of target images must be 1 or must be the number of atlas modalities." );
    }

  // Find all the unique labels in the atlas segmentations
  this->m_LabelSet.clear();
  for( unsigned int i = 0; i < this->m_NumberOfAtlasSegmentations; i++ )
    {
    ImageRegionConstIteratorWithIndex<LabelImageType> It(
      this->m_AtlasSegmentations[i], this->m_AtlasSegmentations[i]->GetRequestedRegion() );
    for( It.GoToBegin(); !It.IsAtEnd(); ++It )
      {
      if( !this->m_MaskImage ||
          this->m_MaskImage->GetPixel( It.GetIndex() ) != NumericTraits<LabelType>::ZeroValue() )
        {
        this->m_LabelSet.insert( It.Get() );
        }
      }
    }

  // Initialize the posterior maps
  this->m_LabelPosteriorProbabilityImages.clear();

  typename LabelSetType::const_iterator labelIt;
  for( labelIt = this->m_LabelSet.begin(); labelIt != this->m_LabelSet.end(); ++labelIt )
    {
    typename ProbabilityImageType::Pointer labelProbabilityImage = ProbabilityImageType::New();
    labelProbabilityImage->CopyInformation( this->m_TargetImage[0] );
    labelProbabilityImage->SetRegions( this->m_TargetImage[0]->GetRequestedRegion() );
    labelProbabilityImage->SetLargestPossibleRegion( this->m_TargetImage[0]->GetLargestPossibleRegion() );
    labelProbabilityImage->Allocate( true );

    this->m_LabelPosteriorProbabilityImages.insert(
      std::pair<LabelType, ProbabilityImagePointer>( *labelIt, labelProbabilityImage ) );
    }

  // Initialize the atlas voting weight images
  if( this->m_RetainAtlasVotingWeightImages )
    {
    this->m_AtlasVotingWeightImages.clear();
    this->m_AtlasVotingWeightImages.resize( this->m_NumberOfAtlases );

    for( SizeValueType i = 0; i < this->m_NumberOfAtlases; i++ )
      {
      this->m_AtlasVotingWeightImages[i] = ProbabilityImageType::New();
      this->m_AtlasVotingWeightImages[i]->CopyInformation( this->m_TargetImage[0] );
      this->m_AtlasVotingWeightImages[i]->SetRegions( this->m_TargetImage[0]->GetRequestedRegion() );
      this->m_AtlasVotingWeightImages[i]->SetLargestPossibleRegion( this->m_TargetImage[0]->GetLargestPossibleRegion() );
      this->m_AtlasVotingWeightImages[i]->Allocate( true );
      }
    }

  // Do the joint intensity fusion
  this->m_JointIntensityFusionImage.clear();
  this->m_JointIntensityFusionImage.resize( this->m_NumberOfAtlasModalities );

  for( SizeValueType i = 0; i < this->m_NumberOfAtlasModalities; i++ )
    {
    this->m_JointIntensityFusionImage[i] = InputImageType::New();
    this->m_JointIntensityFusionImage[i]->CopyInformation( this->m_TargetImage[0] );
    this->m_JointIntensityFusionImage[i]->SetRegions( this->m_TargetImage[0]->GetRequestedRegion() );
    this->m_JointIntensityFusionImage[i]->SetLargestPossibleRegion( this->m_TargetImage[0]->GetLargestPossibleRegion() );
    this->m_JointIntensityFusionImage[i]->Allocate( true );
    }

  // Initialize the weight sum image
  this->m_WeightSumImage = ProbabilityImageType::New();
  this->m_WeightSumImage->CopyInformation( this->m_TargetImage[0] );
  this->m_WeightSumImage->SetRegions( this->m_TargetImage[0]->GetRequestedRegion() );
  this->m_WeightSumImage->SetLargestPossibleRegion( this->m_TargetImage[0]->GetLargestPossibleRegion() );
  this->m_WeightSumImage->Allocate( true );

  // Initialize the count image
  this->m_CountImage = CountImageType::New();
  this->m_CountImage->CopyInformation( this->m_TargetImage[0] );
  this->m_CountImage->SetRegions( this->m_TargetImage[0]->GetRequestedRegion() );
  this->m_CountImage->SetLargestPossibleRegion( this->m_TargetImage[0]->GetLargestPossibleRegion() );
  this->m_CountImage->Allocate( true );

  // Determine the ordered search offset list (or map if an search radius image is specified)

  typename InputImageType::SpacingType spacing = this->m_TargetImage[0]->GetSpacing();

  NeighborhoodOffsetListType orderedNeighborhoodSearchOffsetList;
  orderedNeighborhoodSearchOffsetList.clear();

  this->m_NeighborhoodSearchOffsetSetsMap.clear();

  if( this->m_NeighborhoodSearchRadiusImage.IsNull() )
    {
    ConstNeighborhoodIterator<InputImageType> It( this->GetNeighborhoodSearchRadius(),
      this->GetInput(), this->GetInput()->GetRequestedRegion() );

    DistanceIndexVectorType squaredDistances;
    squaredDistances.resize( this->GetNeighborhoodSearchSize() );

    for( unsigned int n = 0; n < this->GetNeighborhoodSearchSize(); n++ )
      {
      NeighborhoodOffsetType offset = ( It.GetNeighborhood() ).GetOffset( n );

      squaredDistances[n].first = n;
      squaredDistances[n].second = 0.0;
      for( unsigned int d = 0; d < ImageDimension; d++ )
        {
        squaredDistances[n].second += itk::Math::sqr ( offset[d] * spacing[d] );
        }
      }
    std::sort( squaredDistances.begin(), squaredDistances.end(), DistanceIndexComparator() );

    for( unsigned int n = 0; n < this->GetNeighborhoodSearchSize(); n++ )
      {
      orderedNeighborhoodSearchOffsetList.push_back( ( It.GetNeighborhood() ).GetOffset( squaredDistances[n].first ) );
      }
    this->SetNeighborhoodSearchOffsetList( orderedNeighborhoodSearchOffsetList );
    }
  else
    {
    ImageRegionConstIterator<RadiusImageType> ItR( this->m_NeighborhoodSearchRadiusImage,
      this->m_NeighborhoodSearchRadiusImage->GetRequestedRegion() );

    for( ItR.GoToBegin(); !ItR.IsAtEnd(); ++ItR )
      {
      RadiusValueType localSearchRadius = ItR.Get();
      if( localSearchRadius > 0 &&
        this->m_NeighborhoodSearchOffsetSetsMap.find( localSearchRadius ) ==
          this->m_NeighborhoodSearchOffsetSetsMap.end() )
        {
        NeighborhoodRadiusType localNeighborhoodSearchRadius;
        localNeighborhoodSearchRadius.Fill( localSearchRadius );

        std::vector<NeighborhoodOffsetType> localNeighborhoodSearchOffsetList;

        ConstNeighborhoodIterator<InputImageType> It( localNeighborhoodSearchRadius,
          this->GetInput(), this->GetInput()->GetRequestedRegion() );

        RadiusValueType localNeighborhoodSearchSize = ( It.GetNeighborhood() ).Size();

        DistanceIndexVectorType squaredDistances;
        squaredDistances.resize( localNeighborhoodSearchSize );

        for( unsigned int n = 0; n < localNeighborhoodSearchSize; n++ )
          {
          NeighborhoodOffsetType offset = ( It.GetNeighborhood() ).GetOffset( n );

          squaredDistances[n].first = n;
          squaredDistances[n].second = 0.0;
          for( unsigned int d = 0; d < ImageDimension; d++ )
            {
            squaredDistances[n].second += itk::Math::sqr ( offset[d] * spacing[d] );
            }
          }
        std::sort( squaredDistances.begin(), squaredDistances.end(), DistanceIndexComparator() );

        for( unsigned int n = 0; n < localNeighborhoodSearchSize; n++ )
          {
          localNeighborhoodSearchOffsetList.push_back( ( It.GetNeighborhood() ).GetOffset( squaredDistances[n].first ) );
          }
        this->m_NeighborhoodSearchOffsetSetsMap[localSearchRadius] = localNeighborhoodSearchOffsetList;
        }
      }
    }

  this->AllocateOutputs();
}

template <typename TInputImage, typename TOutputImage>
void
WeightedVotingFusionImageFilter<TInputImage, TOutputImage>
::ThreadedGenerateData( const RegionType &region, ThreadIdType threadId )
{
  if( !this->m_IsWeightedAveragingComplete )
    {
    this->ThreadedGenerateDataForWeightedAveraging( region, threadId );
    }
  else
    {
    this->ThreadedGenerateDataForReconstruction( region, threadId );
    }
}

template <typename TInputImage, typename TOutputImage>
void
WeightedVotingFusionImageFilter<TInputImage, TOutputImage>
::ThreadedGenerateDataForWeightedAveraging( const RegionType & region, ThreadIdType threadId )
{
  ProgressReporter progress( this, threadId, region.GetNumberOfPixels(), 100 );

  typename OutputImageType::Pointer output = this->GetOutput();

  SizeValueType numberOfTargetModalities = this->m_TargetImage.size();

  MatrixType absoluteAtlasPatchDifferences( this->m_NumberOfAtlases,
    this->GetNeighborhoodPatchSize() * numberOfTargetModalities );

  MatrixType originalAtlasPatchIntensities( this->m_NumberOfAtlases,
    this->GetNeighborhoodPatchSize() * this->m_NumberOfAtlasModalities );

  std::vector<SizeValueType> minimumAtlasOffsetIndices( this->m_NumberOfAtlases );

  bool useOnlyFirstAtlasImage = true;
  if( numberOfTargetModalities == this->m_NumberOfAtlasModalities )
    {
    useOnlyFirstAtlasImage = false;
    }

  std::vector<NeighborhoodOffsetType> searchNeighborhoodOffsetList = this->GetNeighborhoodSearchOffsetList();

  // Iterate over the input region
  ConstNeighborhoodIteratorType ItN( this->GetNeighborhoodPatchRadius(), this->m_TargetImage[0], region );
  for( ItN.GoToBegin(); !ItN.IsAtEnd(); ++ItN )
    {
    progress.CompletedPixel();

    IndexType currentCenterIndex = ItN.GetIndex();

    if( this->m_MaskImage &&
        this->m_MaskImage->GetPixel( currentCenterIndex ) == NumericTraits<LabelType>::ZeroValue() )
      {
      continue;
      }

    // Do not do the following check from Paul's original code.  Since we're incorporating
    // joint intensity fusion, we want to calculate at every voxel (except outside of a
    // possible mask) even if there are no segmentation labels at that voxel.

//     // Check to see if there are any non-zero labels
//     if( this->m_NumberOfAtlasSegmentations > 0 )
//       {
//       bool nonBackgroundLabelExistAtThisVoxel = false;
//       for( SizeValueType i = 0; i < this->m_NumberOfAtlasSegmentations; i++ )
//         {
//         if( this->m_AtlasSegmentations[i]->GetPixel( currentCenterIndex ) > 0 )
//           {
//           nonBackgroundLabelExistAtThisVoxel = true;
//           break;
//           }
//         }
//       if( ! nonBackgroundLabelExistAtThisVoxel )
//         {
//         continue;
//         }
//       }

    // Determine the search neighborhood offset list for the current center voxel
    if( this->m_NeighborhoodSearchRadiusImage.IsNotNull() )
      {
      RadiusValueType localSearchRadius =
        this->m_NeighborhoodSearchRadiusImage->GetPixel( currentCenterIndex );
      if( localSearchRadius <= 0 )
        {
        continue;
        }
      searchNeighborhoodOffsetList = this->m_NeighborhoodSearchOffsetSetsMap[localSearchRadius];
      }
    SizeValueType searchNeighborhoodSize = searchNeighborhoodOffsetList.size();

    InputImagePixelVectorType normalizedTargetPatch =
      this->VectorizeImageListPatch( this->m_TargetImage, currentCenterIndex, true );

    absoluteAtlasPatchDifferences.fill( 0.0 );
    originalAtlasPatchIntensities.fill( 0.0 );

    // In each atlas, search for a patch that matches the target patch
    for( SizeValueType i = 0; i < this->m_NumberOfAtlases; i++ )
      {

      RealType minimumPatchSimilarity = NumericTraits<RealType>::max();
      SizeValueType minimumPatchOffsetIndex = 0;

      for( SizeValueType j = 0; j < searchNeighborhoodSize; j++ )
        {
        IndexType searchIndex = currentCenterIndex + searchNeighborhoodOffsetList[j];

        if( !output->GetRequestedRegion().IsInside( searchIndex ) )
          {
          continue;
          }

        RealType patchSimilarity = this->ComputeNeighborhoodPatchSimilarity(
          this->m_AtlasImages[i], searchIndex, normalizedTargetPatch, useOnlyFirstAtlasImage );

        if( patchSimilarity < minimumPatchSimilarity )
          {
          minimumPatchSimilarity = patchSimilarity;
          minimumPatchOffsetIndex = j;
          }
        }

      // Once the patch has been found, normalize it and then compute the absolute
      // difference with target patch

      IndexType minimumIndex = currentCenterIndex +
        searchNeighborhoodOffsetList[minimumPatchOffsetIndex];
      InputImagePixelVectorType normalizedMinimumAtlasPatch;
      if( numberOfTargetModalities == this->m_NumberOfAtlasModalities )
        {
        normalizedMinimumAtlasPatch =
          this->VectorizeImageListPatch( this->m_AtlasImages[i], minimumIndex, true );
        }
      else
        {
        normalizedMinimumAtlasPatch =
          this->VectorizeImagePatch( this->m_AtlasImages[i][0], minimumIndex, true );
        }

      typename InputImagePixelVectorType::const_iterator itA = normalizedMinimumAtlasPatch.begin();
      typename InputImagePixelVectorType::const_iterator itT = normalizedTargetPatch.begin();
      while( itA != normalizedMinimumAtlasPatch.end() )
        {
        RealType value = std::fabs( *itA - *itT );
        absoluteAtlasPatchDifferences(i, itA - normalizedMinimumAtlasPatch.begin()) = value;

        ++itA;
        ++itT;
        }

      InputImagePixelVectorType originalMinimumAtlasPatch =
        this->VectorizeImageListPatch( this->m_AtlasImages[i], minimumIndex, false );

      typename InputImagePixelVectorType::const_iterator itO = originalMinimumAtlasPatch.begin();
      while( itO != originalMinimumAtlasPatch.end() )
        {
        originalAtlasPatchIntensities(i, itO - originalMinimumAtlasPatch.begin()) = *itO;
        ++itO;
        }

      minimumAtlasOffsetIndices[i] = minimumPatchOffsetIndex;
      }

    // Allocate Mx
    MatrixType Mx( this->m_NumberOfAtlases, this->m_NumberOfAtlases );

    // Compute Mx values
    for( SizeValueType i = 0; i < this->m_NumberOfAtlases; i++ )
      {
      for( SizeValueType j = 0; j <= i; j++ )
        {
        RealType mxValue = 0.0;

        for( unsigned int k = 0; k < this->GetNeighborhoodPatchSize() * numberOfTargetModalities; k++ )
          {
          mxValue += absoluteAtlasPatchDifferences[i][k] * absoluteAtlasPatchDifferences[j][k];
          }
        mxValue /= static_cast<RealType>( this->GetNeighborhoodPatchSize() - 1 );

        if( this->m_Beta == 2.0 )
          {
          mxValue *= mxValue;
          }
        else
          {
          mxValue = std::pow( mxValue, this->m_Beta );
          }

        if( !std::isfinite( mxValue ) )
          {
          mxValue = 0.0;
          }

        Mx(i, j) = Mx(j, i) = mxValue;
        }
      }

    // Compute the weights by solving for the inverse of Mx
    MatrixType MxBar( this->m_NumberOfAtlases, this->m_NumberOfAtlases, 0.0 );
    MxBar.fill_diagonal( this->m_Alpha );
    MxBar += Mx;

    // Define a vector of all ones
    VectorType ones( this->m_NumberOfAtlases, 1.0 );
    VectorType W( this->m_NumberOfAtlases, 1.0 );

    if( this->m_ConstrainSolutionToNonnegativeWeights )
      {
      W = this->NonNegativeLeastSquares( MxBar, ones, 1e-6 );
      }
    else
      {
      vnl_cholesky cholesky( MxBar, vnl_cholesky::estimate_condition );
      if( cholesky.rcond() > itk::Math::sqrteps )
        {
        // well-conditioned matrix
        W = cholesky.solve( ones );
        }
      else
        {
        // ill-conditioned matrix
        W = vnl_svd<RealType>( MxBar ).solve( ones );
        }

      for(double & i : W)
        {
        if( i < 0.0 )
          {
          i = 0.0;
          }
        }
      }

    // Normalize the weights
    W *= 1.0 / dot_product( W, ones );

    // Do joint intensity fusion
    VectorType estimatedNeighborhoodIntensities = W;

    estimatedNeighborhoodIntensities.post_multiply( originalAtlasPatchIntensities );

    for( SizeValueType i = 0; i < this->m_NumberOfAtlasModalities; i++ )
      {
      for( SizeValueType j = 0; j < this->GetNeighborhoodPatchSize(); j++ )
        {
        IndexType neighborhoodIndex = ItN.GetIndex( j );

        if( !output->GetRequestedRegion().IsInside( neighborhoodIndex ) )
          {
          continue;
          }

        if( this->m_MaskImage &&
            this->m_MaskImage->GetPixel( neighborhoodIndex ) == NumericTraits<LabelType>::ZeroValue() )
          {
          continue;
          }

        RealType estimatedValue = (
          estimatedNeighborhoodIntensities[i * this->GetNeighborhoodPatchSize() + j] +
          this->m_JointIntensityFusionImage[i]->GetPixel( neighborhoodIndex ) );

        if( !std::isfinite( estimatedValue ) )
          {
          estimatedValue = 0.0;
          }

        this->m_JointIntensityFusionImage[i]->SetPixel( neighborhoodIndex,
           static_cast<InputImagePixelType>( estimatedValue ) );
        if( i == 0 )
          {
          this->m_CountImage->SetPixel( neighborhoodIndex,
            this->m_CountImage->GetPixel( neighborhoodIndex ) + 1 );
          }
        }
      }

    if( this->m_NumberOfAtlasSegmentations > 0 )
      {
      // Perform voting using Hongzhi's averaging scheme. Iterate over all segmentation patches
      for( SizeValueType n = 0; n < this->GetNeighborhoodPatchSize(); n++ )
        {
        IndexType neighborhoodIndex = ItN.GetIndex( n );
        if( !output->GetRequestedRegion().IsInside( neighborhoodIndex ) )
          {
          continue;
          }

        for( SizeValueType i = 0; i < this->m_NumberOfAtlasSegmentations; i++ )
          {
          // The segmentation at the corresponding patch location in atlas i
          IndexType minimumIndex = neighborhoodIndex +
            searchNeighborhoodOffsetList[minimumAtlasOffsetIndices[i]];

          if( !output->GetRequestedRegion().IsInside( minimumIndex ) )
            {
            continue;
            }

          LabelType label = this->m_AtlasSegmentations[i]->GetPixel( minimumIndex );

          if( this->m_LabelSet.find( label ) == this->m_LabelSet.end() )
            {
            continue;
            }

          // Add that weight the posterior map for voxel at idx
          this->m_LabelPosteriorProbabilityImages[label]->SetPixel( neighborhoodIndex,
            this->m_LabelPosteriorProbabilityImages[label]->GetPixel( neighborhoodIndex ) + W[i] );
          this->m_WeightSumImage->SetPixel( neighborhoodIndex,
            this->m_WeightSumImage->GetPixel( neighborhoodIndex ) + W[i] );

          if( this->m_RetainAtlasVotingWeightImages )
            {
            this->m_AtlasVotingWeightImages[i]->SetPixel( neighborhoodIndex,
              this->m_AtlasVotingWeightImages[i]->GetPixel( neighborhoodIndex ) + W[i] );
            }
          }
        }
      }
    }
}

template <typename TInputImage, typename TOutputImage>
void
WeightedVotingFusionImageFilter<TInputImage, TOutputImage>
::ThreadedGenerateDataForReconstruction( const RegionType &region, ThreadIdType threadId )
{
  ProgressReporter progress( this, threadId, 2 * region.GetNumberOfPixels(), 100 );

  typename OutputImageType::Pointer output = this->GetOutput();

  // Perform voting at each voxel
  ImageRegionIteratorWithIndex<OutputImageType> It( output, region );

  for( It.GoToBegin(); !It.IsAtEnd(); ++It )
    {
    progress.CompletedPixel();

    IndexType index = It.GetIndex();

    if( this->m_MaskImage &&
        this->m_MaskImage->GetPixel( It.GetIndex() ) == NumericTraits<LabelType>::ZeroValue() )
      {
      continue;
      }

    RealType maxPosteriorProbability = 0.0;
    LabelType winningLabel = NumericTraits<LabelType>::ZeroValue();

    typename LabelSetType::const_iterator labelIt;
    for( labelIt = this->m_LabelSet.begin(); labelIt != this->m_LabelSet.end(); ++labelIt )
      {
      // check if the label is excluded
      typename LabelExclusionMap::const_iterator xIt = this->m_LabelExclusionImages.find( *labelIt );
      bool isLabelExcluded = ( xIt != m_LabelExclusionImages.end() && xIt->second->GetPixel( index ) != 0 );

      if( !isLabelExcluded )
        {
        typename ProbabilityImageType::PixelType posteriorProbability =
          this->m_LabelPosteriorProbabilityImages[*labelIt]->GetPixel( index );

        // Vote!
        if( maxPosteriorProbability < posteriorProbability )
          {
          maxPosteriorProbability = posteriorProbability;
          winningLabel = *labelIt;
          }
        }
      }
    It.Set( winningLabel );
    }

  ImageRegionIteratorWithIndex<ProbabilityImageType> ItW( this->m_WeightSumImage,
    region );

  for( ItW.GoToBegin(); !ItW.IsAtEnd(); ++ItW )
    {
    progress.CompletedPixel();

    typename ProbabilityImageType::PixelType weightSum = ItW.Get();

    IndexType index = ItW.GetIndex();

    if( weightSum < 0.1 )
      {
      continue;
      }

    if( this->m_RetainLabelPosteriorProbabilityImages )
      {
      typename LabelSetType::const_iterator labelIt;
      for( labelIt = this->m_LabelSet.begin(); labelIt != this->m_LabelSet.end(); ++labelIt )
        {
        typename ProbabilityImageType::PixelType labelProbability =
          this->m_LabelPosteriorProbabilityImages[*labelIt]->GetPixel( index );
        this->m_LabelPosteriorProbabilityImages[*labelIt]->SetPixel( index, labelProbability / weightSum );
        }
      }

    if( this->m_RetainAtlasVotingWeightImages )
      {
      for( SizeValueType i = 0; i < this->m_NumberOfAtlases; i++ )
        {
        typename ProbabilityImageType::PixelType votingWeight =
          this->m_AtlasVotingWeightImages[i]->GetPixel( index );
        this->m_AtlasVotingWeightImages[i]->SetPixel( index, votingWeight / weightSum );
        }
      }
    }
}

template <typename TInputImage, typename TOutputImage>
void
WeightedVotingFusionImageFilter<TInputImage, TOutputImage>
::AfterThreadedGenerateData()
{
  // Clear posterior maps if not kept
  if( !this->m_RetainLabelPosteriorProbabilityImages )
    {
    this->m_LabelPosteriorProbabilityImages.clear();
    }

  // Normalize the joint intensity fusion images.
  for( SizeValueType i = 0; i < this->m_NumberOfAtlasModalities; i++ )
    {
    ImageRegionIterator<InputImageType> ItJ( this->m_JointIntensityFusionImage[i],
      this->m_JointIntensityFusionImage[i]->GetRequestedRegion() );
    ImageRegionIterator<CountImageType> ItC( this->m_CountImage,
      this->m_CountImage->GetRequestedRegion() );

    for( ItJ.GoToBegin(), ItC.GoToBegin(); !ItJ.IsAtEnd(); ++ItJ, ++ItC )
      {
      typename CountImageType::PixelType count = ItC.Get();

      if( count > 0 )
        {
        ItJ.Set( ItJ.Get() / static_cast<RealType>( count ) );
        }
      }
    }
}

template <typename TInputImage, typename TOutputImage>
typename WeightedVotingFusionImageFilter<TInputImage, TOutputImage>::VectorType
WeightedVotingFusionImageFilter<TInputImage, TOutputImage>
::NonNegativeLeastSquares( const MatrixType A, const VectorType y, const RealType tolerance )
{
  // Algorithm based on
  // Lawson, Charles L.; Hanson, Richard J. (1995). Solving Least Squares Problems. SIAM.
  // cf https://en.wikipedia.org/wiki/Non-negative_least_squares

  SizeValueType m = A.rows();
  SizeValueType n = A.cols();

  // This fortran implementation sets a maximum iteration number of 3 times the
  // number of columns:
  //    http://www.netlib.org/lawson-hanson/all

  const SizeValueType maximumNumberOfIterations = 3 * n;

  // Initialization

  VectorType P( n, 0 );
  VectorType R( n, 1 );
  VectorType x( n, 0 );
  VectorType s( n, 0 );
  VectorType w = A.transpose() * ( y - A * x );

  RealType wMaxValue = w.max_value();
  SizeValueType maxIndex = NumericTraits<SizeValueType>::max();
  wMaxValue = NumericTraits<RealType>::NonpositiveMin();
  for( SizeValueType i = 0; i < n; i++ )
    {
    if( R[i] == 1 && wMaxValue < w[i] )
      {
      maxIndex = i;
      wMaxValue = w[i];
      }
    }

  // Outer loop

  SizeValueType numberOfIterations = 0;
  while( R.sum() > 0 && wMaxValue > tolerance &&
    numberOfIterations++ < maximumNumberOfIterations )
    {
    P[maxIndex] = 1;
    R[maxIndex] = 0;

    SizeValueType sizeP = P.sum();

    MatrixType AP( m, sizeP, 0 );
    SizeValueType jIndex = 0;
    for( SizeValueType j = 0; j < n; j++ )
      {
      if( P[j] == 1 )
        {
        AP.set_column( jIndex++, A.get_column( j ) );
        }
      }

    VectorType sP = vnl_svd<RealType>( AP ).pinverse() * y;

    SizeValueType iIndex = 0;

    for( SizeValueType i = 0; i < n; i++ )
      {
      if( R[i] != 0 )
        {
        s[i] = 0;
        }
      else
        {
        s[i] = sP[iIndex++];
        }
      }

    // Inner loop
    while( sP.min_value() <= tolerance && sizeP > 0 )
      {
      RealType alpha = NumericTraits<RealType>::max();

      for( SizeValueType i = 0; i < n; i++ )
        {
        if( P[i] == 1 && s[i] <= tolerance )
          {
          RealType value = x[i] / ( x[i] - s[i] );
          if( value < alpha )
            {
            alpha = value;
            }
          }
        }

      x += alpha * ( s - x );

      for( SizeValueType i = 0; i < n; i++ )
        {
        if( P[i] == 1 && std::fabs( x[i] ) < tolerance )
          {
          P[i] = 0;
          R[i] = 1;
          }
        }

      sizeP = P.sum();
      if( sizeP == 0 )
        {
        break;
        }

      AP.set_size( m, sizeP );
      jIndex = 0;
      for( SizeValueType j = 0; j < n; j++ )
        {
        if( P[j] == 1 )
          {
          AP.set_column( jIndex++, A.get_column( j ) );
          }
        }

      sP = vnl_svd<RealType>( AP ).pinverse() * y;

      iIndex = 0;
      for( SizeValueType i = 0; i < n; i++ )
        {
        if( R[i] != 0 )
          {
          s[i] = 0;
          }
        else
          {
          s[i] = sP[iIndex++];
          }
        }
      }

    x = s;
    w = A.transpose() * ( y - A * x );

    maxIndex = NumericTraits<SizeValueType>::max();
    wMaxValue = NumericTraits<RealType>::NonpositiveMin();
    for( SizeValueType i = 0; i < n; i++ )
      {
      if( R[i] == 1 && wMaxValue < w[i] )
        {
        maxIndex = i;
        wMaxValue = w[i];
        }
      }
    }

  return x;
}

template <typename TInputImage, typename TOutputImage>
void
WeightedVotingFusionImageFilter<TInputImage, TOutputImage>
::PrintSelf( std::ostream &os, Indent indent ) const
{
  Superclass::PrintSelf( os, indent );

  os << "Number of atlases = " << this->m_NumberOfAtlases << std::endl;
  os << "Number of atlas segmentations = " << this->m_NumberOfAtlasSegmentations << std::endl;
  os << "Number of atlas modalities = " << this->m_NumberOfAtlasModalities << std::endl;
  os << "Alpha = " << this->m_Alpha << std::endl;
  os << "Beta = " << this->m_Beta << std::endl;
  if( this->m_ConstrainSolutionToNonnegativeWeights )
    {
    os << "Constrain solution to positive weights using NNLS." << std::endl;
    }

  os << "Label set: ";
  typename LabelSetType::const_iterator labelIt;
  for( labelIt = this->m_LabelSet.begin(); labelIt != this->m_LabelSet.end(); ++labelIt )
    {
    os << *labelIt << " ";
    }
  os << std::endl;
}

} // end namespace itk

#endif