mitkTrackingForestHandler.h

Go to the documentation of this file.

/*===================================================================

The Medical Imaging Interaction Toolkit (MITK)

Copyright (c) German Cancer Research Center,
Division of Medical and Biological Informatics.
All rights reserved.

This software is distributed WITHOUT ANY WARRANTY; without
even the implied warranty of MERCHANTABILITY or FITNESS FOR
A PARTICULAR PURPOSE.

See LICENSE.txt or http://www.mitk.org for details.

===================================================================*/

#ifndef _TrackingForestHandler
#define _TrackingForestHandler

#include "mitkBaseData.h"

#include <mitkFiberBundle.h>
#include <itkAnalyticalDiffusionQballReconstructionImageFilter.h>
#include <itkOrientationDistributionFunction.h>
#include <itkMersenneTwisterRandomVariateGenerator.h>
#include <itkOrientationDistributionFunction.h>
#include <chrono>

#undef DIFFERENCE
#define VIGRA_STATIC_LIB
#include <vigra/random_forest.hxx>
#include <vigra/multi_array.hxx>
#include <vigra/random_forest_hdf5_impex.hxx>

#define _USE_MATH_DEFINES
#include <math.h>

namespace mitk
{

template< int ShOrder=6, int NumberOfSignalFeatures=100 >
class TrackingForestHandler
{

public:

    TrackingForestHandler();
    ~TrackingForestHandler();

    typedef itk::Image<unsigned char, 3>                                        ItkUcharImgType;
    typedef itk::Image< itk::Vector< float, NumberOfSignalFeatures*2 > , 3 >    InterpolatedRawImageType;
    typedef itk::Image< Vector< float, NumberOfSignalFeatures > , 3 >           FeatureImageType;

    void SetRawData( std::vector< Image::Pointer > images ){ m_RawData = images; }
    void AddRawData( Image::Pointer img ){ m_RawData.push_back(img); }
    void SetTractograms( std::vector< FiberBundle::Pointer > tractograms )
    {
        m_Tractograms.clear();
        for (auto fib : tractograms)
        {
            m_Tractograms.push_back(fib->GetDeepCopy());
        }
    }
    void SetMaskImages( std::vector< ItkUcharImgType::Pointer > images ){ m_MaskImages = images; }
    void SetWhiteMatterImages( std::vector< ItkUcharImgType::Pointer > images ){ m_WhiteMatterImages = images; }

    void StartTraining();
    void SaveForest(std::string forestFile);
    void LoadForest(std::string forestFile);

    // training parameters

    void SetNumTrees(int num){ m_NumTrees = num; }
    void SetMaxTreeDepth(int depth){ m_MaxTreeDepth = depth; }
    void SetStepSize(double step){ m_WmSampleDistance = step; }
    void SetGrayMatterSamplesPerVoxel(int samples){ m_GmSamplesPerVoxel = samples; }
    void SetSampleFraction(double fraction){ m_SampleFraction = fraction; }
    std::shared_ptr< vigra::RandomForest<int> > GetForest(){ return m_Forest; }

    void InitForTracking();     
    vnl_vector_fixed<double,3> Classify(itk::Point<double, 3>& pos, int& candidates, vnl_vector_fixed<double,3>& olddir, double angularThreshold, double& w, ItkUcharImgType::Pointer mask=nullptr);  

    bool IsForestValid();   

protected:

    // tracking
    void InputDataValidForTracking();                                                   
    typename FeatureImageType::PixelType GetFeatureValues(itk::Point<float, 3> itkP);   

    // training
    void InputDataValidForTraining();       
    void PreprocessInputDataForTraining();  
    void CalculateFeaturesForTraining();    
    void TrainForest();                     
    typename InterpolatedRawImageType::PixelType GetImageValues(itk::Point<float, 3> itkP, typename InterpolatedRawImageType::Pointer image);   


    std::vector< Image::Pointer >               m_RawData;  
    std::shared_ptr< vigra::RandomForest<int> > m_Forest;   
    std::chrono::time_point<std::chrono::system_clock> m_StartTime;
    std::chrono::time_point<std::chrono::system_clock> m_EndTime;

    // only for training
    std::vector< FiberBundle::Pointer >                         m_Tractograms;              
    std::vector< ItkUcharImgType::Pointer >                     m_MaskImages;               
    std::vector< ItkUcharImgType::Pointer >                     m_WhiteMatterImages;        
    std::vector< typename InterpolatedRawImageType::Pointer >   m_InterpolatedRawImages;    
    double                                                      m_WmSampleDistance;         
    int                                                         m_NumTrees;                 
    int                                                         m_MaxTreeDepth;             
    double                                                      m_SampleFraction;           
    unsigned int                                                m_NumberOfSamples;          
    std::vector< unsigned int >                                 m_GmSamples;                
    int                                                         m_GmSamplesPerVoxel;        
    vigra::MultiArray<2, double>                                m_FeatureData;              

    // only for tracking
    typename FeatureImageType::Pointer                          m_FeatureImage;             
    vigra::MultiArray<2, double>                                m_LabelData;                
    std::vector< int >                                          m_DirectionIndices;         
    itk::OrientationDistributionFunction< double, NumberOfSignalFeatures*2 >  m_DirContainer;   
};

}

#include "mitkTrackingForestHandler.cpp"

#endif