diff --git a/Modules/DiffusionImaging/DiffusionIO/files.cmake b/Modules/DiffusionImaging/DiffusionIO/files.cmake
index 20174ef77d..209e2cf4ff 100644
--- a/Modules/DiffusionImaging/DiffusionIO/files.cmake
+++ b/Modules/DiffusionImaging/DiffusionIO/files.cmake
@@ -1,36 +1,38 @@
 set(CPP_FILES
   mitkDiffusionModuleActivator.cpp
   mitkNrrdTbssImageWriterFactory.cpp
   #mitkFiberBundleIOFactory.cpp
   mitkConnectomicsNetworkReader.cpp
   mitkConnectomicsNetworkWriter.cpp
   mitkConnectomicsNetworkCSVWriter.cpp
   mitkConnectomicsNetworkMatrixWriter.cpp
   mitkConnectomicsNetworkSerializer.cpp
   mitkConnectomicsNetworkDefinitions.cpp
   mitkNrrdTbssRoiImageIOFactory.cpp
   #mitkFiberBundleWriterFactory.cpp
   mitkNrrdTbssRoiImageWriterFactory.cpp
   mitkFiberTrackingObjectFactory.cpp
   mitkConnectomicsObjectFactory.cpp
   mitkQuantificationObjectFactory.cpp
   mitkNrrdTbssImageIOFactory.cpp
 
   mitkDiffusionIOMimeTypes.cpp
 
   mitkFiberBundleTckReader.cpp
   mitkFiberBundleTrackVisReader.cpp
   mitkFiberBundleTrackVisWriter.cpp
   mitkFiberBundleVtkReader.cpp
   mitkFiberBundleVtkWriter.cpp
   mitkFiberBundleSerializer.cpp
   mitkFiberBundleMapper2D.cpp
   mitkFiberBundleMapper3D.cpp
+  
   mitkTractographyForestReader.cpp
   mitkTractographyForestWriter.cpp
+  mitkTractographyForestSerializer.cpp
 
   mitkPlanarFigureCompositeWriter.cpp
   mitkPlanarFigureCompositeReader.cpp
   mitkPlanarFigureCompositeSerializer.cpp
 )
 
diff --git a/Modules/DiffusionImaging/DiffusionIO/mitkPlanarFigureCompositeSerializer.h b/Modules/DiffusionImaging/DiffusionIO/mitkPlanarFigureCompositeSerializer.h
index 24184e2a5f..28e9b29a5e 100644
--- a/Modules/DiffusionImaging/DiffusionIO/mitkPlanarFigureCompositeSerializer.h
+++ b/Modules/DiffusionImaging/DiffusionIO/mitkPlanarFigureCompositeSerializer.h
@@ -1,39 +1,39 @@
 /*===================================================================
 
 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 mitkPlanarFigureCompositeSerializer_h_included
 #define mitkPlanarFigureCompositeSerializer_h_included
 
 #include "mitkBaseDataSerializer.h"
 
 namespace mitk
 {
 /**
-  \brief Serializes mitk::Surface for mitk::SceneIO
+  \brief Serializes mitk::PlanarFigureComposite for mitk::SceneIO
 */
 class PlanarFigureCompositeSerializer : public BaseDataSerializer
 {
   public:
     mitkClassMacro( PlanarFigureCompositeSerializer, BaseDataSerializer );
     itkFactorylessNewMacro(Self)
     itkCloneMacro(Self)
     virtual std::string Serialize() override;
   protected:
     PlanarFigureCompositeSerializer();
     virtual ~PlanarFigureCompositeSerializer();
 };
 } // namespace
 #endif
diff --git a/Modules/DiffusionImaging/DiffusionIO/mitkTractographyForestSerializer.cpp b/Modules/DiffusionImaging/DiffusionIO/mitkTractographyForestSerializer.cpp
new file mode 100644
index 0000000000..b26c6ce347
--- /dev/null
+++ b/Modules/DiffusionImaging/DiffusionIO/mitkTractographyForestSerializer.cpp
@@ -0,0 +1,73 @@
+/*===================================================================
+
+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.
+
+===================================================================*/
+
+#include "mitkTractographyForestSerializer.h"
+#include "mitkTractographyForest.h"
+#include "mitkTractographyForestWriter.h"
+#include <mitkTractographyForest.h>
+
+#include <itksys/SystemTools.hxx>
+#include <mitkIOUtil.h>
+
+
+MITK_REGISTER_SERIALIZER(TractographyForestSerializer)
+
+
+mitk::TractographyForestSerializer::TractographyForestSerializer()
+{
+}
+
+
+mitk::TractographyForestSerializer::~TractographyForestSerializer()
+{
+}
+
+
+std::string mitk::TractographyForestSerializer::Serialize()
+{
+  const TractographyForest* fb = dynamic_cast<const TractographyForest*>( m_Data.GetPointer() );
+  if (fb == nullptr)
+  {
+    MITK_ERROR << " Object at " << (const void*) this->m_Data
+              << " is not an mitk::TractographyForest. Cannot serialize as TractographyForest.";
+    return "";
+  }
+
+  std::string filename( this->GetUniqueFilenameInWorkingDirectory() );
+  filename += "_";
+  filename += m_FilenameHint;
+  filename += ".rf";
+
+  std::string fullname(m_WorkingDirectory);
+  fullname += "/";
+  fullname += itksys::SystemTools::ConvertToOutputPath(filename.c_str());
+
+  try
+  {
+    mitk::IOUtil::Save(fb,fullname);
+  }
+  catch (std::exception& e)
+  {
+    MITK_ERROR << " Error serializing object at " << (const void*) this->m_Data
+              << " to "
+              << fullname
+              << ": "
+              << e.what();
+    return "";
+  }
+  return filename;
+}
+
diff --git a/Modules/DiffusionImaging/DiffusionIO/mitkPlanarFigureCompositeSerializer.h b/Modules/DiffusionImaging/DiffusionIO/mitkTractographyForestSerializer.h
similarity index 64%
copy from Modules/DiffusionImaging/DiffusionIO/mitkPlanarFigureCompositeSerializer.h
copy to Modules/DiffusionImaging/DiffusionIO/mitkTractographyForestSerializer.h
index 24184e2a5f..ac336e6b17 100644
--- a/Modules/DiffusionImaging/DiffusionIO/mitkPlanarFigureCompositeSerializer.h
+++ b/Modules/DiffusionImaging/DiffusionIO/mitkTractographyForestSerializer.h
@@ -1,39 +1,39 @@
 /*===================================================================
 
 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 mitkPlanarFigureCompositeSerializer_h_included
-#define mitkPlanarFigureCompositeSerializer_h_included
+#ifndef mitkTractographyForestSerializer_h_included
+#define mitkTractographyForestSerializer_h_included
 
 #include "mitkBaseDataSerializer.h"
 
 namespace mitk
 {
 /**
-  \brief Serializes mitk::Surface for mitk::SceneIO
+  \brief Serializes mitk::TractographyForest for mitk::SceneIO
 */
-class PlanarFigureCompositeSerializer : public BaseDataSerializer
+class TractographyForestSerializer : public BaseDataSerializer
 {
   public:
-    mitkClassMacro( PlanarFigureCompositeSerializer, BaseDataSerializer );
+    mitkClassMacro( TractographyForestSerializer, BaseDataSerializer );
     itkFactorylessNewMacro(Self)
     itkCloneMacro(Self)
     virtual std::string Serialize() override;
   protected:
-    PlanarFigureCompositeSerializer();
-    virtual ~PlanarFigureCompositeSerializer();
+    TractographyForestSerializer();
+    virtual ~TractographyForestSerializer();
 };
 } // namespace
 #endif
