itkPatchBasedDenoisingImageFilter.h

Go to the documentation of this file.

/*=========================================================================
 *
 *  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 itkPatchBasedDenoisingImageFilter_h
#define itkPatchBasedDenoisingImageFilter_h

#include "itkPatchBasedDenoisingBaseImageFilter.h"
#include "itkImageRegionIterator.h"
#include "itkNeighborhoodAlgorithm.h"
#include "itkVector.h"
#include "itkVectorImage.h"
#include "itkRGBPixel.h"
#include "itkRGBAPixel.h"
#include "itkDiffusionTensor3D.h"
#include "itkFixedArray.h"
#include "itkMatrix.h"
#include "itkRegionConstrainedSubsampler.h"
#include <type_traits>

#include <vector>

namespace itk
{
template <typename TInputImage, typename TOutputImage>
class ITK_TEMPLATE_EXPORT PatchBasedDenoisingImageFilter :
  public PatchBasedDenoisingBaseImageFilter<TInputImage, TOutputImage>
{
public:
  ITK_DISALLOW_COPY_AND_ASSIGN(PatchBasedDenoisingImageFilter);

  using Self = PatchBasedDenoisingImageFilter;
  using Superclass = PatchBasedDenoisingBaseImageFilter<TInputImage, TOutputImage>;
  using Pointer = SmartPointer<Self>;
  using ConstPointer = SmartPointer<const Self>;
  using OutputImagePointer = typename Superclass::OutputImagePointer;

  itkNewMacro(Self);

  itkTypeMacro(PatchBasedDenoisingImageFilter,
    PatchBasedDenoisingBaseImageFilter);

  using InputImageType = typename Superclass::InputImageType;
  using OutputImageType = typename Superclass::OutputImageType;

  static constexpr unsigned int ImageDimension = Superclass::ImageDimension;

  using InputImageRegionType = typename InputImageType::RegionType;

  using OutputImageRegionIteratorType = ImageRegionIterator<OutputImageType>;
  using InputImageRegionConstIteratorType = ImageRegionConstIterator<InputImageType>;

  using PixelType = typename Superclass::PixelType;
  using PixelValueType = typename Superclass::PixelValueType;

  using RealType = typename NumericTraits< PixelType >::RealType;
  using RealValueType = typename NumericTraits< PixelValueType >::RealType;
  using PixelArrayType = Array<PixelValueType>;
  using RealArrayType = Array<RealValueType>;
  using ShortArrayType = Array<unsigned short>;

  using ListAdaptorType = typename Superclass::ListAdaptorType;
  using PatchRadiusType = typename Superclass::PatchRadiusType;
  using InputImagePatchIterator = typename Superclass::InputImagePatchIterator;
  using PatchSampleType = ListAdaptorType;
  using PatchWeightsType = typename Superclass::PatchWeightsType;

  using BaseSamplerType = itk::Statistics::RegionConstrainedSubsampler<
      PatchSampleType, InputImageRegionType >;
  using BaseSamplerPointer = typename BaseSamplerType::Pointer;
  using InstanceIdentifier = typename BaseSamplerType::InstanceIdentifier;

  using EigenValuesArrayType = FixedArray< PixelValueType, 3 >;
  using EigenVectorsMatrixType = Matrix< PixelValueType, 3, 3 >;
  using EigenValuesCacheType = std::vector<EigenValuesArrayType>;
  using EigenVectorsCacheType = std::vector<EigenVectorsMatrixType>;

  struct ThreadDataStruct
    {
    ShortArrayType validDerivatives;
    RealArrayType entropyFirstDerivative;
    RealArrayType entropySecondDerivative;
    ShortArrayType validNorms;
    RealArrayType minNorm;
    RealArrayType maxNorm;
    BaseSamplerPointer sampler;
    EigenValuesCacheType eigenValsCache;
    EigenVectorsCacheType eigenVecsCache;
    };

  itkSetMacro(UseSmoothDiscPatchWeights, bool);
  itkBooleanMacro(UseSmoothDiscPatchWeights);
  itkGetConstMacro(UseSmoothDiscPatchWeights, bool);

  void SetKernelBandwidthSigma(const RealArrayType& kernelSigma);
  itkGetConstMacro(KernelBandwidthSigma, RealArrayType);

  itkSetClampMacro(KernelBandwidthFractionPixelsForEstimation, double, 0.01, 1.0);
  itkGetConstReferenceMacro(KernelBandwidthFractionPixelsForEstimation, double);

  itkSetMacro(ComputeConditionalDerivatives, bool);
  itkBooleanMacro(ComputeConditionalDerivatives);
  itkGetConstMacro(ComputeConditionalDerivatives, bool);

  itkSetMacro(UseFastTensorComputations, bool);
  itkBooleanMacro(UseFastTensorComputations);
  itkGetConstMacro(UseFastTensorComputations, bool);

  static constexpr unsigned int MaxSigmaUpdateIterations = 20;

  itkSetClampMacro(KernelBandwidthMultiplicationFactor, double, 0.01, 100);
  itkGetConstReferenceMacro(KernelBandwidthMultiplicationFactor, double);

  void SetNoiseSigma(const RealType& sigma);

  itkGetConstMacro(NoiseSigma, RealType);

  itkSetObjectMacro(Sampler, BaseSamplerType);
  itkGetModifiableObjectMacro(Sampler, BaseSamplerType);

  itkGetConstMacro(NumIndependentComponents, unsigned int);

protected:
  PatchBasedDenoisingImageFilter();
  ~PatchBasedDenoisingImageFilter() override;
  void PrintSelf(std::ostream& os, Indent indent) const override;

  virtual void EmptyCaches();

  void AllocateUpdateBuffer() override;

  void CopyInputToOutput() override;

  void GenerateInputRequestedRegion() override;

  template<typename T, typename U = void> using DisableIfMultiComponent =
    typename std::enable_if<std::is_same<T, typename NumericTraits<T>::ValueType>::value, U >;

  template<typename T, typename U = void> using EnableIfMultiComponent =
    typename std::enable_if<!std::is_same<T, typename NumericTraits<T>::ValueType>::value, U >;


  template< typename T>
    typename DisableIfMultiComponent<T, T>::type
  GetComponent(const T pix,
               unsigned int itkNotUsed( idx ) ) const
  {
    // The enable if idiom is used to overload this method for both
    // scalars and multi-component types. By exploiting that
    // NumericTraits' ValueType type alias (defines the per-element type
    // for multi-component types ) is different then the parameterize
    // type, the bracket operator is used only for multi-component
    // types.
    return pix;
  }

  template< typename T>
    typename EnableIfMultiComponent<T, typename NumericTraits<T>::ValueType>::type
  GetComponent(const T& pix,
               unsigned int idx ) const
  {
    return pix[idx];
  }

  template< typename T >
  void
  SetComponent( T &pix,
                unsigned int itkNotUsed( idx ),
                typename DisableIfMultiComponent<T, RealValueType>::type val) const
  {
    pix = val;
  }

  template< typename T >
  void
  SetComponent( T &pix,
                unsigned int idx,
                typename EnableIfMultiComponent<T, RealValueType>::type val) const
  {
    pix[idx] =  val;
  }

  void ComputeMinMax(const Image< DiffusionTensor3D<PixelValueType> , ImageDimension>* img)
  {
    if( this->GetComponentSpace() == Superclass::RIEMANNIAN )
      {
      DispatchedRiemannianMinMax(img);
      }
    else
      {
      DispatchedArrayMinMax(img);
      }
  }

  template< typename TImageType >
    typename DisableIfMultiComponent<typename TImageType::PixelType>::type
  ComputeMinMax(const TImageType* img)
  {
    DispatchedMinMax(img);
  }

  template< typename TImageType >
    typename EnableIfMultiComponent<typename TImageType::PixelType>::type
  ComputeMinMax(const TImageType* img)
  {
    DispatchedArrayMinMax(img);
  }

  void ComputeDifferenceAndWeightedSquaredNorm(const DiffusionTensor3D<PixelValueType>& a,
                                               const DiffusionTensor3D<PixelValueType>& b,
                                               const RealArrayType& weight,
                                               bool useCachedComputations,
                                               SizeValueType cacheIndex,
                                               EigenValuesCacheType& eigenValsCache,
                                               EigenVectorsCacheType& eigenVecsCache,
                                               RealType& diff, RealArrayType& norm)
  {
    if( this->GetComponentSpace() == Superclass::RIEMANNIAN )
      {
      ComputeLogMapAndWeightedSquaredGeodesicDifference(a, b, weight,
                                                        useCachedComputations, cacheIndex,
                                                        eigenValsCache, eigenVecsCache,
                                                        diff, norm);
      }
    else
      {
      ComputeSignedEuclideanDifferenceAndWeightedSquaredNorm(a, b, weight,
                                                             useCachedComputations, cacheIndex,
                                                             eigenValsCache, eigenVecsCache,
                                                             diff, norm);
      }
  }

  template <typename PixelT>
  void ComputeDifferenceAndWeightedSquaredNorm(const PixelT& a,
                                               const PixelT& b,
                                               const RealArrayType& weight,
                                               bool useCachedComputations,
                                               SizeValueType cacheIndex,
                                               EigenValuesCacheType& eigenValsCache,
                                               EigenVectorsCacheType& eigenVecsCache,
                                               RealType& diff, RealArrayType& norm)
  {
    ComputeSignedEuclideanDifferenceAndWeightedSquaredNorm(a, b, weight,
                                                           useCachedComputations, cacheIndex,
                                                           eigenValsCache, eigenVecsCache,
                                                           diff, norm);
  }

  RealType AddUpdate(const DiffusionTensor3D<RealValueType>& a,
                     const RealType& b)
  {
    if( this->GetComponentSpace() == Superclass::RIEMANNIAN )
      {
      return this->AddExponentialMapUpdate(a, b);
      }
    else
      {
      return this->AddEuclideanUpdate(a, b);
      }
  }

  template <typename RealT>
  RealType AddUpdate(const RealT& a,
                     const RealType& b)
  {
    return this->AddEuclideanUpdate(a, b);
  }

  virtual void EnforceConstraints();

  void Initialize() override;

  virtual void InitializeKernelSigma();

  void InitializePatchWeights() override;

  virtual void InitializePatchWeightsSmoothDisc();

  void InitializeIteration() override;

  void ComputeKernelBandwidthUpdate() override; // derived from base class;

  // define here

  virtual ThreadDataStruct ThreadedComputeSigmaUpdate(const InputImageRegionType& regionToProcess,
                                                      const int itkNotUsed(threadId),
                                                      ThreadDataStruct threadData);

  virtual RealArrayType ResolveSigmaUpdate();

  void ComputeImageUpdate() override;

  virtual ThreadDataStruct ThreadedComputeImageUpdate(const InputImageRegionType& regionToProcess,
                                                      const int threadId,
                                                      ThreadDataStruct threadData);

  virtual RealType ComputeGradientJointEntropy(InstanceIdentifier id,
                                               typename ListAdaptorType::Pointer& inList,
                                               BaseSamplerPointer& sampler,
                                               ThreadDataStruct& threadData);

  void ApplyUpdate() override;

  virtual void ThreadedApplyUpdate(const InputImageRegionType& regionToProcess,
                                   const int itkNotUsed(threadId) );

  void PostProcessOutput() override;

  virtual void SetThreadData(int threadId, const ThreadDataStruct& data);

  virtual ThreadDataStruct GetThreadData(int threadId);

private:
  static ITK_THREAD_RETURN_FUNCTION_CALL_CONVENTION ComputeSigmaUpdateThreaderCallback( void *arg );

  static ITK_THREAD_RETURN_FUNCTION_CALL_CONVENTION ComputeImageUpdateThreaderCallback( void *arg );

  static ITK_THREAD_RETURN_FUNCTION_CALL_CONVENTION ApplyUpdateThreaderCallback( void *arg );

  template <typename TInputImageType>
  void DispatchedMinMax(const TInputImageType* img);

  template <typename TInputImageType>
  void DispatchedArrayMinMax(const TInputImageType* img);

  template <typename TInputImageType>
  void DispatchedVectorMinMax(const TInputImageType* img);

  template <typename TInputImageType>
  void DispatchedRiemannianMinMax(const TInputImageType* img);

  static ITK_THREAD_RETURN_FUNCTION_CALL_CONVENTION RiemannianMinMaxThreaderCallback( void *arg );

  ThreadDataStruct ThreadedRiemannianMinMax(const InputImageRegionType& regionToProcess,
                                            const int itkNotUsed(threadId),
                                            const InputImageType* img,
                                            ThreadDataStruct threadData);

  virtual void ResolveRiemannianMinMax();

  void ComputeSignedEuclideanDifferenceAndWeightedSquaredNorm(const PixelType& a, const PixelType& b,
                                                              const RealArrayType& weight,
                                                              bool useCachedComputations,
                                                              SizeValueType cacheIndex,
                                                              EigenValuesCacheType& eigenValsCache,
                                                              EigenVectorsCacheType& eigenVecsCache,
                                                              RealType& diff, RealArrayType& norm);

  void ComputeLogMapAndWeightedSquaredGeodesicDifference(const DiffusionTensor3D<PixelValueType>& spdMatrixA,
                                                         const DiffusionTensor3D<PixelValueType>& spdMatrixB,
                                                         const RealArrayType& weight,
                                                         bool useCachedComputations,
                                                         SizeValueType cacheIndex,
                                                         EigenValuesCacheType& eigenValsCache,
                                                         EigenVectorsCacheType& eigenVecsCache,
                                                         RealType& symMatrixLogMap, RealArrayType& geodesicDist);

  template <typename TensorValueT>
  void Compute3x3EigenAnalysis(const DiffusionTensor3D<TensorValueT>& spdMatrix,
                               FixedArray< TensorValueT, 3 >&  eigenVals,
                               Matrix< TensorValueT, 3, 3 >& eigenVecs);

  RealType AddEuclideanUpdate(const RealType& a, const RealType& b);

  RealType AddExponentialMapUpdate(const DiffusionTensor3D<RealValueType>& spdMatrix,
                                   const DiffusionTensor3D<RealValueType>& symMatrix);

  struct ThreadFilterStruct
    {
    PatchBasedDenoisingImageFilter *Filter;
    InputImageType *Img;
    };

  std::vector<ThreadDataStruct> m_ThreadData;

  typename OutputImageType::Pointer m_UpdateBuffer;

  unsigned int m_NumPixelComponents{ 0 };
  unsigned int m_NumIndependentComponents{ 0 };
  unsigned int m_TotalNumberPixels{ 0 };

  bool m_UseSmoothDiscPatchWeights{ true };

  bool m_UseFastTensorComputations{ true };

  RealArrayType  m_KernelBandwidthSigma;
  bool           m_KernelBandwidthSigmaIsSet{ false };
  RealArrayType  m_IntensityRescaleInvFactor;
  PixelType      m_ZeroPixel;
  PixelArrayType m_ImageMin;
  PixelArrayType m_ImageMax;
  double         m_KernelBandwidthFractionPixelsForEstimation{ 0.20 };
  bool           m_ComputeConditionalDerivatives{ false };
  double         m_MinSigma;
  double         m_MinProbability;
  unsigned int   m_SigmaUpdateDecimationFactor;
  double         m_SigmaUpdateConvergenceTolerance{ 0.01 };
  ShortArrayType m_SigmaConverged;
  double         m_KernelBandwidthMultiplicationFactor{ 1.0 };

  RealType m_NoiseSigma;
  RealType m_NoiseSigmaSquared;
  bool     m_NoiseSigmaIsSet{ false };

  BaseSamplerPointer                m_Sampler;
  typename ListAdaptorType::Pointer m_SearchSpaceList;
};
} // end namespace itk

#ifndef ITK_MANUAL_INSTANTIATION
#include "itkPatchBasedDenoisingImageFilter.hxx"
#endif

#endif