itkImpactImageToImageMetric.h

Go to the documentation of this file.

/*=========================================================================
 *
 *  Copyright UMC Utrecht and contributors
 *
 *  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 itkImpactImageToImageMetric_h
#define itkImpactImageToImageMetric_h
 
#include "itkAdvancedImageToImageMetric.h"
#include "itkBSplineInterpolateImageFunction.h"
#include "itkBSplineInterpolateVectorImageFunction.h"
#include "itkImpactModelConfiguration.h"
#include "ImpactTensorUtils.h"
#include "ImpactLoss.h"
 
#include <torch/script.h>
#include <torch/torch.h>
 
#include <string>
#include <vector>
#include <random>
 
namespace itk
{

 
template <typename TFixedImage, typename TMovingImage>
class ITK_TEMPLATE_EXPORT ImpactImageToImageMetric : public AdvancedImageToImageMetric<TFixedImage, TMovingImage>
{
public:
  ITK_DISALLOW_COPY_AND_MOVE(ImpactImageToImageMetric);

  using Self = ImpactImageToImageMetric;
  using Superclass = AdvancedImageToImageMetric<TFixedImage, TMovingImage>;
  using Pointer = SmartPointer<Self>;
  using ConstPointer = SmartPointer<const Self>;

  itkNewMacro(Self);

  itkTypeMacro(ImpactImageToImageMetric, AdvancedImageToImageMetric);

  using typename Superclass::CoordinateRepresentationType;
  using typename Superclass::MovingImageType;
  using typename Superclass::MovingImagePixelType;
  using typename Superclass::MovingImageConstPointer;
  using typename Superclass::FixedImageType;
  using typename Superclass::FixedImageConstPointer;
  using typename Superclass::FixedImageRegionType;
  using typename Superclass::TransformType;
  using typename Superclass::TransformPointer;
  using typename Superclass::InputPointType;
  using typename Superclass::OutputPointType;
  using typename Superclass::TransformJacobianType;
  using typename Superclass::NumberOfParametersType;
  using typename Superclass::InterpolatorType;
  using typename Superclass::InterpolatorPointer;
  using typename Superclass::RealType;
  using typename Superclass::GradientPixelType;
  using typename Superclass::GradientImageType;
  using typename Superclass::GradientImagePointer;
  using typename Superclass::FixedImageMaskType;
  using typename Superclass::FixedImageMaskPointer;
  using typename Superclass::MovingImageMaskType;
  using typename Superclass::MovingImageMaskPointer;
  using typename Superclass::MeasureType;
  using typename Superclass::DerivativeType;
  using typename Superclass::DerivativeValueType;
  using typename Superclass::ParametersType;
  using typename Superclass::FixedImagePixelType;
  using typename Superclass::MovingImageRegionType;
  using typename Superclass::ImageSamplerType;
  using typename Superclass::ImageSamplerPointer;
  using typename Superclass::ImageSampleContainerType;
  using typename Superclass::ImageSampleContainerPointer;
  using typename Superclass::FixedImageLimiterType;
  using typename Superclass::MovingImageLimiterType;
  using typename Superclass::FixedImageLimiterOutputType;
  using typename Superclass::MovingImageLimiterOutputType;
  using typename Superclass::MovingImageDerivativeScalesType;
  using typename Superclass::ThreadInfoType;

  itkStaticConstMacro(FixedImageDimension, unsigned int, FixedImageType::ImageDimension);

  itkStaticConstMacro(MovingImageDimension, unsigned int, MovingImageType::ImageDimension);

  virtual MeasureType
  GetValueSingleThreaded(const ParametersType & parameters) const;

  MeasureType
  GetValue(const ParametersType & parameters) const override;

  void
  GetDerivative(const ParametersType & parameters, DerivativeType & derivative) const override;

  void
  GetValueAndDerivativeSingleThreaded(const ParametersType & parameters,
                                      MeasureType &          value,
                                      DerivativeType &       derivative) const;

  void
  GetValueAndDerivative(const ParametersType & parameters,
                        MeasureType &          value,
                        DerivativeType &       derivative) const override;

  void
  Initialize() override;

  itkSetMacro(FixedModelsConfiguration, std::vector<ImpactModelConfiguration>);
  itkGetConstReferenceMacro(FixedModelsConfiguration, std::vector<ImpactModelConfiguration>);

  itkSetMacro(MovingModelsConfiguration, std::vector<ImpactModelConfiguration>);
  itkGetConstReferenceMacro(MovingModelsConfiguration, std::vector<ImpactModelConfiguration>);

  itkSetMacro(SubsetFeatures, std::vector<unsigned int>);
  itkGetConstMacro(SubsetFeatures, std::vector<unsigned int>);

  itkSetMacro(LayersWeight, std::vector<float>);
  itkGetConstMacro(LayersWeight, std::vector<float>);
 

  itkSetMacro(Distance, std::vector<std::string>);
  itkGetConstMacro(Distance, std::vector<std::string>);

  itkSetMacro(PCA, std::vector<unsigned int>);
  itkGetConstMacro(PCA, std::vector<unsigned int>);

  itkSetMacro(Device, torch::Device);
  itkGetConstMacro(Device, torch::Device);

  itkSetMacro(UseMixedPrecision, bool);
  itkGetConstMacro(UseMixedPrecision, bool);

  itkSetMacro(WriteFeatureMaps, bool);
  itkGetConstMacro(WriteFeatureMaps, bool);

  itkSetMacro(FeatureMapsPath, std::string);
  itkGetConstMacro(FeatureMapsPath, std::string);

  itkSetMacro(Mode, std::string);
  itkGetConstMacro(Mode, std::string);

  itkSetMacro(CurrentLevel, unsigned int);
  itkGetConstMacro(CurrentLevel, unsigned int);

  itkSetMacro(Seed, unsigned int);
  itkGetConstMacro(Seed, unsigned int);

  itkSetMacro(FeaturesMapUpdateInterval, int);
  itkGetConstMacro(FeaturesMapUpdateInterval, int);
 
protected:
  ImpactImageToImageMetric();
  ~ImpactImageToImageMetric() override = default;

  void
  InitializeThreadingParameters() const override;


  struct LossPerThreadStruct
  {
    std::vector<std::unique_ptr<ImpactLoss::Loss>> m_Losses;
    std::vector<float>                             m_LayersWeight;
    SizeValueType                                  m_NumberOfPixelsCounted;
    int                                            m_NumberOfParameters;
    std::mt19937                                   m_RandomGenerator;
 
    void
    init(std::vector<std::string> distanceName, std::vector<float> layersWeight, unsigned int seed)
    {
      if (seed > 0)
      {
        m_RandomGenerator = std::mt19937(seed);
      }
      else
      {
        m_RandomGenerator = std::mt19937(time(nullptr));
      }
      m_LayersWeight = layersWeight;
      for (std::string name : distanceName)
      {
        m_Losses.push_back(ImpactLoss::LossFactory::Instance().Create(name));
      }
    }
 
    void
    setNumberOfParameters(int numberOfParameters)
    {
      m_NumberOfParameters = numberOfParameters;
      for (int l = 0; l < m_LayersWeight.size(); ++l)
      {
        m_Losses[l]->setNumberOfParameters(numberOfParameters);
      }
    }
 
    void
    reset()
    {
      m_NumberOfPixelsCounted = 0;
      for (std::unique_ptr<ImpactLoss::Loss> & loss : m_Losses)
      {
        loss->reset();
      }
    }
 
    double
    GetValue()
    {
      MeasureType value = MeasureType{};
      for (int l = 0; l < m_LayersWeight.size(); ++l)
      {
        value += m_LayersWeight[l] * m_Losses[l]->GetValue(static_cast<double>(m_NumberOfPixelsCounted));
      }
      return value;
    }
 
    DerivativeType
    GetDerivative()
    {
      DerivativeType derivative = DerivativeType(m_NumberOfParameters);
      derivative.Fill(DerivativeValueType{});
      for (int l = 0; l < m_LayersWeight.size(); ++l)
      {
        torch::Tensor d = m_LayersWeight[l] * m_Losses[l]->GetDerivative(static_cast<double>(m_NumberOfPixelsCounted));
        for (int i = 0; i < d.size(0); ++i)
        {
          derivative[i] += d[i].item<float>();
        }
      }
      return derivative;
    }
 
    LossPerThreadStruct &
    operator+=(const LossPerThreadStruct & other)
    {
      const auto * lossPerThreadStructOther = dynamic_cast<const LossPerThreadStruct *>(&other);
      if (lossPerThreadStructOther)
      {
        m_NumberOfPixelsCounted += lossPerThreadStructOther->m_NumberOfPixelsCounted;
        for (int i = 0; i < lossPerThreadStructOther->m_Losses.size(); ++i)
        {
          *m_Losses[i] += *lossPerThreadStructOther->m_Losses[i];
        }
      }
      return *this;
    }
  };

  using typename Superclass::FixedImageIndexType;
  using typename Superclass::FixedImageIndexValueType;
  using typename Superclass::MovingImageIndexType;
  using typename Superclass::FixedImagePointType;
  using typename Superclass::MovingImagePointType;
  using typename Superclass::MovingImageContinuousIndexType;
  using typename Superclass::BSplineInterpolatorType;
  using typename Superclass::MovingImageDerivativeType;
  using typename Superclass::NonZeroJacobianIndicesType;

  bool
  SampleCheck(const FixedImagePointType &             fixedImageCenterCoordinate,
              const std::vector<std::vector<float>> & patchIndex) const;

  bool
  SampleCheck(const FixedImagePointType & fixedImageCenterCoordinate) const;

  void
  ThreadedGetValue(ThreadIdType threadID) const override;

  void
  AfterThreadedGetValue(MeasureType & value) const override;

  void
  ThreadedGetValueAndDerivative(ThreadIdType threadID) const override;

  void
  AfterThreadedGetValueAndDerivative(MeasureType & value, DerivativeType & derivative) const override;

  unsigned int
  ComputeValue(const std::vector<FixedImagePointType> & fixedPoints, LossPerThreadStruct & loss) const;

  unsigned int
  ComputeValueStatic(const std::vector<FixedImagePointType> & fixedPoints, LossPerThreadStruct & loss) const;

  unsigned int
  ComputeValueAndDerivativeJacobian(const std::vector<FixedImagePointType> & fixedPoints,
                                    LossPerThreadStruct &                    loss) const;

  unsigned int
  ComputeValueAndDerivativeStatic(const std::vector<FixedImagePointType> & fixedPoints,
                                  LossPerThreadStruct &                    loss) const;
 

  void
  UpdateFixedFeaturesMaps();

  void
  UpdateMovingFeaturesMaps();
 
private:
  using FixedInterpolatorType = BSplineInterpolateImageFunction<FixedImageType, CoordinateRepresentationType, double>;
  using FeaturesImageType = itk::VectorImage<float, FixedImageDimension>;
  using FeaturesInterpolatorType = BSplineInterpolateVectorImageFunction<
    FeaturesImageType,
    BSplineInterpolateImageFunction<itk::Image<float, FixedImageDimension>, CoordinateRepresentationType, float>>;

  struct FeaturesMaps
  {
    typename FeaturesImageType::Pointer m_FeaturesMaps;
    FeaturesInterpolatorType            m_FeaturesMapsInterpolator;
 
    FeaturesMaps(typename FeaturesImageType::Pointer featuresMaps)
      : m_FeaturesMaps(featuresMaps)
    {
      m_FeaturesMapsInterpolator = FeaturesInterpolatorType();
      m_FeaturesMapsInterpolator.SetInputImage(featuresMaps);
    }
  };
 
  using FeaturesMaps = typename ImpactImageToImageMetric<TFixedImage, TMovingImage>::FeaturesMaps;

  torch::Tensor
  EvaluateFixedImagesPatchValue(const FixedImagePointType &             fixedImageCenterCoordinate,
                                const std::vector<std::vector<float>> & patchIndex,
                                const std::vector<int64_t> &            patchSize) const;

  torch::Tensor
  EvaluateMovingImagesPatchValue(const FixedImagePointType &             fixedImageCenterCoordinate,
                                 const std::vector<std::vector<float>> & patchIndex,
                                 const std::vector<int64_t> &            patchSize) const;

  torch::Tensor
  EvaluateMovingImagesPatchValuesAndJacobians(const FixedImagePointType &             fixedImageCenterCoordinate,
                                              torch::Tensor &                         movingImagesPatchesJacobians,
                                              const std::vector<std::vector<float>> & patchIndex,
                                              const std::vector<int64_t> &            patchSize,
                                              int                                     s) const;

  template <typename ImagePointType>
  std::vector<ImagePointType>
  GeneratePatchIndex(const std::vector<ImpactModelConfiguration> &               modelConfig,
                     std::mt19937 &                                              randomGenerator,
                     const std::vector<ImagePointType> &                         fixedPointsTmp,
                     std::vector<std::vector<std::vector<std::vector<float>>>> & patchIndex) const;

  std::vector<ImpactModelConfiguration> m_FixedModelsConfiguration;
  std::vector<ImpactModelConfiguration> m_MovingModelsConfiguration;
 
  std::vector<unsigned int> m_SubsetFeatures;
  std::vector<unsigned int> m_PCA;
  std::vector<float>        m_LayersWeight;
  std::vector<std::string>  m_Distance;
  int                       m_FeaturesMapUpdateInterval;
  std::string               m_Mode;
  bool                      m_WriteFeatureMaps;
  std::string               m_FeatureMapsPath;
  torch::Device             m_Device = torch::Device(torch::kCPU);
  bool                      m_UseMixedPrecision;
  unsigned int              m_CurrentLevel;
  unsigned int              m_Seed;
 
 
  std::vector<FeaturesMaps>  m_FixedFeaturesMaps;
  std::vector<FeaturesMaps>  m_MovingFeaturesMaps;
  std::vector<torch::Tensor> m_PrincipalComponents;
 
  std::vector<std::vector<unsigned int>> m_FeaturesIndexes;
 

  InterpolatorPointer m_FixedInterpolator = [this] {
    const auto interpolator = FixedInterpolatorType::New();
    interpolator->SetSplineOrder(3);
    return interpolator;
  }();

  std::vector<unsigned int>
  GetSubsetOfFeatures(const std::vector<unsigned int> & featuresIndex, std::mt19937 & randomGenerator, int n) const;

  itkPadStruct(ITK_CACHE_LINE_ALIGNMENT, LossPerThreadStruct, PaddedLossPerThreadStruct);
 
  itkAlignedTypedef(ITK_CACHE_LINE_ALIGNMENT, PaddedLossPerThreadStruct, AlignedLossPerThreadStruct);

  mutable std::unique_ptr<AlignedLossPerThreadStruct[]> m_LossThreadStruct{ nullptr };
 
  mutable int m_LossThreadStructSize = 0;
};
 
} // end namespace itk
 
#ifndef ITK_MANUAL_INSTANTIATION
#  include "itkImpactImageToImageMetric.hxx"
#endif
 
#endif // end #ifndef itkImpactImageToImageMetric_h