diff --git a/Modules/DiffusionImaging/FiberTracking/Algorithms/TrackingHandlers/mitkTrackingHandlerOdf.cpp b/Modules/DiffusionImaging/FiberTracking/Algorithms/TrackingHandlers/mitkTrackingHandlerOdf.cpp
index c8361cceab..5e81179c97 100644
--- a/Modules/DiffusionImaging/FiberTracking/Algorithms/TrackingHandlers/mitkTrackingHandlerOdf.cpp
+++ b/Modules/DiffusionImaging/FiberTracking/Algorithms/TrackingHandlers/mitkTrackingHandlerOdf.cpp
@@ -1,274 +1,273 @@
 /*===================================================================
 
 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.
 
 ===================================================================*/
 
 #include "mitkTrackingHandlerOdf.h"
 #include <itkDiffusionQballGeneralizedFaImageFilter.h>
 #include <itkImageRegionIterator.h>
 #include <itkPointShell.h>
 #include <omp.h>
 #include <math.h>
 
 namespace mitk
 {
 
 TrackingHandlerOdf::TrackingHandlerOdf()
   : m_GfaThreshold(0.2)
   , m_OdfThreshold(0.1)
   , m_SharpenOdfs(false)
   , m_NumProbSamples(1)
 {
 
 }
 
 TrackingHandlerOdf::~TrackingHandlerOdf()
 {
 }
 
 void TrackingHandlerOdf::InitForTracking()
 {
   MITK_INFO << "Initializing ODF tracker.";
 
   if (m_NeedsDataInit)
   {
     m_OdfHemisphereIndices.clear();
     itk::OrientationDistributionFunction< float, QBALL_ODFSIZE > odf;
     vnl_vector_fixed<double,3> ref; ref.fill(0); ref[0]=1;
 
     for (int i=0; i<QBALL_ODFSIZE; i++)
       if (dot_product(ref, odf.GetDirection(i))>0)
         m_OdfHemisphereIndices.push_back(i);
     m_OdfFloatDirs.set_size(m_OdfHemisphereIndices.size(), 3);
 
     for (unsigned int i=0; i<m_OdfHemisphereIndices.size(); i++)
     {
       m_OdfFloatDirs[i][0] = odf.GetDirection(m_OdfHemisphereIndices[i])[0];
       m_OdfFloatDirs[i][1] = odf.GetDirection(m_OdfHemisphereIndices[i])[1];
       m_OdfFloatDirs[i][2] = odf.GetDirection(m_OdfHemisphereIndices[i])[2];
     }
 
     if (m_GfaImage.IsNull())
     {
       MITK_INFO << "Calculating GFA image.";
       typedef itk::DiffusionQballGeneralizedFaImageFilter<float,float,QBALL_ODFSIZE> GfaFilterType;
       GfaFilterType::Pointer gfaFilter = GfaFilterType::New();
       gfaFilter->SetInput(m_OdfImage);
       gfaFilter->SetComputationMethod(GfaFilterType::GFA_STANDARD);
       gfaFilter->Update();
       m_GfaImage = gfaFilter->GetOutput();
     }
 
     m_NeedsDataInit = false;
   }
 
   std::cout << "TrackingHandlerOdf - GFA threshold: " << m_GfaThreshold << std::endl;
   std::cout << "TrackingHandlerOdf - ODF threshold: " << m_OdfThreshold << std::endl;
   if (m_SharpenOdfs)
     std::cout << "TrackingHandlerOdf - Sharpening ODfs" << std::endl;
 }
 
 int TrackingHandlerOdf::SampleOdf(vnl_vector< float >& probs, vnl_vector< float >& angles)
 {
   boost::random::discrete_distribution<int, float> dist(probs.begin(), probs.end());
   int sampled_idx = 0;
   int max_sample_idx = -1;
   float max_prob = 0;
   int trials = 0;
 
   for (int i=0; i<m_NumProbSamples; i++)  // we sample m_NumProbSamples times and retain the sample with maximum probabilty
   {
     trials++;
 #pragma omp critical
     {
       boost::random::variate_generator<boost::random::mt19937&, boost::random::discrete_distribution<int,float>> sampler(m_Rng, dist);
       sampled_idx = sampler();
     }
     if (probs[sampled_idx]>max_prob && probs[sampled_idx]>m_OdfThreshold && fabs(angles[sampled_idx])>=m_AngularThreshold)
     {
       max_prob = probs[sampled_idx];
       max_sample_idx = sampled_idx;
     }
     else if ( (probs[sampled_idx]<=m_OdfThreshold || fabs(angles[sampled_idx])<m_AngularThreshold) && trials<50) // we allow 50 trials to exceed the ODF threshold
       i--;
   }
 
   return max_sample_idx;
 }
 
 vnl_vector_fixed<float,3> TrackingHandlerOdf::ProposeDirection(const itk::Point<float, 3>& pos, std::deque<vnl_vector_fixed<float, 3> >& olddirs, itk::Index<3>& oldIndex)
 {
 
   vnl_vector_fixed<float,3> output_direction; output_direction.fill(0);
 
   itk::Index<3> idx;
   m_OdfImage->TransformPhysicalPointToIndex(pos, idx);
 
   if ( !m_OdfImage->GetLargestPossibleRegion().IsInside(idx) )
     return output_direction;
 
   // check GFA threshold for termination
   float gfa = GetImageValue<float>(pos, m_GfaImage, m_Interpolate);
   if (gfa<m_GfaThreshold)
     return output_direction;
 
   vnl_vector_fixed<float,3> last_dir;
   if (!olddirs.empty())
     last_dir = olddirs.back();
 
   if (!m_Interpolate && oldIndex==idx)
     return last_dir;
 
   ItkOdfImageType::PixelType odf_values = GetImageValue<float, QBALL_ODFSIZE>(pos, m_OdfImage, m_Interpolate);
   vnl_vector< float > probs; probs.set_size(m_OdfHemisphereIndices.size());
   vnl_vector< float > angles; angles.set_size(m_OdfHemisphereIndices.size()); angles.fill(1.0);
 
   // Find ODF maximum and remove <0 values
   float max_odf_val = 0;
   float min_odf_val = 999;
   int max_idx_d = -1;
   int c = 0;
   for (int i : m_OdfHemisphereIndices)
   {
     if (odf_values[i]<0)
       odf_values[i] = 0;
 
     if (odf_values[i]>max_odf_val)
     {
       max_odf_val = odf_values[i];
       max_idx_d = c;
     }
     if (odf_values[i]<min_odf_val)
       min_odf_val = odf_values[i];
 
     probs[c] = odf_values[i];
     c++;
   }
 
   if (m_SharpenOdfs)
   {
     // sharpen ODF
     probs -= min_odf_val;
     probs /= (max_odf_val-min_odf_val);
     for (unsigned int i=0; i<probs.size(); i++)
       probs[i] = pow(probs[i], 4);
     float odf_sum = probs.sum();
     if (odf_sum>0)
     {
       probs /= odf_sum;
       max_odf_val /= odf_sum;
     }
   }
 
   // no previous direction
   if (max_odf_val>m_OdfThreshold && (olddirs.empty() || last_dir.magnitude()<=0.5))
   {
     if (m_Mode==MODE::DETERMINISTIC)  // return maximum peak
     {
       output_direction = m_OdfFloatDirs.get_row(max_idx_d);
       return output_direction * max_odf_val;
     }
     else if (m_Mode==MODE::PROBABILISTIC) // sample from complete ODF
     {
       int max_sample_idx = SampleOdf(probs, angles);
       if (max_sample_idx>=0)
         output_direction = m_OdfFloatDirs.get_row(max_sample_idx) * probs[max_sample_idx];
       return output_direction;
     }
   }
   else if (max_odf_val<=m_OdfThreshold) // return (0,0,0)
   {
     return output_direction;
   }
 
   // correct previous direction
   if (m_FlipX)
     last_dir[0] *= -1;
   if (m_FlipY)
     last_dir[1] *= -1;
   if (m_FlipZ)
     last_dir[2] *= -1;
 
   // calculate angles between previous direction and ODF directions
   angles = m_OdfFloatDirs*last_dir;
 
   float probs_sum = 0;
   float max_prob = 0;
   for (unsigned int i=0; i<m_OdfHemisphereIndices.size(); i++)
   {
     float odf_val = probs[i];
     float angle = angles[i];
     float abs_angle = fabs(angle);
 
     odf_val *= abs_angle; // weight probabilities according to deviation from last direction
     if (m_Mode==MODE::DETERMINISTIC && odf_val>max_prob && odf_val>m_OdfThreshold)
     {
       // use maximum peak of the ODF weighted with the directional prior
       max_prob = odf_val;
       vnl_vector_fixed<float,3> d = m_OdfFloatDirs.get_row(i);
       if (angle<0)
         d *= -1;
       output_direction = odf_val*d;
     }
     else if (m_Mode==MODE::PROBABILISTIC)
     {
       // update ODF probabilties with the ODF values pow(abs_angle, m_DirPriorPower)
       probs[i] = odf_val;
       probs_sum += probs[i];
     }
   }
 
   // do probabilistic sampling
   if (m_Mode==MODE::PROBABILISTIC && probs_sum>0.0001)
   {
-//    probs /= probs_sum;
     int max_sample_idx = SampleOdf(probs, angles);
     if (max_sample_idx>=0)
     {
       output_direction = m_OdfFloatDirs.get_row(max_sample_idx);
       if (angles[max_sample_idx]<0)
         output_direction *= -1;
       output_direction *= probs[max_sample_idx];
     }
   }
 
   // check hard angular threshold
   float mag = output_direction.magnitude();
   if (mag>=0.0001)
   {
     output_direction.normalize();
     float a = dot_product(output_direction, last_dir);
     if (a<m_AngularThreshold)
       output_direction.fill(0);
   }
   else
     output_direction.fill(0);
 
   if (m_FlipX)
     output_direction[0] *= -1;
   if (m_FlipY)
     output_direction[1] *= -1;
   if (m_FlipZ)
     output_direction[2] *= -1;
 
   return output_direction;
 }
 
 void TrackingHandlerOdf::SetNumProbSamples(int NumProbSamples)
 {
   m_NumProbSamples = NumProbSamples;
 }
 
 }
 
diff --git a/Modules/DiffusionImaging/FiberTracking/Algorithms/TrackingHandlers/mitkTrackingHandlerRandomForest.cpp b/Modules/DiffusionImaging/FiberTracking/Algorithms/TrackingHandlers/mitkTrackingHandlerRandomForest.cpp
index 482bd32d07..ad9edb8b68 100644
--- a/Modules/DiffusionImaging/FiberTracking/Algorithms/TrackingHandlers/mitkTrackingHandlerRandomForest.cpp
+++ b/Modules/DiffusionImaging/FiberTracking/Algorithms/TrackingHandlers/mitkTrackingHandlerRandomForest.cpp
@@ -1,928 +1,932 @@
 /*===================================================================
 
 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_cpp
 #define _TrackingForestHandler_cpp
 
 #include "mitkTrackingHandlerRandomForest.h"
 #include <itkTractDensityImageFilter.h>
 #include <mitkDiffusionPropertyHelper.h>
 
 namespace mitk
 {
 template< int ShOrder, int NumberOfSignalFeatures >
 TrackingHandlerRandomForest< ShOrder, NumberOfSignalFeatures >::TrackingHandlerRandomForest()
   : m_WmSampleDistance(-1)
   , m_NumTrees(30)
   , m_MaxTreeDepth(25)
   , m_SampleFraction(1.0)
   , m_NumberOfSamples(0)
   , m_GmSamplesPerVoxel(-1)
   , m_NumPreviousDirections(1)
   , m_BidirectionalFiberSampling(false)
   , m_ZeroDirWmFeatures(true)
   , m_MaxNumWmSamples(-1)
 {
   vnl_vector_fixed<float,3> ref; ref.fill(0); ref[0]=1;
 
   itk::OrientationDistributionFunction< float, 200 > odf;
   m_DirectionContainer.clear();
   for (unsigned int i = 0; i<odf.GetNumberOfComponents(); i++)
   {
     vnl_vector_fixed<float,3> odf_dir;
     odf_dir[0] = odf.GetDirection(i)[0];
     odf_dir[1] = odf.GetDirection(i)[1];
     odf_dir[2] = odf.GetDirection(i)[2];
     if (dot_product(ref, odf_dir)>0)            // only used directions on one hemisphere
       m_DirectionContainer.push_back(odf_dir);  // store indices for later mapping the classifier output to the actual direction
   }
 
 
   m_OdfFloatDirs.set_size(m_DirectionContainer.size(), 3);
 
   for (unsigned int i=0; i<m_DirectionContainer.size(); i++)
   {
       m_OdfFloatDirs[i][0] = m_DirectionContainer.at(i)[0];
       m_OdfFloatDirs[i][1] = m_DirectionContainer.at(i)[1];
       m_OdfFloatDirs[i][2] = m_DirectionContainer.at(i)[2];
   }
 }
 
 template< int ShOrder, int NumberOfSignalFeatures >
 TrackingHandlerRandomForest< ShOrder, NumberOfSignalFeatures >::~TrackingHandlerRandomForest()
 {
 }
 
 template< int ShOrder, int NumberOfSignalFeatures >
 typename TrackingHandlerRandomForest< ShOrder, NumberOfSignalFeatures >::DwiFeatureImageType::PixelType TrackingHandlerRandomForest< ShOrder, NumberOfSignalFeatures >::GetDwiFeaturesAtPosition(itk::Point<float, 3> itkP, typename DwiFeatureImageType::Pointer image, bool interpolate)
 {
   // transform physical point to index coordinates
   itk::Index<3> idx;
   itk::ContinuousIndex< float, 3> cIdx;
   image->TransformPhysicalPointToIndex(itkP, idx);
   image->TransformPhysicalPointToContinuousIndex(itkP, cIdx);
 
   typename DwiFeatureImageType::PixelType pix; pix.Fill(0.0);
   if ( image->GetLargestPossibleRegion().IsInside(idx) )
   {
     pix = image->GetPixel(idx);
     if (!interpolate)
       return pix;
   }
   else
     return pix;
 
   float frac_x = cIdx[0] - idx[0];
   float frac_y = cIdx[1] - idx[1];
   float frac_z = cIdx[2] - idx[2];
   if (frac_x<0)
   {
     idx[0] -= 1;
     frac_x += 1;
   }
   if (frac_y<0)
   {
     idx[1] -= 1;
     frac_y += 1;
   }
   if (frac_z<0)
   {
     idx[2] -= 1;
     frac_z += 1;
   }
   frac_x = 1-frac_x;
   frac_y = 1-frac_y;
   frac_z = 1-frac_z;
 
   // int coordinates inside image?
   if (idx[0] >= 0 && idx[0] < static_cast<itk::IndexValueType>(image->GetLargestPossibleRegion().GetSize(0) - 1) &&
       idx[1] >= 0 && idx[1] < static_cast<itk::IndexValueType>(image->GetLargestPossibleRegion().GetSize(1) - 1) &&
       idx[2] >= 0 && idx[2] < static_cast<itk::IndexValueType>(image->GetLargestPossibleRegion().GetSize(2) - 1))
   {
     // trilinear interpolation
     vnl_vector_fixed<float, 8> interpWeights;
     interpWeights[0] = (  frac_x)*(  frac_y)*(  frac_z);
     interpWeights[1] = (1-frac_x)*(  frac_y)*(  frac_z);
     interpWeights[2] = (  frac_x)*(1-frac_y)*(  frac_z);
     interpWeights[3] = (  frac_x)*(  frac_y)*(1-frac_z);
     interpWeights[4] = (1-frac_x)*(1-frac_y)*(  frac_z);
     interpWeights[5] = (  frac_x)*(1-frac_y)*(1-frac_z);
     interpWeights[6] = (1-frac_x)*(  frac_y)*(1-frac_z);
     interpWeights[7] = (1-frac_x)*(1-frac_y)*(1-frac_z);
 
     pix = image->GetPixel(idx) * interpWeights[0];
     typename DwiFeatureImageType::IndexType tmpIdx = idx; tmpIdx[0]++;
     pix +=  image->GetPixel(tmpIdx) * interpWeights[1];
     tmpIdx = idx; tmpIdx[1]++;
     pix +=  image->GetPixel(tmpIdx) * interpWeights[2];
     tmpIdx = idx; tmpIdx[2]++;
     pix +=  image->GetPixel(tmpIdx) * interpWeights[3];
     tmpIdx = idx; tmpIdx[0]++; tmpIdx[1]++;
     pix +=  image->GetPixel(tmpIdx) * interpWeights[4];
     tmpIdx = idx; tmpIdx[1]++; tmpIdx[2]++;
     pix +=  image->GetPixel(tmpIdx) * interpWeights[5];
     tmpIdx = idx; tmpIdx[2]++; tmpIdx[0]++;
     pix +=  image->GetPixel(tmpIdx) * interpWeights[6];
     tmpIdx = idx; tmpIdx[0]++; tmpIdx[1]++; tmpIdx[2]++;
     pix +=  image->GetPixel(tmpIdx) * interpWeights[7];
   }
 
   return pix;
 }
 
 template< int ShOrder, int NumberOfSignalFeatures >
 void TrackingHandlerRandomForest< ShOrder, NumberOfSignalFeatures >::InputDataValidForTracking()
 {
   if (m_InputDwis.empty())
     mitkThrow() << "No diffusion-weighted images set!";
   if (!IsForestValid())
     mitkThrow() << "No or invalid random forest detected!";
 }
 
 template< int ShOrder, int NumberOfSignalFeatures>
 template<typename T>
 typename std::enable_if< NumberOfSignalFeatures <=99, T >::type TrackingHandlerRandomForest< ShOrder, NumberOfSignalFeatures >::InitDwiImageFeatures(mitk::Image::Pointer mitk_dwi)
 {
   MITK_INFO << "Calculating spherical harmonics features";
   typedef itk::AnalyticalDiffusionQballReconstructionImageFilter<short,short,float,ShOrder, 2*NumberOfSignalFeatures> InterpolationFilterType;
 
   typename InterpolationFilterType::Pointer filter = InterpolationFilterType::New();
   filter->SetBValue(mitk::DiffusionPropertyHelper::GetReferenceBValue(mitk_dwi));
   filter->SetGradientImage( mitk::DiffusionPropertyHelper::GetOriginalGradientContainer(mitk_dwi), mitk::DiffusionPropertyHelper::GetItkVectorImage(mitk_dwi) );
   filter->SetLambda(0.006);
   filter->SetNormalizationMethod(InterpolationFilterType::QBAR_RAW_SIGNAL);
   filter->Update();
 
   m_DwiFeatureImages.push_back(filter->GetCoefficientImage());
   return true;
 }
 
 template< int ShOrder, int NumberOfSignalFeatures>
 template<typename T>
 typename std::enable_if< NumberOfSignalFeatures >=100, T >::type TrackingHandlerRandomForest< ShOrder, NumberOfSignalFeatures >::InitDwiImageFeatures(mitk::Image::Pointer mitk_dwi)
 {
   MITK_INFO << "Interpolating raw dwi signal features";
   typedef itk::AnalyticalDiffusionQballReconstructionImageFilter<short,short,float,ShOrder, 2*NumberOfSignalFeatures> InterpolationFilterType;
 
   typename InterpolationFilterType::Pointer filter = InterpolationFilterType::New();
   filter->SetBValue(mitk::DiffusionPropertyHelper::GetReferenceBValue(mitk_dwi));
   filter->SetGradientImage( mitk::DiffusionPropertyHelper::GetOriginalGradientContainer(mitk_dwi), mitk::DiffusionPropertyHelper::GetItkVectorImage(mitk_dwi) );
   filter->SetLambda(0.006);
   filter->SetNormalizationMethod(InterpolationFilterType::QBAR_RAW_SIGNAL);
   filter->Update();
 
   typename DwiFeatureImageType::Pointer dwiFeatureImage = DwiFeatureImageType::New();
   dwiFeatureImage->SetSpacing(filter->GetOutput()->GetSpacing());
   dwiFeatureImage->SetOrigin(filter->GetOutput()->GetOrigin());
   dwiFeatureImage->SetDirection(filter->GetOutput()->GetDirection());
   dwiFeatureImage->SetLargestPossibleRegion(filter->GetOutput()->GetLargestPossibleRegion());
   dwiFeatureImage->SetBufferedRegion(filter->GetOutput()->GetLargestPossibleRegion());
   dwiFeatureImage->SetRequestedRegion(filter->GetOutput()->GetLargestPossibleRegion());
   dwiFeatureImage->Allocate();
 
   // get signal values and store them in the feature image
   vnl_vector_fixed<double,3> ref; ref.fill(0); ref[0]=1;
   itk::OrientationDistributionFunction< float, 2*NumberOfSignalFeatures > odf;
   itk::ImageRegionIterator< typename InterpolationFilterType::OutputImageType > it(filter->GetOutput(), filter->GetOutput()->GetLargestPossibleRegion());
   while(!it.IsAtEnd())
   {
     typename DwiFeatureImageType::PixelType pix;
     int f = 0;
     for (unsigned int i = 0; i<odf.GetNumberOfComponents(); i++)
     {
       if (dot_product(ref, odf.GetDirection(i))>0)            // only used directions on one hemisphere
       {
         pix[f] = it.Get()[i];
         f++;
       }
     }
     dwiFeatureImage->SetPixel(it.GetIndex(), pix);
     ++it;
   }
   m_DwiFeatureImages.push_back(dwiFeatureImage);
   return true;
 }
 
 template< int ShOrder, int NumberOfSignalFeatures >
 void TrackingHandlerRandomForest< ShOrder, NumberOfSignalFeatures >::InitForTracking()
 {
   MITK_INFO << "Initializing random forest tracker.";
   if (m_NeedsDataInit)
   {
     InputDataValidForTracking();
     m_DwiFeatureImages.clear();
     InitDwiImageFeatures<>(m_InputDwis.at(0));
     m_NeedsDataInit = false;
   }
 }
 
 template< int ShOrder, int NumberOfSignalFeatures >
 vnl_vector_fixed<float,3> TrackingHandlerRandomForest< ShOrder, NumberOfSignalFeatures >::ProposeDirection(const itk::Point<float, 3>& pos, std::deque<vnl_vector_fixed<float, 3> >& olddirs, itk::Index<3>& oldIndex)
 {
 
   vnl_vector_fixed<float,3> output_direction; output_direction.fill(0);
 
   itk::Index<3> idx;
   m_DwiFeatureImages.at(0)->TransformPhysicalPointToIndex(pos, idx);
 
   bool check_last_dir = false;
   vnl_vector_fixed<float,3> last_dir;
   if (!olddirs.empty())
   {
     last_dir = olddirs.back();
     if (last_dir.magnitude()>0.5)
       check_last_dir = true;
   }
 
   if (!m_Interpolate && oldIndex==idx)
     return last_dir;
 
   // store feature pixel values in a vigra data type
   vigra::MultiArray<2, float> featureData = vigra::MultiArray<2, float>( vigra::Shape2(1,m_Forest->GetNumFeatures()) );
   typename DwiFeatureImageType::PixelType dwiFeaturePixel = GetDwiFeaturesAtPosition(pos, m_DwiFeatureImages.at(0), m_Interpolate);
   for (unsigned int f=0; f<NumberOfSignalFeatures; f++)
     featureData(0,f) = dwiFeaturePixel[f];
 
   vnl_matrix_fixed<double,3,3> direction_matrix = m_DwiFeatureImages.at(0)->GetDirection().GetVnlMatrix();
   vnl_matrix_fixed<double,3,3> inverse_direction_matrix = m_DwiFeatureImages.at(0)->GetInverseDirection().GetVnlMatrix();
 
   // append normalized previous direction(s) to feature vector
   int i = 0;
   vnl_vector_fixed<double,3> ref; ref.fill(0); ref[0]=1;
 
   for (auto d : olddirs)
   {
     vnl_vector_fixed<double,3> tempD;
     tempD[0] = d[0]; tempD[1] = d[1]; tempD[2] = d[2];
 
     if (m_FlipX)
         tempD[0] *= -1;
     if (m_FlipY)
         tempD[1] *= -1;
     if (m_FlipZ)
         tempD[2] *= -1;
 
     tempD = inverse_direction_matrix * tempD;
     last_dir[0] = tempD[0];
     last_dir[1] = tempD[1];
     last_dir[2] = tempD[2];
 
     int c = 0;
     for (int f=NumberOfSignalFeatures+3*i; f<NumberOfSignalFeatures+3*(i+1); f++)
     {
       if (dot_product(ref, tempD)<0)
         featureData(0,f) = -tempD[c];
       else
         featureData(0,f) = tempD[c];
       c++;
     }
     i++;
   }
 
   // additional feature images
   if (m_AdditionalFeatureImages.size()>0)
   {
     int c = 0;
     for (auto img : m_AdditionalFeatureImages.at(0))
     {
       float v = GetImageValue<float>(pos, img, false);
       featureData(0,NumberOfSignalFeatures+m_NumPreviousDirections*3+c) = v;
       c++;
     }
   }
 
   // perform classification
   vigra::MultiArray<2, float> probs(vigra::Shape2(1, m_Forest->GetNumClasses()));
   m_Forest->PredictProbabilities(featureData, probs);
 
   vnl_vector< float > angles = m_OdfFloatDirs*last_dir;
   vnl_vector< float > probs2; probs2.set_size(m_DirectionContainer.size()); probs2.fill(0.0); // used for probabilistic direction sampling
   float probs_sum = 0;
 
   float pNonFib = 0;     // probability that we left the white matter
   float w = 0;           // weight of the predicted direction
 
   for (int i=0; i<m_Forest->GetNumClasses(); i++)   // for each class (number of possible directions + out-of-wm class)
   {
     if (probs(0,i)>0)   // if probability of respective class is 0, do nothing
     {
       // get label of class (does not correspond to the loop variable i)
       unsigned int classLabel = m_Forest->IndexToClassLabel(i);
 
       if (classLabel<m_DirectionContainer.size())   // does class label correspond to a direction or to the out-of-wm class?
       {
+        float angle = angles[classLabel];
+        float abs_angle = fabs(angle);
+
         if (m_Mode==MODE::PROBABILISTIC)
         {
-          if (!check_last_dir || fabs(angles[classLabel])>=m_AngularThreshold)
-            probs2[classLabel] = probs(0,i);
+          probs2[classLabel] = probs(0,i);
+          if (check_last_dir)
+            probs2[classLabel] *= abs_angle;
           probs_sum += probs2[classLabel];
         }
         else if (m_Mode==MODE::DETERMINISTIC)
         {
           vnl_vector_fixed<float,3> d = m_DirectionContainer.at(classLabel);  // get direction vector assiciated with the respective direction index
           if (check_last_dir)   // do we have a previous streamline direction or did we just start?
           {
-            // TODO: check if hard curvature threshold is necessary.
-            // alternatively try square of dot product as weight.
-            // TODO: check if additional weighting with dot product as directional prior is necessary. are there alternatives on the classification level?
-
-            float dot = angles[classLabel];    // claculate angle between the candidate direction vector and the previous streamline direction
-            if (fabs(dot)>=m_AngularThreshold)         // is angle between the directions smaller than our hard threshold?
+            if (abs_angle>=m_AngularThreshold)         // is angle between the directions smaller than our hard threshold?
             {
-              if (dot<0)                          // make sure we don't walk backwards
+              if (angle<0)                          // make sure we don't walk backwards
                 d *= -1;
-              float w_i = probs(0,i)*fabs(dot);
+              float w_i = probs(0,i)*abs_angle;
               output_direction += w_i*d; // weight contribution to output direction with its probability and the angular deviation from the previous direction
               w += w_i;           // increase output weight of the final direction
             }
           }
           else
           {
             output_direction += probs(0,i)*d;
             w += probs(0,i);
           }
         }
       }
       else
         pNonFib += probs(0,i);  // probability that we are not in the white matter anymore
     }
   }
 
 
   if (m_Mode==MODE::PROBABILISTIC && pNonFib<0.5)
   {
-    probs2 /= probs_sum;
     boost::random::discrete_distribution<int, float> dist(probs2.begin(), probs2.end());
-
     int sampled_idx = 0;
-#pragma omp critical
+
+    for (int i=0; i<50; i++)  // we allow 50 trials to exceed m_AngularThreshold
     {
+  #pragma omp critical
+      {
         boost::random::variate_generator<boost::random::mt19937&, boost::random::discrete_distribution<int,float>> sampler(m_Rng, dist);
         sampled_idx = sampler();
+      }
+
+      if ( probs2[sampled_idx]>0.1 && (!check_last_dir || (check_last_dir && fabs(angles[sampled_idx])>=m_AngularThreshold)) )
+        break;
     }
 
     output_direction = m_DirectionContainer.at(sampled_idx);
     w = probs2[sampled_idx];
     if (check_last_dir && angles[sampled_idx]<0)                          // make sure we don't walk backwards
         output_direction *= -1;
   }
 
   // if we did not find a suitable direction, make sure that we return (0,0,0)
   if (pNonFib>w && w>0)
     output_direction.fill(0.0);
   else
   {
     vnl_vector_fixed<double,3> tempD;
     tempD[0] = output_direction[0]; tempD[1] = output_direction[1]; tempD[2] = output_direction[2];
     tempD = direction_matrix * tempD;
     output_direction[0] = tempD[0];
     output_direction[1] = tempD[1];
     output_direction[2] = tempD[2];
 
     if (m_FlipX)
       output_direction[0] *= -1;
     if (m_FlipY)
       output_direction[1] *= -1;
     if (m_FlipZ)
       output_direction[2] *= -1;
   }
 
   return output_direction * w;
 }
 
 template< int ShOrder, int NumberOfSignalFeatures >
 void TrackingHandlerRandomForest< ShOrder, NumberOfSignalFeatures >::StartTraining()
 {
   m_StartTime = std::chrono::system_clock::now();
   InputDataValidForTraining();
   InitForTraining();
   CalculateTrainingSamples();
 
   MITK_INFO << "Maximum tree depths: " << m_MaxTreeDepth;
   MITK_INFO << "Sample fraction per tree: " << m_SampleFraction;
   MITK_INFO << "Number of trees: " << m_NumTrees;
 
   DefaultSplitType splitter;
   splitter.UsePointBasedWeights(true);
   splitter.SetWeights(m_Weights);
   splitter.UseRandomSplit(false);
   splitter.SetPrecision(mitk::eps);
   splitter.SetMaximumTreeDepth(m_MaxTreeDepth);
 
   std::vector< std::shared_ptr< vigra::RandomForest<int> > > trees;
   int count = 0;
 #pragma omp parallel for
   for (int i = 0; i < m_NumTrees; ++i)
   {
     std::shared_ptr< vigra::RandomForest<int> > lrf = std::make_shared< vigra::RandomForest<int> >();
     lrf->set_options().use_stratification(vigra::RF_NONE); // How the data should be made equal
     lrf->set_options().sample_with_replacement(true); // if sampled with replacement or not
     lrf->set_options().samples_per_tree(m_SampleFraction); // Fraction of samples that are used to train a tree
     lrf->set_options().tree_count(1); // Number of trees that are calculated;
     lrf->set_options().min_split_node_size(5); // Minimum number of datapoints that must be in a node
     lrf->ext_param_.max_tree_depth = m_MaxTreeDepth;
 
     lrf->learn(m_FeatureData, m_LabelData,vigra::rf::visitors::VisitorBase(),splitter);
 #pragma omp critical
     {
       count++;
       MITK_INFO << "Tree " << count << " finished training.";
       trees.push_back(lrf);
     }
   }
 
   for (int i = 1; i < m_NumTrees; ++i)
     trees.at(0)->trees_.push_back(trees.at(i)->trees_[0]);
 
   std::shared_ptr< vigra::RandomForest<int> > forest = trees.at(0);
   forest->options_.tree_count_ = m_NumTrees;
   m_Forest = mitk::TractographyForest::New(forest);
   MITK_INFO << "Training finsihed";
 
   m_EndTime = std::chrono::system_clock::now();
   std::chrono::hours   hh = std::chrono::duration_cast<std::chrono::hours>(m_EndTime - m_StartTime);
   std::chrono::minutes mm = std::chrono::duration_cast<std::chrono::minutes>(m_EndTime - m_StartTime);
   mm %= 60;
   MITK_INFO << "Training took " << hh.count() << "h and " << mm.count() << "m";
 }
 
 template< int ShOrder, int NumberOfSignalFeatures >
 void TrackingHandlerRandomForest< ShOrder, NumberOfSignalFeatures >::InputDataValidForTraining()
 {
   if (m_InputDwis.empty())
     mitkThrow() << "No diffusion-weighted images set!";
   if (m_Tractograms.empty())
     mitkThrow() << "No tractograms set!";
   if (m_InputDwis.size()!=m_Tractograms.size())
     mitkThrow() << "Unequal number of diffusion-weighted images and tractograms detected!";
 }
 
 template< int ShOrder, int NumberOfSignalFeatures >
 bool TrackingHandlerRandomForest< ShOrder, NumberOfSignalFeatures >::IsForestValid()
 {
   int additional_features = 0;
   if (m_AdditionalFeatureImages.size()>0)
     additional_features = m_AdditionalFeatureImages.at(0).size();
 
   if (!m_Forest)
     MITK_INFO << "No forest available!";
   else
   {
     if (m_Forest->GetNumTrees() <= 0)
       MITK_ERROR << "Forest contains no trees!";
     if ( m_Forest->GetNumFeatures() != static_cast<int>(NumberOfSignalFeatures+3*m_NumPreviousDirections+additional_features) )
       MITK_ERROR << "Wrong number of features in forest: got " << m_Forest->GetNumFeatures() << ", expected " << (NumberOfSignalFeatures+3*m_NumPreviousDirections+additional_features);
   }
 
   if(m_Forest && m_Forest->GetNumTrees()>0 && m_Forest->GetNumFeatures() == static_cast<int>(NumberOfSignalFeatures+3*m_NumPreviousDirections+additional_features))
     return true;
 
   return false;
 }
 
 template< int ShOrder, int NumberOfSignalFeatures >
 void TrackingHandlerRandomForest< ShOrder, NumberOfSignalFeatures >::InitForTraining()
 {
   MITK_INFO << "Spherical signal interpolation and sampling ...";
   for (unsigned int i=0; i<m_InputDwis.size(); i++)
   {
     InitDwiImageFeatures<>(m_InputDwis.at(i));
 
     if (i>=m_AdditionalFeatureImages.size())
     {
       m_AdditionalFeatureImages.push_back(std::vector< ItkFloatImgType::Pointer >());
     }
 
     if (i>=m_FiberVolumeModImages.size())
     {
       ItkFloatImgType::Pointer img = ItkFloatImgType::New();
       img->SetSpacing( m_DwiFeatureImages.at(i)->GetSpacing() );
       img->SetOrigin( m_DwiFeatureImages.at(i)->GetOrigin() );
       img->SetDirection( m_DwiFeatureImages.at(i)->GetDirection() );
       img->SetLargestPossibleRegion( m_DwiFeatureImages.at(i)->GetLargestPossibleRegion() );
       img->SetBufferedRegion( m_DwiFeatureImages.at(i)->GetLargestPossibleRegion() );
       img->SetRequestedRegion( m_DwiFeatureImages.at(i)->GetLargestPossibleRegion() );
       img->Allocate();
       img->FillBuffer(1);
       m_FiberVolumeModImages.push_back(img);
     }
 
     if (m_FiberVolumeModImages.at(i)==nullptr)
     {
       m_FiberVolumeModImages.at(i) = ItkFloatImgType::New();
       m_FiberVolumeModImages.at(i)->SetSpacing( m_DwiFeatureImages.at(i)->GetSpacing() );
       m_FiberVolumeModImages.at(i)->SetOrigin( m_DwiFeatureImages.at(i)->GetOrigin() );
       m_FiberVolumeModImages.at(i)->SetDirection( m_DwiFeatureImages.at(i)->GetDirection() );
       m_FiberVolumeModImages.at(i)->SetLargestPossibleRegion( m_DwiFeatureImages.at(i)->GetLargestPossibleRegion() );
       m_FiberVolumeModImages.at(i)->SetBufferedRegion( m_DwiFeatureImages.at(i)->GetLargestPossibleRegion() );
       m_FiberVolumeModImages.at(i)->SetRequestedRegion( m_DwiFeatureImages.at(i)->GetLargestPossibleRegion() );
       m_FiberVolumeModImages.at(i)->Allocate();
       m_FiberVolumeModImages.at(i)->FillBuffer(1);
     }
 
     if (i>=m_MaskImages.size())
     {
       ItkUcharImgType::Pointer newMask = ItkUcharImgType::New();
       newMask->SetSpacing( m_DwiFeatureImages.at(i)->GetSpacing() );
       newMask->SetOrigin( m_DwiFeatureImages.at(i)->GetOrigin() );
       newMask->SetDirection( m_DwiFeatureImages.at(i)->GetDirection() );
       newMask->SetLargestPossibleRegion( m_DwiFeatureImages.at(i)->GetLargestPossibleRegion() );
       newMask->SetBufferedRegion( m_DwiFeatureImages.at(i)->GetLargestPossibleRegion() );
       newMask->SetRequestedRegion( m_DwiFeatureImages.at(i)->GetLargestPossibleRegion() );
       newMask->Allocate();
       newMask->FillBuffer(1);
       m_MaskImages.push_back(newMask);
     }
 
     if (m_MaskImages.at(i)==nullptr)
     {
       m_MaskImages.at(i) = ItkUcharImgType::New();
       m_MaskImages.at(i)->SetSpacing( m_DwiFeatureImages.at(i)->GetSpacing() );
       m_MaskImages.at(i)->SetOrigin( m_DwiFeatureImages.at(i)->GetOrigin() );
       m_MaskImages.at(i)->SetDirection( m_DwiFeatureImages.at(i)->GetDirection() );
       m_MaskImages.at(i)->SetLargestPossibleRegion( m_DwiFeatureImages.at(i)->GetLargestPossibleRegion() );
       m_MaskImages.at(i)->SetBufferedRegion( m_DwiFeatureImages.at(i)->GetLargestPossibleRegion() );
       m_MaskImages.at(i)->SetRequestedRegion( m_DwiFeatureImages.at(i)->GetLargestPossibleRegion() );
       m_MaskImages.at(i)->Allocate();
       m_MaskImages.at(i)->FillBuffer(1);
     }
   }
 
   MITK_INFO << "Resampling fibers and calculating number of samples ...";
   m_NumberOfSamples = 0;
   m_SampleUsage.clear();
   for (unsigned int t=0; t<m_Tractograms.size(); t++)
   {
     ItkUcharImgType::Pointer mask = m_MaskImages.at(t);
     ItkUcharImgType::Pointer wmmask;
     if (t<m_WhiteMatterImages.size() && m_WhiteMatterImages.at(t)!=nullptr)
       wmmask = m_WhiteMatterImages.at(t);
     else
     {
       MITK_INFO << "No white-matter mask found. Using fiber envelope.";
       itk::TractDensityImageFilter< ItkUcharImgType >::Pointer env = itk::TractDensityImageFilter< ItkUcharImgType >::New();
       env->SetFiberBundle(m_Tractograms.at(t));
       env->SetInputImage(mask);
       env->SetBinaryOutput(true);
       env->SetUseImageGeometry(true);
       env->Update();
       wmmask = env->GetOutput();
       if (t>=m_WhiteMatterImages.size())
         m_WhiteMatterImages.push_back(wmmask);
       else
         m_WhiteMatterImages.at(t) = wmmask;
     }
 
     // Calculate white-matter samples
     if (m_WmSampleDistance<0)
     {
       typename DwiFeatureImageType::Pointer image = m_DwiFeatureImages.at(t);
       float minSpacing = 1;
       if(image->GetSpacing()[0]<image->GetSpacing()[1] && image->GetSpacing()[0]<image->GetSpacing()[2])
         minSpacing = image->GetSpacing()[0];
       else if (image->GetSpacing()[1] < image->GetSpacing()[2])
         minSpacing = image->GetSpacing()[1];
       else
         minSpacing = image->GetSpacing()[2];
       m_WmSampleDistance = minSpacing*0.5;
     }
 
     m_Tractograms.at(t)->ResampleLinear(m_WmSampleDistance);
 
     int wmSamples = m_Tractograms.at(t)->GetNumberOfPoints()-2*m_Tractograms.at(t)->GetNumFibers();
     if (m_BidirectionalFiberSampling)
       wmSamples *= 2;
     if (m_ZeroDirWmFeatures)
       wmSamples *= (m_NumPreviousDirections+1);
 
     MITK_INFO << "White matter samples available: " << wmSamples;
 
     // upper limit for samples
     if (m_MaxNumWmSamples>0 && wmSamples>m_MaxNumWmSamples)
     {
       if ((float)m_MaxNumWmSamples/wmSamples > 0.8)
       {
         m_SampleUsage.push_back(std::vector<bool>(wmSamples, true));
         m_NumberOfSamples += wmSamples;
       }
       else
       {
         m_SampleUsage.push_back(std::vector<bool>(wmSamples, false));
         m_NumberOfSamples += m_MaxNumWmSamples;
         wmSamples = m_MaxNumWmSamples;
         MITK_INFO << "Limiting white matter samples to: " << m_MaxNumWmSamples;
 
         itk::Statistics::MersenneTwisterRandomVariateGenerator::Pointer randgen = itk::Statistics::MersenneTwisterRandomVariateGenerator::New();
         randgen->SetSeed();
         int c = 0;
         while (c<m_MaxNumWmSamples)
         {
           int idx = randgen->GetIntegerVariate(m_MaxNumWmSamples-1);
           if (m_SampleUsage[t][idx]==false)
           {
             m_SampleUsage[t][idx]=true;
             ++c;
           }
         }
       }
     }
     else
     {
       m_SampleUsage.push_back(std::vector<bool>(wmSamples, true));
       m_NumberOfSamples += wmSamples;
     }
 
     // calculate gray-matter samples
     itk::ImageRegionConstIterator<ItkUcharImgType> it(wmmask, wmmask->GetLargestPossibleRegion());
     int OUTOFWM = 0;
     while(!it.IsAtEnd())
     {
       if (it.Get()==0 && mask->GetPixel(it.GetIndex())>0)
         OUTOFWM++;
       ++it;
     }
     MITK_INFO << "Non-white matter voxels: " << OUTOFWM;
 
     if (m_GmSamplesPerVoxel>0)
     {
       m_GmSamples.push_back(m_GmSamplesPerVoxel);
       m_NumberOfSamples += m_GmSamplesPerVoxel*OUTOFWM;
     }
     else if (OUTOFWM>0)
     {
       int gm_per_voxel = 0.5+(float)wmSamples/(float)OUTOFWM;
       if (gm_per_voxel<=0)
         gm_per_voxel = 1;
       m_GmSamples.push_back(gm_per_voxel);
       m_NumberOfSamples += m_GmSamples.back()*OUTOFWM;
       MITK_INFO << "Non-white matter samples per voxel: " << m_GmSamples.back();
     }
     else
     {
       m_GmSamples.push_back(0);
     }
     MITK_INFO << "Non-white matter samples: " << m_GmSamples.back()*OUTOFWM;
   }
   MITK_INFO << "Number of samples: " << m_NumberOfSamples;
 }
 
 template< int ShOrder, int NumberOfSignalFeatures >
 void TrackingHandlerRandomForest< ShOrder, NumberOfSignalFeatures >::CalculateTrainingSamples()
 {
   vnl_vector_fixed<float,3> ref; ref.fill(0); ref[0]=1;
 
   m_FeatureData.reshape( vigra::Shape2(m_NumberOfSamples, NumberOfSignalFeatures+m_NumPreviousDirections*3+m_AdditionalFeatureImages.at(0).size()) );
   m_LabelData.reshape( vigra::Shape2(m_NumberOfSamples,1) );
   m_Weights.reshape( vigra::Shape2(m_NumberOfSamples,1) );
   MITK_INFO << "Number of features: " << m_FeatureData.shape(1);
 
   itk::Statistics::MersenneTwisterRandomVariateGenerator::Pointer m_RandGen = itk::Statistics::MersenneTwisterRandomVariateGenerator::New();
   m_RandGen->SetSeed();
   MITK_INFO <<  "Creating training data ...";
   unsigned int sampleCounter = 0;
   for (unsigned int t=0; t<m_Tractograms.size(); t++)
   {
     ItkFloatImgType::Pointer fiber_folume = m_FiberVolumeModImages.at(t);
     typename DwiFeatureImageType::Pointer image = m_DwiFeatureImages.at(t);
     ItkUcharImgType::Pointer wmMask = m_WhiteMatterImages.at(t);
     ItkUcharImgType::Pointer mask;
     if (t<m_MaskImages.size())
       mask = m_MaskImages.at(t);
 
     // non-white matter samples
     itk::ImageRegionConstIterator<ItkUcharImgType> it(wmMask, wmMask->GetLargestPossibleRegion());
     while(!it.IsAtEnd())
     {
       if (it.Get()==0 && (mask.IsNull() || (mask.IsNotNull() && mask->GetPixel(it.GetIndex())>0)))
       {
         typename DwiFeatureImageType::PixelType pix = image->GetPixel(it.GetIndex());
 
         // random directions
         for (unsigned int i=0; i<m_GmSamples.at(t); i++)
         {
           // diffusion signal features
           for (unsigned int f=0; f<NumberOfSignalFeatures; f++)
           {
             m_FeatureData(sampleCounter,f) = pix[f];
             if(m_FeatureData(sampleCounter,f)!=m_FeatureData(sampleCounter,f))
               m_FeatureData(sampleCounter,f) = 0;
           }
 
           // direction feature
           int num_zero_dirs = m_RandGen->GetIntegerVariate(m_NumPreviousDirections);  // how many directions should be zero?
           for (unsigned int i=0; i<m_NumPreviousDirections; i++)
           {
             int c=0;
             vnl_vector_fixed<float,3> probe;
             if (static_cast<int>(i)<num_zero_dirs)
               probe.fill(0.0);
             else
             {
               probe[0] = m_RandGen->GetVariate()*2-1;
               probe[1] = m_RandGen->GetVariate()*2-1;
               probe[2] = m_RandGen->GetVariate()*2-1;
               probe.normalize();
               if (dot_product(ref, probe)<0)
                 probe *= -1;
             }
             for (unsigned int f=NumberOfSignalFeatures+3*i; f<NumberOfSignalFeatures+3*(i+1); f++)
             {
               m_FeatureData(sampleCounter,f) = probe[c];
               c++;
             }
           }
 
           // additional feature images
           int add_feat_c = 0;
           for (auto img : m_AdditionalFeatureImages.at(t))
           {
             itk::Point<float, 3> itkP;
             image->TransformIndexToPhysicalPoint(it.GetIndex(), itkP);
             float v = GetImageValue<float>(itkP, img, false);
             m_FeatureData(sampleCounter,NumberOfSignalFeatures+m_NumPreviousDirections*3+add_feat_c) = v;
             add_feat_c++;
           }
 
           // label and sample weight
           m_LabelData(sampleCounter,0) = m_DirectionContainer.size();
           m_Weights(sampleCounter,0) = 1.0;
           sampleCounter++;
         }
       }
       ++it;
     }
 
     unsigned int num_gm_samples = sampleCounter;
     // white matter samples
     mitk::FiberBundle::Pointer fib = m_Tractograms.at(t);
     vtkSmartPointer< vtkPolyData > polyData = fib->GetFiberPolyData();
     vnl_vector_fixed<float,3> zero_dir; zero_dir.fill(0.0);
     for (int i=0; i<fib->GetNumFibers(); i++)
     {
       vtkCell* cell = polyData->GetCell(i);
       int numPoints = cell->GetNumberOfPoints();
       vtkPoints* points = cell->GetPoints();
       float fiber_weight = fib->GetFiberWeight(i);
 
       for (int n = 0; n <= static_cast<int>(m_NumPreviousDirections); ++n)
       {
         if (!m_ZeroDirWmFeatures)
           n = m_NumPreviousDirections;
         for (bool reverse : {false, true})
         {
           for (int j=1; j<numPoints-1; j++)
           {
             if (!m_SampleUsage[t][sampleCounter-num_gm_samples])
               continue;
 
             itk::Point<float, 3> itkP1, itkP2;
 
             int num_nonzero_dirs = m_NumPreviousDirections;
             if (!reverse)
               num_nonzero_dirs = std::min(n, j);
             else
               num_nonzero_dirs = std::min(n, numPoints-j-1);
 
             vnl_vector_fixed<float,3> dir;
             // zero directions
             for (unsigned int k=0; k<m_NumPreviousDirections-num_nonzero_dirs; k++)
             {
               dir.fill(0.0);
               int c = 0;
               for (unsigned int f=NumberOfSignalFeatures+3*k; f<NumberOfSignalFeatures+3*(k+1); f++)
               {
                 m_FeatureData(sampleCounter,f) = dir[c];
                 c++;
               }
             }
 
             // nonzero directions
             for (int k=0; k<num_nonzero_dirs; k++)
             {
               double* p = nullptr;
               int n_idx = num_nonzero_dirs-k;
               if (!reverse)
               {
                 p = points->GetPoint(j-n_idx);
                 itkP1[0] = p[0]; itkP1[1] = p[1]; itkP1[2] = p[2];
                 p = points->GetPoint(j-n_idx+1);
                 itkP2[0] = p[0]; itkP2[1] = p[1]; itkP2[2] = p[2];
               }
               else
               {
                 p = points->GetPoint(j+n_idx);
                 itkP1[0] = p[0]; itkP1[1] = p[1]; itkP1[2] = p[2];
                 p = points->GetPoint(j+n_idx-1);
                 itkP2[0] = p[0]; itkP2[1] = p[1]; itkP2[2] = p[2];
               }
 
               dir[0]=itkP1[0]-itkP2[0];
               dir[1]=itkP1[1]-itkP2[1];
               dir[2]=itkP1[2]-itkP2[2];
 
               if (dir.magnitude()<0.0001)
                 mitkThrow() << "streamline error!";
               dir.normalize();
               if (dir[0]!=dir[0] || dir[1]!=dir[1] || dir[2]!=dir[2])
                 mitkThrow() << "ERROR: NaN direction!";
 
               if (dot_product(ref, dir)<0)
                 dir *= -1;
 
               int c = 0;
               for (unsigned int f=NumberOfSignalFeatures+3*(k+m_NumPreviousDirections-num_nonzero_dirs); f<NumberOfSignalFeatures+3*(k+1+m_NumPreviousDirections-num_nonzero_dirs); f++)
               {
                 m_FeatureData(sampleCounter,f) = dir[c];
                 c++;
               }
             }
 
             // get target direction
             double* p = points->GetPoint(j);
             itkP1[0] = p[0]; itkP1[1] = p[1]; itkP1[2] = p[2];
             if (reverse)
             {
               p = points->GetPoint(j-1);
               itkP2[0] = p[0]; itkP2[1] = p[1]; itkP2[2] = p[2];
             }
             else
             {
               p = points->GetPoint(j+1);
               itkP2[0] = p[0]; itkP2[1] = p[1]; itkP2[2] = p[2];
             }
             dir[0]=itkP2[0]-itkP1[0];
             dir[1]=itkP2[1]-itkP1[1];
             dir[2]=itkP2[2]-itkP1[2];
 
             if (dir.magnitude()<0.0001)
               mitkThrow() << "streamline error!";
             dir.normalize();
             if (dir[0]!=dir[0] || dir[1]!=dir[1] || dir[2]!=dir[2])
               mitkThrow() << "ERROR: NaN direction!";
             if (dot_product(ref, dir)<0)
               dir *= -1;
 
             // image features
             float volume_mod = GetImageValue<float>(itkP1, fiber_folume, false);
 
             // diffusion signal features
             typename DwiFeatureImageType::PixelType pix = GetDwiFeaturesAtPosition(itkP1, image, m_Interpolate);
             for (unsigned int f=0; f<NumberOfSignalFeatures; f++)
               m_FeatureData(sampleCounter,f) = pix[f];
 
             // additional feature images
             int add_feat_c = 0;
             for (auto img : m_AdditionalFeatureImages.at(t))
             {
               float v = GetImageValue<float>(itkP1, img, false);
               add_feat_c++;
               m_FeatureData(sampleCounter,NumberOfSignalFeatures+2+add_feat_c) = v;
             }
 
             // set label values
             float angle = 0;
             float m = dir.magnitude();
             if (m>0.0001)
             {
               int l = 0;
               for (auto d : m_DirectionContainer)
               {
                 float a = fabs(dot_product(dir, d));
                 if (a>angle)
                 {
                   m_LabelData(sampleCounter,0) = l;
                   m_Weights(sampleCounter,0) = fiber_weight*volume_mod;
                   angle = a;
                 }
                 l++;
               }
             }
 
             sampleCounter++;
           }
           if (!m_BidirectionalFiberSampling)  // don't sample fibers backward
             break;
         }
       }
     }
   }
   m_Tractograms.clear();
   MITK_INFO << "done";
 }
 
 }
 
 #endif
diff --git a/Modules/DiffusionImaging/FiberTracking/IODataStructures/mitkTractographyForest.cpp b/Modules/DiffusionImaging/FiberTracking/IODataStructures/mitkTractographyForest.cpp
index 670a6e8d91..7187027c20 100755
--- a/Modules/DiffusionImaging/FiberTracking/IODataStructures/mitkTractographyForest.cpp
+++ b/Modules/DiffusionImaging/FiberTracking/IODataStructures/mitkTractographyForest.cpp
@@ -1,116 +1,119 @@
 /*===================================================================
 
 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.
 
 ===================================================================*/
 
 
 #define _USE_MATH_DEFINES
 #include "mitkTractographyForest.h"
 #include <mitkExceptionMacro.h>
+#include <mitkGeometry3D.h>
 
 using namespace std;
 
 namespace mitk
 {
 
 TractographyForest::TractographyForest( std::shared_ptr< vigra::RandomForest<int> > forest )
 {
   m_Forest = forest;
+  mitk::Geometry3D::Pointer geometry = mitk::Geometry3D::New();
+  SetGeometry(geometry);
 }
 
 TractographyForest::~TractographyForest()
 {
 
 }
 
 void TractographyForest::PredictProbabilities(vigra::MultiArray<2, float>& features, vigra::MultiArray<2, float>& probabilities) const
 {
   m_Forest->predictProbabilities(features, probabilities);
 }
 
 int TractographyForest::GetNumFeatures() const
 {
   if (!HasForest())
     mitkThrow() << "Forest is NULL";
   return m_Forest->feature_count();
 }
 
 int TractographyForest::GetNumTrees() const
 {
   if (!HasForest())
     mitkThrow() << "Forest is NULL";
   return m_Forest->tree_count();
 }
 
 int TractographyForest::GetNumClasses() const
 {
   if (!HasForest())
     mitkThrow() << "Forest is NULL";
   return m_Forest->class_count();
 }
 
 int TractographyForest::GetMaxTreeDepth() const
 {
   if (!HasForest())
     mitkThrow() << "Forest is NULL";
   return m_Forest->ext_param_.max_tree_depth;
 }
 
 int TractographyForest::IndexToClassLabel(int idx) const
 {
   if (!HasForest())
     mitkThrow() << "Forest is NULL";
   unsigned int classLabel = 0;
   m_Forest->ext_param_.to_classlabel(idx, classLabel);
   return classLabel;
 }
 
 bool TractographyForest::HasForest() const
 {
   return m_Forest!=nullptr;
 }
 
 void TractographyForest::PrintSelf(std::ostream &os, itk::Indent indent) const
 {
 //  Superclass::PrintSelf(os, indent);
   os << indent << this->GetNameOfClass() << ":\n";
   os << indent << "Number of features: " << GetNumFeatures() << std::endl;
   os << indent << "Number of classes: " << GetNumClasses() << std::endl;
   os << indent << "Number of trees: " << GetNumTrees() << std::endl;
   os << indent << "Maximum tree depth: " << GetMaxTreeDepth() << std::endl;
 }
 
 /* ESSENTIAL IMPLEMENTATION OF SUPERCLASS METHODS */
 void TractographyForest::UpdateOutputInformation()
 {
 
 }
 void TractographyForest::SetRequestedRegionToLargestPossibleRegion()
 {
 
 }
 bool TractographyForest::RequestedRegionIsOutsideOfTheBufferedRegion()
 {
   return false;
 }
 bool TractographyForest::VerifyRequestedRegion()
 {
   return true;
 }
 void TractographyForest::SetRequestedRegion(const itk::DataObject* )
 {
 
 }
 
 }