diff --git a/Modules/Multilabel/mitkLabelSetImage.cpp b/Modules/Multilabel/mitkLabelSetImage.cpp
index 1f111035dc..8bdc0fa195 100644
--- a/Modules/Multilabel/mitkLabelSetImage.cpp
+++ b/Modules/Multilabel/mitkLabelSetImage.cpp
@@ -1,877 +1,889 @@
 /*===================================================================
 
 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 "mitkLabelSetImage.h"
 
 #include "mitkImageAccessByItk.h"
 #include "mitkInteractionConst.h"
 #include "mitkRenderingManager.h"
 #include "mitkImageCast.h"
 #include "mitkImageReadAccessor.h"
 #include "mitkLookupTableProperty.h"
 #include "mitkPadImageFilter.h"
 
 #include <vtkTransform.h>
 #include <vtkTransformPolyDataFilter.h>
 #include <vtkCell.h>
 
 #include <itkImageRegionIterator.h>
 #include <itkQuadEdgeMesh.h>
 #include <itkTriangleMeshToBinaryImageFilter.h>
 //#include <itkRelabelComponentImageFilter.h>
 
 #include <itkCommand.h>
 
 mitk::LabelSetImage::LabelSetImage() :
 mitk::Image(),
 m_ActiveLayer(0),
 m_ExteriorLabel(nullptr)
 {
   // Iniitlaize Background Label
   mitk::Color color;
   color.Set(0, 0, 0);
   m_ExteriorLabel = mitk::Label::New();
   m_ExteriorLabel->SetColor(color);
   m_ExteriorLabel->SetName("Exterior");
   m_ExteriorLabel->SetOpacity(0.0);
   m_ExteriorLabel->SetLocked(false);
   m_ExteriorLabel->SetValue(0);
 }
 
 mitk::LabelSetImage::LabelSetImage(const mitk::LabelSetImage & other) :
 Image(other),
 m_ActiveLayer(other.GetActiveLayer()),
 m_ExteriorLabel(other.GetExteriorLabel()->Clone())
 {
   for (unsigned int i = 0; i < other.GetNumberOfLayers(); i++)
   {
     // Clone LabelSet data
     mitk::LabelSet::Pointer lsClone = other.GetLabelSet(i)->Clone();
     // add modified event listener to LabelSet (listen to LabelSet changes)
     itk::SimpleMemberCommand<Self>::Pointer command = itk::SimpleMemberCommand<Self>::New();
     command->SetCallbackFunction(this, &mitk::LabelSetImage::OnLabelSetModified);
     lsClone->AddObserver(itk::ModifiedEvent(), command);
     m_LabelSetContainer.push_back(lsClone);
 
     // clone layer Image data
     mitk::Image::Pointer liClone = other.GetLayerImage(i)->Clone();
     m_LayerContainer.push_back(liClone);
   }
 }
 
 void mitk::LabelSetImage::OnLabelSetModified()
 {
   Superclass::Modified();
 }
 
 void mitk::LabelSetImage::SetExteriorLabel(mitk::Label * label)
 {
   m_ExteriorLabel = label;
 }
 
 mitk::Label* mitk::LabelSetImage::GetExteriorLabel()
 {
   return m_ExteriorLabel;
 }
 
 const mitk::Label* mitk::LabelSetImage::GetExteriorLabel() const
 {
   return m_ExteriorLabel;
 }
 
 void mitk::LabelSetImage::Initialize(const mitk::Image* other)
 {
   mitk::PixelType pixelType(mitk::MakeScalarPixelType<LabelSetImage::PixelType>());
   if (other->GetDimension() == 2)
   {
     const unsigned int dimensions[] = { other->GetDimension(0), other->GetDimension(1), 1 };
     Superclass::Initialize(pixelType, 3, dimensions);
   }
   else
   {
     Superclass::Initialize(pixelType, other->GetDimension(), other->GetDimensions());
   }
 
   auto originalGeometry = other->GetTimeGeometry()->Clone();
   this->SetTimeGeometry(originalGeometry);
 
   // Add a inital LabelSet ans corresponding image data to the stack
   AddLayer();
 }
 
 mitk::LabelSetImage::~LabelSetImage()
 {
   m_LabelSetContainer.clear();
 }
 
 mitk::Image* mitk::LabelSetImage::GetLayerImage(unsigned int layer)
 {
   return m_LayerContainer[layer];
 }
 
 const mitk::Image* mitk::LabelSetImage::GetLayerImage(unsigned int layer) const
 {
   return m_LayerContainer[layer];
 }
 
+void mitk::LabelSetImage::SetLayerImage(unsigned int layer, mitk::Image *image)
+{
+    if (image != nullptr)
+    {
+        m_LayerContainer[layer] = image;
+    }
+    else
+    {
+        mitkThrow() << "Tried to assign null image to layer.";
+    }
+}
+
 unsigned int mitk::LabelSetImage::GetActiveLayer() const
 {
   return m_ActiveLayer;
 }
 
 unsigned int mitk::LabelSetImage::GetNumberOfLayers() const
 {
   return m_LabelSetContainer.size();
 }
 
 void mitk::LabelSetImage::RemoveLayer()
 {
   int layerToDelete = GetActiveLayer();
   // remove all observers from active label set
   GetLabelSet(layerToDelete)->RemoveAllObservers();
 
   // set the active layer to one below, if exists.
   SetActiveLayer(layerToDelete-1);
 
   // remove labelset and image data
   m_LabelSetContainer.erase(m_LabelSetContainer.begin() + layerToDelete);
   m_LayerContainer.erase(m_LayerContainer.begin() + layerToDelete);
 
   this->Modified();
 }
 
 template<typename TPixel, unsigned int VDimensions>
 void SetToZero(itk::Image< TPixel, VDimensions> * source)
 {
   source->FillBuffer(0);
 }
 
 unsigned int mitk::LabelSetImage::AddLayer(mitk::LabelSet::Pointer lset)
 {
   mitk::Image::Pointer newImage = mitk::Image::New();
   newImage->Initialize( this->GetPixelType(), this->GetDimension(), this->GetDimensions(), this->GetImageDescriptor()->GetNumberOfChannels() );
   newImage->SetTimeGeometry(this->GetTimeGeometry()->Clone());
 
   if (newImage->GetDimension() < 4)
   {
     AccessByItk(newImage, SetToZero);
   }
   else
   {
     AccessFixedDimensionByItk(newImage, SetToZero, 4);
   }
 
   unsigned int newLabelSetId = this->AddLayer(newImage, lset);
 
   return newLabelSetId;
 }
 
 unsigned int mitk::LabelSetImage::AddLayer(mitk::Image::Pointer layerImage, mitk::LabelSet::Pointer lset)
 {
   unsigned int newLabelSetId = m_LayerContainer.size();
 
   // Add labelset to layer
   mitk::LabelSet::Pointer ls;
   if (lset.IsNotNull())
   {
     ls = lset;
   }
   else
   {
     ls = mitk::LabelSet::New();
     ls->AddLabel(GetExteriorLabel());
     ls->SetActiveLabel(0 /*Exterior Label*/);
   }
 
   ls->SetLayer(newLabelSetId);
   // Add exterior Label to label set
   //mitk::Label::Pointer exteriorLabel = CreateExteriorLabel();
 
   // push a new working image for the new layer
   m_LayerContainer.push_back(layerImage);
 
   // push a new labelset for the new layer
   m_LabelSetContainer.push_back(ls);
 
   // add modified event listener to LabelSet (listen to LabelSet changes)
   itk::SimpleMemberCommand<Self>::Pointer command = itk::SimpleMemberCommand<Self>::New();
   command->SetCallbackFunction(this, &mitk::LabelSetImage::OnLabelSetModified);
   ls->AddObserver(itk::ModifiedEvent(), command);
 
   SetActiveLayer(newLabelSetId);
   //MITK_INFO << GetActiveLayer();
   this->Modified();
   return newLabelSetId;
 }
 
 void mitk::LabelSetImage::AddLabelSetToLayer(const unsigned int layerIdx, const mitk::LabelSet::Pointer labelSet)
 {
   if (m_LayerContainer.size() <= layerIdx)
   {
     mitkThrow() << "Trying to add labelSet to non-existing layer.";
   }
 
   if (layerIdx < m_LabelSetContainer.size())
   {
     m_LabelSetContainer[layerIdx] = labelSet;
   }
   else
   {
     while (layerIdx >= m_LabelSetContainer.size())
     {
       mitk::LabelSet::Pointer defaultLabelSet = mitk::LabelSet::New();
       defaultLabelSet->AddLabel(GetExteriorLabel());
       defaultLabelSet->SetActiveLabel(0 /*Exterior Label*/);
       defaultLabelSet->SetLayer(m_LabelSetContainer.size());
       m_LabelSetContainer.push_back(defaultLabelSet);
     }
     m_LabelSetContainer.push_back(labelSet);
   }
 }
 
 void mitk::LabelSetImage::SetActiveLayer(unsigned int layer)
 {
   try
   {
     if ((layer != GetActiveLayer()) && (layer < this->GetNumberOfLayers()))
     {
       BeforeChangeLayerEvent.Send();
 
       AccessByItk_1(this, ImageToLayerContainerProcessing, GetActiveLayer());
       m_ActiveLayer = layer; // only at this place m_ActiveLayer should be manipulated!!! Use Getter and Setter
       AccessByItk_1(this, LayerContainerToImageProcessing, GetActiveLayer());
 
       AfterChangeLayerEvent.Send();
     }
   }
   catch (itk::ExceptionObject& e)
   {
     mitkThrow() << e.GetDescription();
   }
   this->Modified();
 }
 
 void mitk::LabelSetImage::Concatenate(mitk::LabelSetImage* other)
 {
   const unsigned int* otherDims = other->GetDimensions();
   const unsigned int* thisDims = this->GetDimensions();
   if ((otherDims[0] != thisDims[0]) || (otherDims[1] != thisDims[1]) || (otherDims[2] != thisDims[2]))
     mitkThrow() << "Dimensions do not match.";
 
   try
   {
     int numberOfLayers = other->GetNumberOfLayers();
     for (int layer = 0; layer < numberOfLayers; ++layer)
     {
       this->SetActiveLayer(layer);
       AccessByItk_1(this, ConcatenateProcessing, other);
       mitk::LabelSet * ls = other->GetLabelSet(layer);
       auto it = ls->IteratorConstBegin();
       auto end = ls->IteratorConstEnd();
       it++;// skip exterior
       while (it != end)
       {
         GetLabelSet()->AddLabel((it->second));
         //AddLabelEvent.Send();
         it++;
       }
     }
   }
   catch (itk::ExceptionObject& e)
   {
     mitkThrow() << e.GetDescription();
   }
   this->Modified();
 }
 
 void mitk::LabelSetImage::ClearBuffer()
 {
   try
   {
     AccessByItk(this, ClearBufferProcessing);
     this->Modified();
   }
   catch (itk::ExceptionObject& e)
   {
     mitkThrow() << e.GetDescription();
   }
 }
 
 bool mitk::LabelSetImage::ExistLabel(PixelType pixelValue) const
 {
   bool exist = false;
   for (unsigned int lidx = 0; lidx < GetNumberOfLayers(); lidx++)
     exist |= m_LabelSetContainer[lidx]->ExistLabel(pixelValue);
   return exist;
 }
 
 bool mitk::LabelSetImage::ExistLabel(PixelType pixelValue, unsigned int layer) const
 {
   bool exist = m_LabelSetContainer[layer]->ExistLabel(pixelValue);
   return exist;
 }
 
 bool mitk::LabelSetImage::ExistLabelSet(unsigned int layer) const
 {
   return layer < m_LabelSetContainer.size();
 }
 
 void mitk::LabelSetImage::MergeLabel(PixelType pixelValue, unsigned int /*layer*/)
 {
   int targetPixelValue = GetActiveLabel()->GetValue();
   try
   {
     AccessByItk_2(this, MergeLabelProcessing, targetPixelValue, pixelValue);
   }
   catch (itk::ExceptionObject& e)
   {
     mitkThrow() << e.GetDescription();
   }
   Modified();
 }
 
 void mitk::LabelSetImage::MergeLabels(std::vector<PixelType> &VectorOfLablePixelValues, PixelType pixelValue, unsigned int layer)
 {
   GetLabelSet(layer)->SetActiveLabel(pixelValue);
   try
   {
     for (unsigned int idx = 0; idx < VectorOfLablePixelValues.size(); idx++)
     {
       AccessByItk_2(this, MergeLabelProcessing, pixelValue, VectorOfLablePixelValues[idx]);
     }
   }
   catch (itk::ExceptionObject& e)
   {
     mitkThrow() << e.GetDescription();
   }
 }
 
 void mitk::LabelSetImage::RemoveLabels(std::vector<PixelType>& VectorOfLabelPixelValues, unsigned int layer)
 {
   for (unsigned int idx = 0; idx < VectorOfLabelPixelValues.size(); idx++)
   {
     GetLabelSet(layer)->RemoveLabel(VectorOfLabelPixelValues[idx]);
     EraseLabel(VectorOfLabelPixelValues[idx], layer);
   }
 }
 
 void mitk::LabelSetImage::EraseLabels(std::vector<PixelType> &VectorOfLabelPixelValues, unsigned int layer)
 {
   for (unsigned int i = 0; i < VectorOfLabelPixelValues.size(); i++)
   {
     this->EraseLabel(VectorOfLabelPixelValues[i], layer);
   }
 }
 
 void mitk::LabelSetImage::EraseLabel(PixelType pixelValue, unsigned int layer)
 {
   try
   {
     AccessByItk_2(this, EraseLabelProcessing, pixelValue, layer);
   }
   catch (itk::ExceptionObject& e)
   {
     mitkThrow() << e.GetDescription();
   }
   Modified();
 }
 
 mitk::Label *mitk::LabelSetImage::GetActiveLabel(unsigned int layer)
 {
   if (m_LabelSetContainer.size() > layer)
     return m_LabelSetContainer[layer]->GetActiveLabel();
   else
     return nullptr;
 }
 
 mitk::Label *mitk::LabelSetImage::GetLabel(PixelType pixelValue, unsigned int layer) const
 {
   if (m_LabelSetContainer.size() <= layer)
     return nullptr;
   else
     return m_LabelSetContainer[layer]->GetLabel(pixelValue);
 }
 
 mitk::LabelSet* mitk::LabelSetImage::GetLabelSet(unsigned int layer)
 {
   if (m_LabelSetContainer.size() <= layer)
     return nullptr;
   else
     return m_LabelSetContainer[layer].GetPointer();
 }
 
 const mitk::LabelSet* mitk::LabelSetImage::GetLabelSet(unsigned int layer) const
 {
   if (m_LabelSetContainer.size() <= layer)
     return nullptr;
   else
     return m_LabelSetContainer[layer].GetPointer();
 }
 
 mitk::LabelSet* mitk::LabelSetImage::GetActiveLabelSet()
 {
   if (m_LabelSetContainer.size() == 0)
     return nullptr;
   else
     return m_LabelSetContainer[GetActiveLayer()].GetPointer();
 }
 
 void mitk::LabelSetImage::UpdateCenterOfMass(PixelType pixelValue, unsigned int layer)
 {
   AccessByItk_2(this, CalculateCenterOfMassProcessing, pixelValue, layer);
 }
 
 unsigned int mitk::LabelSetImage::GetNumberOfLabels(unsigned int layer) const
 {
   return m_LabelSetContainer[layer]->GetNumberOfLabels();
 }
 
 unsigned int mitk::LabelSetImage::GetTotalNumberOfLabels() const
 {
   unsigned int totalLabels(0);
   auto layerIter = m_LabelSetContainer.begin();
   for (; layerIter != m_LabelSetContainer.end(); ++layerIter)
     totalLabels += (*layerIter)->GetNumberOfLabels();
   return totalLabels;
 }
 
 void mitk::LabelSetImage::MaskStamp(mitk::Image* mask, bool forceOverwrite)
 {
   try
   {
     mitk::PadImageFilter::Pointer padImageFilter = mitk::PadImageFilter::New();
     padImageFilter->SetInput(0, mask);
     padImageFilter->SetInput(1, this);
     padImageFilter->SetPadConstant(0);
     padImageFilter->SetBinaryFilter(false);
     padImageFilter->SetLowerThreshold(0);
     padImageFilter->SetUpperThreshold(1);
 
     padImageFilter->Update();
 
     mitk::Image::Pointer paddedMask = padImageFilter->GetOutput();
 
     if (paddedMask.IsNull()) return;
 
     AccessByItk_2(this, MaskStampProcessing, paddedMask, forceOverwrite);
   }
   catch (...)
   {
     mitkThrow() << "Could not stamp the provided mask on the selected label.";
   }
 }
 
 mitk::Image::Pointer mitk::LabelSetImage::CreateLabelMask(PixelType index)
 {
   mitk::Image::Pointer mask = mitk::Image::New();
   try
   {
     mask->Initialize(this);
 
     unsigned int byteSize = sizeof(LabelSetImage::PixelType);
     for (unsigned int dim = 0; dim < mask->GetDimension(); ++dim)
     {
       byteSize *= mask->GetDimension(dim);
     }
 
     mitk::ImageWriteAccessor* accessor = new mitk::ImageWriteAccessor(static_cast<mitk::Image*>(mask));
     memset(accessor->GetData(), 0, byteSize);
     delete accessor;
 
     auto geometry = this->GetTimeGeometry()->Clone();
     mask->SetTimeGeometry(geometry);
 
     AccessByItk_2(this, CreateLabelMaskProcessing, mask, index);
   }
   catch (...)
   {
     mitkThrow() << "Could not create a mask out of the selected label.";
   }
 
   return mask;
 }
 
 void mitk::LabelSetImage::InitializeByLabeledImage(mitk::Image::Pointer image)
 {
   if (image.IsNull() || image->IsEmpty() || !image->IsInitialized())
     mitkThrow() << "Invalid labeled image.";
 
   try
   {
     this->Initialize(image);
 
     unsigned int byteSize = sizeof(LabelSetImage::PixelType);
     for (unsigned int dim = 0; dim < image->GetDimension(); ++dim)
     {
       byteSize *= image->GetDimension(dim);
     }
 
     mitk::ImageWriteAccessor* accessor = new mitk::ImageWriteAccessor(static_cast<mitk::Image*>(this));
     memset(accessor->GetData(), 0, byteSize);
     delete accessor;
 
     auto geometry = image->GetTimeGeometry()->Clone();
     this->SetTimeGeometry(geometry);
 
     if (image->GetDimension() == 3)
     {
       AccessTwoImagesFixedDimensionByItk(this, image, InitializeByLabeledImageProcessing, 3);
     }
     else if (image->GetDimension() == 4)
     {
       AccessTwoImagesFixedDimensionByItk(this, image, InitializeByLabeledImageProcessing, 4);
     }
     else
     {
       mitkThrow() << image->GetDimension() << "-dimensional label set images not yet supported";
     }
   }
   catch (...)
   {
     mitkThrow() << "Could not intialize by provided labeled image.";
   }
   this->Modified();
 }
 
 template <typename LabelSetImageType, typename ImageType>
 void mitk::LabelSetImage::InitializeByLabeledImageProcessing(LabelSetImageType* labelSetImage, ImageType* image)
 {
   typedef itk::ImageRegionConstIteratorWithIndex<ImageType> SourceIteratorType;
   typedef itk::ImageRegionIterator<LabelSetImageType> TargetIteratorType;
 
   TargetIteratorType targetIter(labelSetImage, labelSetImage->GetRequestedRegion());
   targetIter.GoToBegin();
 
   SourceIteratorType sourceIter(image, image->GetRequestedRegion());
   sourceIter.GoToBegin();
 
   while (!sourceIter.IsAtEnd())
   {
     PixelType sourceValue = static_cast<PixelType>(sourceIter.Get());
     targetIter.Set(sourceValue);
 
     if (!this->ExistLabel(sourceValue))
     {
       std::stringstream name;
       name << "object-" << sourceValue;
 
       double rgba[4];
       m_LabelSetContainer[this->GetActiveLayer()]->GetLookupTable()->GetTableValue(sourceValue, rgba);
 
       mitk::Color color;
       color.SetRed(rgba[0]);
       color.SetGreen(rgba[1]);
       color.SetBlue(rgba[2]);
 
       auto label = mitk::Label::New();
       label->SetName(name.str().c_str());
       label->SetColor(color);
       label->SetOpacity(rgba[3]);
       label->SetValue(sourceValue);
 
       this->GetLabelSet()->AddLabel(label);
 
       if (GetActiveLabelSet()->GetNumberOfLabels() >= mitk::Label::MAX_LABEL_VALUE || sourceValue >= mitk::Label::MAX_LABEL_VALUE)
         this->AddLayer();
     }
 
     ++sourceIter;
     ++targetIter;
   }
 }
 
 template < typename ImageType >
 void mitk::LabelSetImage::MaskStampProcessing(ImageType* itkImage, mitk::Image* mask, bool forceOverwrite)
 {
   typename ImageType::Pointer itkMask;
   mitk::CastToItkImage(mask, itkMask);
 
   typedef itk::ImageRegionConstIterator< ImageType > SourceIteratorType;
   typedef itk::ImageRegionIterator< ImageType > TargetIteratorType;
 
   SourceIteratorType sourceIter(itkMask, itkMask->GetLargestPossibleRegion());
   sourceIter.GoToBegin();
 
   TargetIteratorType targetIter(itkImage, itkImage->GetLargestPossibleRegion());
   targetIter.GoToBegin();
 
   int activeLabel = this->GetActiveLabel(GetActiveLayer())->GetValue();
 
   while (!sourceIter.IsAtEnd())
   {
     PixelType sourceValue = sourceIter.Get();
     PixelType targetValue = targetIter.Get();
 
     if ((sourceValue != 0) && (forceOverwrite || !this->GetLabel(targetValue)->GetLocked())) // skip exterior and locked labels
     {
       targetIter.Set(activeLabel);
     }
     ++sourceIter;
     ++targetIter;
   }
 
   this->Modified();
 }
 
 template < typename ImageType >
 void mitk::LabelSetImage::CreateLabelMaskProcessing(ImageType* itkImage, mitk::Image* mask, PixelType index)
 {
   typename ImageType::Pointer itkMask;
   mitk::CastToItkImage(mask, itkMask);
 
   typedef itk::ImageRegionConstIterator< ImageType > SourceIteratorType;
   typedef itk::ImageRegionIterator< ImageType > TargetIteratorType;
 
   SourceIteratorType sourceIter(itkImage, itkImage->GetLargestPossibleRegion());
   sourceIter.GoToBegin();
 
   TargetIteratorType targetIter(itkMask, itkMask->GetLargestPossibleRegion());
   targetIter.GoToBegin();
 
   while (!sourceIter.IsAtEnd())
   {
     PixelType sourceValue = sourceIter.Get();
     if (sourceValue == index)
     {
       targetIter.Set(1);
     }
     ++sourceIter;
     ++targetIter;
   }
 }
 
 template < typename ImageType >
 void mitk::LabelSetImage::CalculateCenterOfMassProcessing(ImageType* itkImage, PixelType pixelValue, unsigned int layer)
 {
   // for now, we just retrieve the voxel in the middle
   typedef itk::ImageRegionConstIterator< ImageType > IteratorType;
   IteratorType iter(itkImage, itkImage->GetLargestPossibleRegion());
   iter.GoToBegin();
 
   std::vector< typename ImageType::IndexType > indexVector;
 
   while (!iter.IsAtEnd())
   {
     // TODO fix comparison warning more effective
     if (iter.Get() == pixelValue)
     {
       indexVector.push_back(iter.GetIndex());
     }
     ++iter;
   }
 
   mitk::Point3D pos;
   pos.Fill(0.0);
 
   if (!indexVector.empty())
   {
     typename itk::ImageRegionConstIteratorWithIndex< ImageType >::IndexType centerIndex;
     centerIndex = indexVector.at(indexVector.size() / 2);
     if (centerIndex.GetIndexDimension() == 3)
     {
       pos[0] = centerIndex[0];
       pos[1] = centerIndex[1];
       pos[2] = centerIndex[2];
     }
     else
       return;
   }
 
   GetLabelSet(layer)->GetLabel(pixelValue)->SetCenterOfMassIndex(pos);
   this->GetSlicedGeometry()->IndexToWorld(pos, pos); // TODO: TimeGeometry?
   GetLabelSet(layer)->GetLabel(pixelValue)->SetCenterOfMassCoordinates(pos);
 }
 
 template < typename ImageType >
 void mitk::LabelSetImage::ClearBufferProcessing(ImageType* itkImage)
 {
   itkImage->FillBuffer(0);
 }
 
 // todo: concatenate all layers and not just the active one
 template < typename ImageType >
 void mitk::LabelSetImage::ConcatenateProcessing(ImageType* itkTarget, mitk::LabelSetImage* other)
 {
   typename ImageType::Pointer itkSource = ImageType::New();
   mitk::CastToItkImage(other, itkSource);
 
   typedef itk::ImageRegionConstIterator< ImageType > ConstIteratorType;
   typedef itk::ImageRegionIterator< ImageType >      IteratorType;
 
   ConstIteratorType sourceIter(itkSource, itkSource->GetLargestPossibleRegion());
   IteratorType      targetIter(itkTarget, itkTarget->GetLargestPossibleRegion());
 
   int numberOfTargetLabels = this->GetNumberOfLabels(GetActiveLayer()) - 1; // skip exterior
   sourceIter.GoToBegin();
   targetIter.GoToBegin();
 
   while (!sourceIter.IsAtEnd())
   {
     PixelType sourceValue = sourceIter.Get();
     PixelType targetValue = targetIter.Get();
     if ((sourceValue != 0) && !this->GetLabel(targetValue)->GetLocked()) // skip exterior and locked labels
     {
       targetIter.Set(sourceValue + numberOfTargetLabels);
     }
     ++sourceIter;
     ++targetIter;
   }
 }
 
 template < typename ImageType >
 void mitk::LabelSetImage::LayerContainerToImageProcessing(ImageType* target, unsigned int layer)
 {
   typename ImageType::Pointer itkSource;
   mitk::CastToItkImage(m_LayerContainer[layer], itkSource);
   typedef itk::ImageRegionConstIterator< ImageType >  SourceIteratorType;
   typedef itk::ImageRegionIterator< ImageType >       TargetIteratorType;
 
   SourceIteratorType sourceIter(itkSource, itkSource->GetLargestPossibleRegion());
   sourceIter.GoToBegin();
 
   TargetIteratorType targetIter(target, target->GetLargestPossibleRegion());
   targetIter.GoToBegin();
 
   while (!sourceIter.IsAtEnd())
   {
     targetIter.Set(sourceIter.Get());
     ++sourceIter;
     ++targetIter;
   }
 }
 
 template < typename ImageType >
 void mitk::LabelSetImage::ImageToLayerContainerProcessing(ImageType* source, unsigned int layer) const
 {
   typename ImageType::Pointer itkTarget;
   mitk::CastToItkImage(m_LayerContainer[layer], itkTarget);
 
   typedef itk::ImageRegionConstIterator< ImageType >   SourceIteratorType;
   typedef itk::ImageRegionIterator< ImageType >        TargetIteratorType;
 
   SourceIteratorType sourceIter(source, source->GetLargestPossibleRegion());
   sourceIter.GoToBegin();
 
   TargetIteratorType targetIter(itkTarget, itkTarget->GetLargestPossibleRegion());
   targetIter.GoToBegin();
 
   while (!sourceIter.IsAtEnd())
   {
     targetIter.Set(sourceIter.Get());
     ++sourceIter;
     ++targetIter;
   }
 }
 
 template < typename ImageType >
 void mitk::LabelSetImage::EraseLabelProcessing(ImageType* itkImage, PixelType pixelValue, unsigned int /*layer*/)
 {
   typedef itk::ImageRegionIterator< ImageType > IteratorType;
 
   IteratorType iter(itkImage, itkImage->GetLargestPossibleRegion());
   iter.GoToBegin();
 
   while (!iter.IsAtEnd())
   {
     PixelType value = iter.Get();
 
     if (value == pixelValue)
     {
       iter.Set(0);
     }
     ++iter;
   }
 }
 
 template < typename ImageType >
 void mitk::LabelSetImage::MergeLabelProcessing(ImageType* itkImage, PixelType pixelValue, PixelType index)
 {
   typedef itk::ImageRegionIterator< ImageType > IteratorType;
 
   IteratorType iter(itkImage, itkImage->GetLargestPossibleRegion());
   iter.GoToBegin();
 
   while (!iter.IsAtEnd())
   {
     if (iter.Get() == index)
     {
       iter.Set(pixelValue);
     }
     ++iter;
   }
 }
 
 bool mitk::Equal(const mitk::LabelSetImage& leftHandSide, const mitk::LabelSetImage& rightHandSide, ScalarType eps, bool verbose)
 {
   bool returnValue = true;
 
   /* LabelSetImage members */
 
   MITK_INFO(verbose) << "--- LabelSetImage Equal ---";
 
   // number layers
   returnValue = leftHandSide.GetNumberOfLayers() == rightHandSide.GetNumberOfLayers();
   if (!returnValue)
   {
     MITK_INFO(verbose) << "Number of layers not equal.";
     return false;
   }
 
   // total number labels
   returnValue = leftHandSide.GetTotalNumberOfLabels() == rightHandSide.GetTotalNumberOfLabels();
   if (!returnValue)
   {
     MITK_INFO(verbose) << "Total number of labels not equal.";
     return false;
   }
 
   // active layer
   returnValue = leftHandSide.GetActiveLayer() == rightHandSide.GetActiveLayer();
   if (!returnValue)
   {
     MITK_INFO(verbose) << "Active layer not equal.";
     return false;
   }
 
   // working image data
   returnValue = mitk::Equal((const mitk::Image&)leftHandSide, (const mitk::Image&)rightHandSide, eps, verbose);
   if (!returnValue)
   {
     MITK_INFO(verbose) << "Working image data not equal.";
     return false;
   }
 
   for (unsigned int layerIndex = 0; layerIndex < leftHandSide.GetNumberOfLayers(); layerIndex++)
   {
     // layer image data
     returnValue = mitk::Equal(*leftHandSide.GetLayerImage(layerIndex), *rightHandSide.GetLayerImage(layerIndex), eps, verbose);
     if (!returnValue)
     {
       MITK_INFO(verbose) << "Layer image data not equal.";
       return false;
     }
     // layer labelset data
 
     returnValue = mitk::Equal(*leftHandSide.GetLabelSet(layerIndex), *rightHandSide.GetLabelSet(layerIndex), eps, verbose);
     if (!returnValue)
     {
       MITK_INFO(verbose) << "Layer labelset data not equal.";
       return false;
     }
   }
 
   return returnValue;
 }
diff --git a/Modules/Multilabel/mitkLabelSetImage.h b/Modules/Multilabel/mitkLabelSetImage.h
index 89b11d32b3..6c3800dcdd 100644
--- a/Modules/Multilabel/mitkLabelSetImage.h
+++ b/Modules/Multilabel/mitkLabelSetImage.h
@@ -1,347 +1,349 @@
 /*===================================================================
 
 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 __mitkLabelSetImage_H_
 #define __mitkLabelSetImage_H_
 
 #include <mitkImage.h>
 #include <mitkLabelSet.h>
 
 #include <MitkMultilabelExports.h>
 
 namespace mitk
 {
 
 //##Documentation
 //## @brief LabelSetImage class for handling labels and layers in a segmentation session.
 //##
 //## Handles operations for adding, removing, erasing and editing labels and layers.
 //## @ingroup Data
 
 class MITKMULTILABEL_EXPORT LabelSetImage : public Image
 {
 
 public:
 
   mitkClassMacro(LabelSetImage, Image)
   itkNewMacro(Self)
 
   typedef mitk::Label::PixelType PixelType;
 
   /**
   * \brief BeforeChangeLayerEvent (e.g. used for GUI integration)
   * As soon as active labelset should be changed, the signal emits.
   * Emitted by SetActiveLayer(int layer);
   */
   Message<> BeforeChangeLayerEvent;
 
   /**
   * \brief AfterchangeLayerEvent (e.g. used for GUI integration)
   * As soon as active labelset was changed, the signal emits.
   * Emitted by SetActiveLayer(int layer);
   */
   Message<> AfterChangeLayerEvent;
 
   /**
    * @brief Initialize an empty mitk::LabelSetImage using the information
    *        of an mitk::Image
    * @param image the image which is used for initializing the mitk::LabelSetImage
    */
   using mitk::Image::Initialize;
   virtual void Initialize(const mitk::Image* image) override;
 
   /**
     * \brief  */
   void Concatenate(mitk::LabelSetImage* image);
 
   /**
     * \brief  */
   void ClearBuffer();
 
 
   /**
    * @brief Merges the mitk::Label with a given target value with the active label
    * @param targetPixelValue the value of the mitk::Label that should be merged with the active one
    * @param layer the layer in which the merge should be performed
    */
   void MergeLabel(PixelType targetPixelValue, unsigned int layer = 0);
 
   /**
    * @brief Merges a list of mitk::Labels with the mitk::Label that has a specific value
    * @param VectorOfLablePixelValues the list of labels that should be merge with the specified one
    * @param index the value of the label into which the other should be merged
    * @param layer the layer in which the merge should be performed
    */
   void MergeLabels(std::vector<PixelType>& VectorOfLablePixelValues, PixelType index, unsigned int layer = 0);
 
   /**
     * \brief  */
   void UpdateCenterOfMass(PixelType pixelValue, unsigned int layer =0);
 
 
   /**
    * @brief Removes labels from the mitk::LabelSet of given layer.
    *        Calls mitk::LabelSetImage::EraseLabels() which also removes the labels from within the image.
    * @param VectorOfLabelPixelValues a list of labels to be removed
    * @param layer the layer in which the labels should be removed
    */
   void RemoveLabels(std::vector<PixelType>& VectorOfLabelPixelValues, unsigned int layer = 0);
 
   /**
    * @brief Erases the label with the given value in the given layer from the underlying image.
    *        The label itself will not be erased from the respective mitk::LabelSet. In order to
    *        remove the label itself use mitk::LabelSetImage::RemoveLabels()
    * @param pixelValue the label which will be remove from the image
    * @param layer the layer in which the label should be removed
    */
   void EraseLabel(PixelType pixelValue, unsigned int layer = 0);
 
   /**
    * @brief Similar to mitk::LabelSetImage::EraseLabel() this funtion erase a list of labels from the image
    * @param VectorOfLabelPixelValues the list of labels that should be remove
    * @param layer the layer for which the labels should be removed
    */
   void EraseLabels(std::vector<PixelType>& VectorOfLabelPixelValues, unsigned int layer = 0);
 
   /**
     * \brief  Returns true if the value exists in one of the labelsets*/
   bool ExistLabel(PixelType pixelValue) const;
 
   /**
    * @brief Checks if a label exists in a certain layer
    * @param pixelValue the label value
    * @param layer the layer in which should be searched for the label
    * @return true if the label exists otherwise false
    */
   bool ExistLabel(PixelType pixelValue, unsigned int layer) const;
 
   /**
     * \brief  Returns true if the labelset exists*/
   bool ExistLabelSet(unsigned int layer) const;
 
   /**
    * @brief Returns the active label of a specific layer
    * @param layer the layer ID for which the active label should be returned
    * @return the active label of the specified layer
    */
   mitk::Label* GetActiveLabel(unsigned int layer = 0);
 
   /**
    * @brief Returns the mitk::Label with the given pixelValue and for the given layer
    * @param pixelValue the pixel value of the label
    * @param layer the layer in which the labels should be located
    * @return the mitk::Label if available otherwise NULL
    */
   mitk::Label* GetLabel(PixelType pixelValue, unsigned int layer = 0) const;
 
   /**
    * @brief Returns the currently active mitk::LabelSet
    * @return the mitk::LabelSet of the active layer or NULL if non is present
    */
   mitk::LabelSet* GetActiveLabelSet();
 
   /**
    * @brief Gets the mitk::LabelSet for the given layer
    * @param layer the layer for which the mitk::LabelSet should be retrieved
    * @return the respective mitk::LabelSet or NULL if non exists for the given layer
    */
   mitk::LabelSet * GetLabelSet(unsigned int layer = 0);
   const mitk::LabelSet * GetLabelSet(unsigned int layer = 0) const;
 
   /**
    * @brief Gets the ID of the currently active layer
    * @return the ID of the active layer
    */
   unsigned int GetActiveLayer() const;
 
   /**
    * @brief Get the number of all existing mitk::Labels for a given layer
    * @param layer the layer ID for which the active mitk::Labels should be retrieved
    * @return the number of all existing mitk::Labels for the given layer
    */
   unsigned int GetNumberOfLabels(unsigned int layer = 0) const;
 
   /**
    * @brief Returns the number of all labels summed up across all layers
    * @return the overall number of labels across all layers
    */
   unsigned int GetTotalNumberOfLabels() const;
 
   // This function will need to be ported to an external class
   // it requires knowledge of pixeltype and dimension and includes
   // too much algorithm to be sensibly part of a data class
   ///**
   //  * \brief  */
   //void SurfaceStamp(mitk::Surface* surface, bool forceOverwrite);
 
   /**
     * \brief  */
   mitk::Image::Pointer CreateLabelMask(PixelType index);
 
   /**
    * @brief Initialize a new mitk::LabelSetImage by an given image.
    * For all distinct pixel values of the parameter image new labels will
    * be created. If the number of distinct pixel values exceeds mitk::Label::MAX_LABEL_VALUE
    * a new layer will be created
    * @param image the image which is used for initialization
    */
   void InitializeByLabeledImage(mitk::Image::Pointer image);
 
   /**
     * \brief  */
   void MaskStamp(mitk::Image* mask, bool forceOverwrite);
 
   /**
     * \brief  */
   void SetActiveLayer(unsigned int layer);
 
   /**
     * \brief  */
   unsigned int GetNumberOfLayers() const;
 
   /**
    * @brief Adds a new layer to the LabelSetImage. The new layer will be set as the active one
    * @param layer a mitk::LabelSet which will be set as new layer.
    * @return the layer ID of the new layer
    */
   unsigned int AddLayer(mitk::LabelSet::Pointer layer =nullptr);
 
   /**
   * \brief Add a layer based on a provided mitk::Image
   * \param layerImage is added to the vector of label images
   * \param lset a label set that will be added to the new layer if provided
   *\return the layer ID of the new layer
   */
   unsigned int AddLayer(mitk::Image::Pointer layerImage, mitk::LabelSet::Pointer lset = nullptr);
 
   /**
   * \brief Add a LabelSet to an existing layer
   *
   * This will replace an existing labelSet if one exists. Throws an exceptions if you are trying
   * to add a labelSet to a non-existing layer.
   *
   * If there are no labelSets for layers with an id less than layerIdx default ones will be added
   * for them.
   *
   * \param layerIdx The index of the layer the LabelSet should be added to
   * \param labelSet The LabelSet that should be added
   */
   void AddLabelSetToLayer(const unsigned int layerIdx, const mitk::LabelSet::Pointer labelSet);
 
   /**
    * @brief Removes the active layer and the respective mitk::LabelSet and image information.
    *        The new active layer is the one below, if exists
    */
   void RemoveLayer();
 
   /**
     * \brief  */
   mitk::Image* GetLayerImage(unsigned int layer);
 
   const mitk::Image* GetLayerImage(unsigned int layer) const;
 
+  void SetLayerImage(unsigned int layer, mitk::Image* image);
+
   void OnLabelSetModified();
 
   /**
    * @brief Sets the label which is used as default exterior label when creating a new layer
    * @param label the label which will be used as new exterior label
    */
   void SetExteriorLabel(mitk::Label * label);
 
   /**
    * @brief Gets the mitk::Label which is used as default exterior label
    * @return the exterior mitk::Label
    */
   mitk::Label* GetExteriorLabel();
 
   const mitk::Label* GetExteriorLabel() const;
 
 protected:
 
   mitkCloneMacro(Self)
 
   LabelSetImage();
   LabelSetImage(const LabelSetImage & other);
   virtual ~LabelSetImage();
 
   template < typename ImageType1, typename ImageType2 >
   void ChangeLayerProcessing( ImageType1* source, ImageType2* target );
 
   template < typename ImageType >
   void LayerContainerToImageProcessing( ImageType* source, unsigned int layer);
 
   template < typename ImageType >
   void ImageToLayerContainerProcessing( ImageType* source, unsigned int layer) const;
 
   template < typename ImageType >
   void CalculateCenterOfMassProcessing( ImageType* input, PixelType index, unsigned int layer);
 
   template < typename ImageType >
   void ClearBufferProcessing( ImageType* input);
 
   template < typename ImageType >
   void EraseLabelProcessing( ImageType* input, PixelType index, unsigned int layer);
 
 //  template < typename ImageType >
 //  void ReorderLabelProcessing( ImageType* input, int index, int layer);
 
   template < typename ImageType >
   void MergeLabelProcessing( ImageType* input, PixelType pixelValue, PixelType index);
 
   template < typename ImageType >
   void ConcatenateProcessing( ImageType* input, mitk::LabelSetImage* other);
 
   template < typename ImageType >
   void MaskStampProcessing( ImageType* input, mitk::Image* mask, bool forceOverwrite);
 
   template < typename ImageType >
   void CreateLabelMaskProcessing( ImageType* input, mitk::Image* mask, PixelType index);
 
   template < typename LabelSetImageType, typename ImageType >
   void InitializeByLabeledImageProcessing(LabelSetImageType* input, ImageType* other);
 
   std::vector< LabelSet::Pointer > m_LabelSetContainer;
   std::vector< Image::Pointer > m_LayerContainer;
 
   int m_ActiveLayer;
 
   mitk::Label::Pointer m_ExteriorLabel;
 
 };
 
 /**
 * @brief Equal A function comparing two label set images for beeing equal in meta- and imagedata
 *
 * @ingroup MITKTestingAPI
 *
 * Following aspects are tested for equality:
 *  - LabelSetImage members
 *  - working image data
 *  - layer image data
 *  - labels in label set
 *
 * @param rightHandSide An image to be compared
 * @param leftHandSide An image to be compared
 * @param eps Tolarence for comparison. You can use mitk::eps in most cases.
 * @param verbose Flag indicating if the user wants detailed console output or not.
 * @return true, if all subsequent comparisons are true, false otherwise
 */
 MITKMULTILABEL_EXPORT bool Equal( const mitk::LabelSetImage& leftHandSide, const mitk::LabelSetImage& rightHandSide, ScalarType eps, bool verbose );
 
 
 } // namespace mitk
 
 #endif // __mitkLabelSetImage_H_
diff --git a/Plugins/org.mitk.gui.qt.activelearning/src/internal/QmitkActiveLearning.cpp b/Plugins/org.mitk.gui.qt.activelearning/src/internal/QmitkActiveLearning.cpp
index 6e1ac9f614..f577dccd95 100644
--- a/Plugins/org.mitk.gui.qt.activelearning/src/internal/QmitkActiveLearning.cpp
+++ b/Plugins/org.mitk.gui.qt.activelearning/src/internal/QmitkActiveLearning.cpp
@@ -1,930 +1,982 @@
 /*===================================================================
 
 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.
 
 ===================================================================*/
 
 // Blueberry
 #include <berryISelectionService.h>
 #include <berryIWorkbenchWindow.h>
 
 // Qt
 #include <QColorDialog>
 #include <QMessageBox>
 #include <QFuture>
 #include <QtConcurrent>
 
 // Qmitk
 #include "QmitkActiveLearning.h"
 
 // MITK
 #include <usModuleRegistry.h>
 #include <mitkImage.h>
 #include <mitkImageWriteAccessor.h>
 #include <mitkImagePixelReadAccessor.h>
 #include <mitkLabel.h>
 #include <mitkBaseGeometry.h>
 #include <mitkNumericTypes.h>
 #include <mitkPixelType.h>
 #include <mitkImageCast.h>
 #include <mitkImageAccessByItk.h>
 
 // ITK
 #include <itkDiscreteGaussianImageFilter.h>
 #include <itkGradientMagnitudeRecursiveGaussianImageFilter.h>
 #include <itkLaplacianRecursiveGaussianImageFilter.h>
 #include <itkHessianMatrixEigenvalueImageFilter.h>
 #include <itkStructureTensorEigenvalueImageFilter.h>
 #include <itkHessianRecursiveGaussianImageFilter.h>
 #include <itkSymmetricEigenAnalysisImageFilter.h>
 #include <itkImageRegionIterator.h>
 
 typedef double FeaturePixelType;
 
 // Returns true if list has at least one entry and all entries are valid mitk::Images, otherwise false
 static bool SelectionAllImages(const QList<mitk::DataNode::Pointer>& nodes)
 {
     if (nodes.empty())
     {
         return false;
     }
     for (const auto& node : nodes)
     {
         if(!(node.IsNotNull() && dynamic_cast<mitk::Image*>(node->GetData()) != nullptr)) return false;
     }
     return true;
 }
 
 // QColor to mitk::Color
 static mitk::Color QColorToMitkColor(const QColor& qcolor)
 {
     mitk::Color color;
     color.SetRed((float)qcolor.red() / 255);
     color.SetGreen((float)qcolor.green() / 255);
     color.SetBlue((float)qcolor.blue() / 255);
     return color;
 }
 
 // For debugging
 static void PrintAllLabels(mitk::LabelSetImage* image)
 {
     for (auto it=image->GetActiveLabelSet()->IteratorBegin(); it!=image->GetActiveLabelSet()->IteratorConstEnd(); ++it)
     {
         MITK_INFO << "Key: " << it->first << " - Name: " << it->second->GetName() << " - Value: " << it->second->GetValue() << " - Color: " << it->second->GetColor();
     }
 }
 
 // Make values of labels a consistent range
 static void FillLabelValues(mitk::LabelSetImage* image)
 {
     int value(0);
     for (auto it=image->GetActiveLabelSet()->IteratorBegin(); it!=image->GetActiveLabelSet()->IteratorConstEnd(); ++it)
     {
         it->second->SetValue(value);
         value++;
     }
     image->GetActiveLabelSet()->SetActiveLabel(0);
 }
 
 // Fill image with zeros
 static void FillWithZeros(mitk::Image* image)
 {
     unsigned int size = image->GetPixelType().GetSize();
     for (unsigned int i=0; i<image->GetDimension(); i++)
     {
         size *= image->GetDimension(i);
     }
     for (unsigned int t=0; t<image->GetTimeSteps(); t++)
     {
         mitk::ImageWriteAccessor accessor(image, image->GetVolumeData(0));
         memset(accessor.GetData(), 0, size);
     }
 }
 
 template<typename MatrixElementType>
 static Eigen::Matrix<MatrixElementType, Eigen::Dynamic, Eigen::Dynamic> Transform(const std::vector<mitk::Image::Pointer> images)
 {
     // Find size for output matrix [number of voxels, number of feature images]
     unsigned int size = images[0]->GetDimension(0);
     for (unsigned int i=1; i<3; ++i)
     {
         size *= images[0]->GetDimension(i);
     }
     Eigen::Matrix<MatrixElementType, Eigen::Dynamic, Eigen::Dynamic> outputMatrix(size, images.size());
 
     int j = 0;
     for (auto image : images)
     {
         int i = 0;
         typename itk::Image<MatrixElementType, 3>::Pointer imageItk;
         mitk::CastToItkImage(image, imageItk);
         auto it = itk::ImageRegionConstIterator<itk::Image<MatrixElementType, 3>>(imageItk, imageItk->GetLargestPossibleRegion());
         while (!it.IsAtEnd())
         {
             outputMatrix(i, j) = it.Get();
             ++it;
         }
         ++j;
     }
 
     return outputMatrix;
 }
 
 template<typename MatrixElementType, typename OutputPixelType>
-static void Transform(const Eigen::Matrix<MatrixElementType, Eigen::Dynamic, Eigen::Dynamic> &inputMatrix, mitk::Image::Pointer outputImage)
+static mitk::Image::Pointer Transform(const Eigen::Matrix<MatrixElementType, Eigen::Dynamic, Eigen::Dynamic> &inputMatrix, mitk::Image::Pointer referenceImage)
 {
     typename itk::Image<OutputPixelType, 3>::Pointer imageItk;
+    auto outputImage = mitk::Image::New();
+    outputImage->Initialize(referenceImage);
     mitk::CastToItkImage(outputImage, imageItk);
 
     auto it = itk::ImageRegionIterator<itk::Image<OutputPixelType, 3>>(imageItk, imageItk->GetLargestPossibleRegion());
     int i = 0;
     while (!it.IsAtEnd())
     {
         it.Set(inputMatrix(i, 0));
         ++it;
         ++i;
     }
 
+    int s = 5;
+    for (int i=0; i<s; ++i)
+    {
+        MITK_INFO << "------" << i << "------";
+        for (int j=0; j<s; ++j)
+        {
+            std::stringstream ss;
+            for (int k=0; k<s; ++k)
+            {
+                ss << imageItk->GetPixel({i, j, k}) << " ";
+            }
+            MITK_INFO << ss.str();
+        }
+    }
+
     outputImage = mitk::GrabItkImageMemory(imageItk);
+
+    s = 5;
+    for (int i=0; i<s; ++i)
+    {
+        MITK_INFO << "------" << i << "------";
+        for (int j=0; j<s; ++j)
+        {
+            std::stringstream ss;
+            for (int k=0; k<s; ++k)
+            {
+                ss << outputImage->GetPixelValueByIndex({i, j, k}) << " ";
+            }
+            MITK_INFO << ss.str();
+        }
+    }
+
+    return outputImage;
 }
 
 /* ==================================================================
  * FEATURES
  * =============================================================== */
 
 template<typename PixelType, unsigned int imageDimension>
 static void GaussianSmoothing(const itk::Image<PixelType, imageDimension>* inputImage, const double sigma, mitk::Image::Pointer outputImage)
 {
     typedef itk::Image<PixelType, imageDimension> ImageType;
     typedef itk::Image<FeaturePixelType, imageDimension> FeatureImageType;
     auto filter = itk::DiscreteGaussianImageFilter<ImageType, FeatureImageType>::New();
     filter->SetInput(inputImage);
     filter->SetVariance(sigma*sigma);
     filter->Update();
     mitk::GrabItkImageMemory(filter->GetOutput(), outputImage);
 }
 
 template<typename PixelType, unsigned int imageDimension>
 static void GaussianGradientMagnitude(const itk::Image<PixelType, imageDimension>* inputImage, const double sigma, mitk::Image::Pointer outputImage)
 {
     typedef itk::Image<PixelType, imageDimension> ImageType;
     typedef itk::Image<FeaturePixelType, imageDimension> FeatureImageType;
     auto filter = itk::GradientMagnitudeRecursiveGaussianImageFilter<ImageType, FeatureImageType>::New();
     filter->SetInput(inputImage);
     filter->SetSigma(sigma);
     filter->Update();
     mitk::GrabItkImageMemory(filter->GetOutput(), outputImage);
 }
 
 template<typename PixelType, unsigned int imageDimension>
 static void LaplacianOfGaussian(const itk::Image<PixelType, imageDimension>* inputImage, const double sigma, mitk::Image::Pointer outputImage)
 {
     typedef itk::Image<PixelType, imageDimension> ImageType;
     typedef itk::Image<FeaturePixelType, imageDimension> FeatureImageType;
     auto filter = itk::LaplacianRecursiveGaussianImageFilter<ImageType, FeatureImageType>::New();
     filter->SetInput(inputImage);
     filter->SetSigma(sigma);
     filter->Update();
     mitk::GrabItkImageMemory(filter->GetOutput(), outputImage);
 }
 
 //template<typename PixelType, unsigned int imageDimension>
 //static void StructureTensorEigenvalues(const itk::Image<PixelType, imageDimension>* inputImage, float sigma, std::vector<mitk::Image::Pointer> outputImages)
 //{
 //    typedef itk::Image<PixelType, imageDimension> ImageType;
 //    auto filter = itk::StructureTensorEigenvalueImageFilter<ImageType>::New();
 //    filter->SetInput(inputImage);
 //    filter->SetInnerScale(sigma);
 //    filter->SetOuterScale(sigma);
 //    filter->Update();
 //    for (unsigned int i=0; i<filter->GetNumberOfOutputs(); i++)
 //    {
 //        mitk::GrabItkImageMemory(filter->GetOutput(i), outputImages[i]);
 //    }
 //}
 
 //template<typename PixelType, unsigned int imageDimension>
 //static void HessianEigenvalues(const itk::Image<PixelType, imageDimension>* inputImage, float sigma, std::vector<mitk::Image::Pointer> outputImages)
 //{
 //    typedef itk::Image<PixelType, imageDimension> ImageType;
 //    typedef itk::Image<itk::SymmetricSecondRankTensor<PixelType, imageDimension>, imageDimension> TensorImageType;
 //    auto filter = itk::HessianRecursiveGaussianImageFilter<ImageType, TensorImageType>::New();
 
 //    ImageType::Pointer o1, o2, o3;
 //    o1->Allocate();
 //    o2->Allocate();
 //    o3->Allocate();
 
 //    filter->SetInput(inputImage);
 //    filter->SetSigma(sigma);
 //    filter->Update();
 //    TensorImageType::Pointer tensorImage = filter->GetOutput();
 
 //    itk::ImageRegionIterator<TensorImageType> tensorIt(tensorImage, tensorImage->GetLargestPossibleRegion());
 //    itk::ImageRegionIterator<ImageType> o1It(o1, o1->GetLargestPossibleRegion());
 //    itk::ImageRegionIterator<ImageType> o2It(o2, o2->GetLargestPossibleRegion());
 //    itk::ImageRegionIterator<ImageType> o3It(o3, o3->GetLargestPossibleRegion());
 
 //    while (!tensorIt.IsAtEnd())
 //    {
 //	itk::SymmetricSecondRankTensor<PixelType, imageDimension>::EigenValue
 
 //    for (unsigned int i=0; i<filter->GetNumberOfOutputs(); i++)
 //    {
 //        mitk::GrabItkImageMemory(filter->GetOutput(i), outputImages[i]);
 //    }
 //}
 
 /* ==================================================================
  * PUBLIC SLOTS
  * =============================================================== */
 
 
 void ActiveLearning::Initialize()
 {
     // Get selected nodes and check again if these are all images
     m_Nodes = this->GetDataManagerSelection();
     if (!SelectionAllImages(m_Nodes)) return;
 
     m_Controls.m_InitializePushButton->setDisabled(true);
 
     // Set names to the label (again)
     QString nameList = QString::fromStdString(m_Nodes[0]->GetName());
     if (m_Nodes.length() >= 2)
     {
         for (int i=1; i<m_Nodes.length(); ++i)
         {
             nameList += QString("<br />");
             nameList += QString::fromStdString(m_Nodes[i]->GetName());
         }
     }
     m_Controls.m_InitializeLabel->setText(nameList);
 
     // =======================================
     // PREDICTION IMAGE
     // =======================================
 
     m_PredictionImage = mitk::Image::New();
 
     try
     {
         mitk::Image::Pointer referenceImage = dynamic_cast<mitk::Image*>(m_Nodes[0]->GetData());
         m_PredictionImage->Initialize(referenceImage);
     }
     catch (mitk::Exception& e)
     {
         MITK_ERROR << "Exception caught: " << e.GetDescription();
         QMessageBox::information(m_Parent, "Error", "Could not initialize prediction image");
         return;
     }
 
     FillWithZeros(m_PredictionImage);
 
     m_PredictionNode = mitk::DataNode::New();
     m_PredictionNode->SetData(m_PredictionImage);
     m_PredictionNode->SetName("Predictions");
     m_PredictionNode->SetColor(0., 0., 0.);
 //    m_PredictionNode->SetBoolProperty("helper object", true);
 //    m_PredictionImage->GetExteriorLabel()->SetProperty("name.parent", mitk::StringProperty::New(nodes[0]->GetName().c_str()));
 //    m_PredictionImage->GetExteriorLabel()->SetProperty("name.image", mitk::StringProperty::New("Predictions"));
     this->GetDataStorage()->Add(m_PredictionNode, m_Nodes[0]);
 
     // =======================================
     // SEGMENTATION IMAGE
     // =======================================
 
     m_SegmentationImage = mitk::LabelSetImage::New();
 
     try
     {
         auto referenceImage = dynamic_cast<mitk::Image*>(m_Nodes[0]->GetData());
         m_SegmentationImage->Initialize(referenceImage);
     }
     catch (mitk::Exception& e)
     {
         MITK_ERROR << "Exception caught: " << e.GetDescription();
         QMessageBox::information(m_Parent, "Error", "Could not initialize segmentation image");
         return;
     }
 
     FillWithZeros(m_SegmentationImage);
 
     m_SegmentationNode = mitk::DataNode::New();
     m_SegmentationNode->SetData(m_SegmentationImage);
     m_SegmentationNode->SetName("Segmentation");
     m_SegmentationNode->SetColor(0., 0., 0.);
 //    m_PredictionNode->SetBoolProperty("helper object", true);
     m_SegmentationImage->GetExteriorLabel()->SetProperty("name.parent", mitk::StringProperty::New(m_Nodes[0]->GetName().c_str()));
     m_SegmentationImage->GetExteriorLabel()->SetProperty("name.image", mitk::StringProperty::New("Segmentation"));
     this->GetDataStorage()->Add(m_SegmentationNode, m_Nodes[0]);
 
     // =======================================
     // ANNOTATION IMAGE
     // =======================================
 
     m_AnnotationImage = mitk::LabelSetImage::New();
 
     try
     {
         auto referenceImage = dynamic_cast<mitk::Image*>(m_Nodes[0]->GetData());
         m_AnnotationImage->Initialize(referenceImage);
     }
     catch (mitk::Exception& e)
     {
         MITK_ERROR << "Exception caught: " << e.GetDescription();
         QMessageBox::information(m_Parent, "Error", "Could not initialize annotation image");
         return;
     }
 
     FillWithZeros(m_AnnotationImage);
 
     m_AnnotationNode = mitk::DataNode::New();
     m_AnnotationNode->SetData(m_AnnotationImage);
     m_AnnotationNode->SetName("Labels");
     m_AnnotationNode->SetColor(0., 0., 0.);
 //    m_AnnotationNode->SetBoolProperty("helper object", true);
     m_AnnotationImage->GetExteriorLabel()->SetProperty("name.parent", mitk::StringProperty::New(m_Nodes[0]->GetName().c_str()));
     m_AnnotationImage->GetExteriorLabel()->SetProperty("name.image", mitk::StringProperty::New("Labels"));
     this->GetDataStorage()->Add(m_AnnotationNode, m_Nodes[0]);
 
     // Convert input images to FeaturePixelType
     for (auto node : m_Nodes)
     {
         mitk::Image::Pointer image = dynamic_cast<mitk::Image*>(node->GetData());
         auto itkImage = itk::Image<FeaturePixelType, 3>::New();
         mitk::CastToItkImage(image, itkImage);
         image = mitk::GrabItkImageMemory(itkImage);
         node->SetData(image);
     }
 
     // Calculate features
     for (const auto node : m_Nodes)
     {
         mitk::Image::Pointer currentImage = dynamic_cast<mitk::Image*>(node->GetData());
         QFuture<std::vector<std::pair<mitk::Image::Pointer, std::string>>> future;
         future = QtConcurrent::run(this, &ActiveLearning::CalculateFeatures, currentImage);
         auto futureWatcher = new QFutureWatcher<std::vector<std::pair<mitk::Image::Pointer, std::string>>>();
         futureWatcher->setFuture(future);
         connect(futureWatcher, SIGNAL(finished()), this, SLOT(OnInitializationFinished()));
         m_FeatureCalculationWatchers.push_back(futureWatcher);
     }
 
     // Interactor
     auto activeLearningLib = us::ModuleRegistry::GetModule("MitkCLActiveLearning");
     m_Interactor = mitk::ActiveLearningInteractor::New();
     m_Interactor->LoadStateMachine("Paint.xml", activeLearningLib);
     m_Interactor->SetEventConfig("PaintConfig.xml", activeLearningLib);
     m_Interactor->SetDataNode(m_AnnotationNode);
 
     // Classifier
     m_Classifier = mitk::VigraRandomForestClassifier::New();
     m_Classifier->SetTreeCount(m_NumberOfTrees);
     m_Classifier->SetMaximumTreeDepth(m_MaximumTreeDepth);
 
     // Automatically add first label
     OnAddLabelPushButtonClicked();
 
     m_Active = true;
 }
 
 
 /* ==================================================================
  * PUBLIC
  * =============================================================== */
 
 
 ActiveLearning::ActiveLearning() :
     m_Parent(nullptr),
     m_AnnotationImage(nullptr),
     m_AnnotationNode(nullptr),
     m_PredictionImage(nullptr),
     m_PredictionNode(nullptr),
     m_SegmentationImage(nullptr),
     m_SegmentationNode(nullptr),
     m_Active(false),
     m_NumberOfTrees(50),
     m_MaximumTreeDepth(10),
     m_PredictionMatrix(nullptr)
 {
 
 }
 
 ActiveLearning::~ActiveLearning()
 {
 
 }
 
 void ActiveLearning::CreateQtPartControl( QWidget *parent )
 {
     m_Controls.setupUi(parent);
     m_Parent = parent;
 
     // Label model
     m_LabelListModel = new QStandardItemModel(0, 3, this);
     m_Controls.m_LabelTableView->setModel(m_LabelListModel);
     m_Controls.m_LabelTableView->horizontalHeader()->setDefaultSectionSize(20);
     m_Controls.m_LabelTableView->verticalHeader()->setDefaultSectionSize(20);
     NotEditableDelegate* itemDelegate = new NotEditableDelegate(parent);
     m_Controls.m_LabelTableView->setItemDelegateForColumn(1, itemDelegate);
     connect(m_Controls.m_LabelTableView, SIGNAL(doubleClicked(QModelIndex)),
             this, SLOT(OnColorIconDoubleClicked(QModelIndex)));
     connect(m_Controls.m_LabelTableView->selectionModel(), SIGNAL(selectionChanged(QItemSelection, QItemSelection)),
             this, SLOT(OnLabelListSelectionChanged(QItemSelection, QItemSelection)));
     connect(m_LabelListModel, SIGNAL(dataChanged(QModelIndex, QModelIndex)),
             this, SLOT(OnLabelNameChanged(QModelIndex, QModelIndex)));
 
     // Buttons
     connect(m_Controls.m_InitializePushButton, SIGNAL(clicked()),
             this, SLOT(Initialize()));
     connect(m_Controls.m_AddLabelPushButton, SIGNAL(clicked()),
             this, SLOT(OnAddLabelPushButtonClicked()));
     connect(m_Controls.m_RemoveLabelPushButton, SIGNAL(clicked()),
             this, SLOT(OnRemoveLabelPushButtonClicked()));
     connect(m_Controls.m_PaintToolButton, SIGNAL(clicked()),
             this, SLOT(OnPaintToolButtonClicked()));
     connect(m_Controls.m_EraseToolButton, SIGNAL(clicked()),
             this, SLOT(OnEraseToolButtonClicked()));
     connect(m_Controls.m_UpdatePredictionsPushButton, SIGNAL(clicked()),
             this, SLOT(OnUpdatePredictionsPushButtonClicked()));
 
     // Set start configuration
     m_Controls.m_LabelControlsFrame->setVisible(false);
     SetInitializeReady(false);
 }
 
 void ActiveLearning::ResetLabels()
 {
     // Remove all labels but the first
     for (auto it=m_AnnotationImage->GetActiveLabelSet()->IteratorBegin(); it!=m_AnnotationImage->GetActiveLabelSet()->IteratorConstEnd(); it++)
     {
         if (it->first != 0)
         {
             m_AnnotationImage->GetActiveLabelSet()->RemoveLabel(it->first);
         }
     }
 
     // Fill with labels from list
     for (int i=0; i<m_LabelListModel->rowCount(); i++)
     {
         QString name = m_LabelListModel->item(i, 2)->text();
         m_LabelListModel->item(i, 1)->setText(QString::number(i + 1));
         QColor color = m_LabelListModel->item(i)->background().color();
         m_AnnotationImage->GetActiveLabelSet()->AddLabel(name.toStdString(), QColorToMitkColor(color));
     }
 }
 
 std::vector<std::pair<mitk::Image::Pointer, std::string> > ActiveLearning::CalculateFeatures(const mitk::Image::Pointer inputImage)
 {
     std::vector<std::pair<mitk::Image::Pointer, std::string>> result;
 
     // TODO: Get features from preference page
     std::vector<float> sigmas = {0.7, 1.6};
 
     for (auto sigma : sigmas)
     {
         std::stringstream ss;
 
         auto gaussImage = mitk::Image::New();
         AccessByItk_n(inputImage, GaussianSmoothing, (sigma, gaussImage));
         ss << "GaussianSmoothing (" << std::fixed << std::setprecision(2) << sigma << ")";
         result.push_back(std::pair<mitk::Image::Pointer, std::string>(gaussImage, ss.str()));
         ss.str("");
 
         auto gradMagImage = mitk::Image::New();
         AccessByItk_n(inputImage, GaussianGradientMagnitude, (sigma, gradMagImage));
         ss << "GaussianGradientMagnitude (" << std::fixed << std::setprecision(2) << sigma << ")";
         result.push_back(std::pair<mitk::Image::Pointer, std::string>(gradMagImage, ss.str()));
         ss.str("");
 
         auto logImage = mitk::Image::New();
         AccessByItk_n(inputImage, LaplacianOfGaussian, (sigma, logImage));
         ss << "LaplacianOfGaussian (" << std::fixed << std::setprecision(2) << sigma << ")";
         result.push_back(std::pair<mitk::Image::Pointer, std::string>(logImage, ss.str()));
         ss.str("");
 
 //        auto structImage1 = mitk::Image::New();
 //        auto structImage2 = mitk::Image::New();
 //        auto structImage3 = mitk::Image::New();
 //        std::vector<mitk::Image::Pointer> structImages = {structImage1, structImage2, structImage3};
 //        AccessByItk_n(inputImage, StructureTensorEigenvalues, (sigma, structImages));
 //        ss << "StructureTensorEigenvalue1 (" << std::fixed << std::setprecision(2) << sigma << ")";
 //        result.push_back(std::pair<mitk::Image::Pointer, std::string>(structImage1, ss.str()));
 //        ss.str("");
 //        ss << "StructureTensorEigenvalue2 (" << std::fixed << std::setprecision(2) << sigma << ")";
 //        result.push_back(std::pair<mitk::Image::Pointer, std::string>(structImage2, ss.str()));
 //        ss.str("");
 //        if (inputImage->GetDimension() == 3)
 //        {
 //            ss << "StructureTensorEigenvalue3 (" << std::fixed << std::setprecision(2) << sigma << ")";
 //            result.push_back(std::pair<mitk::Image::Pointer, std::string>(structImage3, ss.str()));
 //            ss.str("");
 //        }
 
 //        auto hessianImage1 = mitk::Image::New();
 //        auto hessianImage2 = mitk::Image::New();
 //        auto hessianImage3 = mitk::Image::New();
 //        std::vector<mitk::Image::Pointer> hessianImages = {hessianImage1, hessianImage2, hessianImage3};
 //        AccessByItk_n(inputImage, HessianEigenvalues, (sigma, hessianImages));
 //        ss << "HessianEigenvalue1 (" << std::fixed << std::setprecision(2) << sigma << ")";
 //        result.push_back(std::pair<mitk::Image::Pointer, std::string>(hessianImage1, ss.str()));
 //        ss.str("");
 //        ss << "HessianEigenvalue2 (" << std::fixed << std::setprecision(2) << sigma << ")";
 //        result.push_back(std::pair<mitk::Image::Pointer, std::string>(hessianImage2, ss.str()));
 //        ss.str("");
 //        if (inputImage->GetDimension() == 3)
 //        {
 //            ss << "HessianEigenvalue3 (" << std::fixed << std::setprecision(2) << sigma << ")";
 //            result.push_back(std::pair<mitk::Image::Pointer, std::string>(hessianImage3, ss.str()));
 //            ss.str("");
 //        }
     }
 
     return result;
 }
 
 std::pair<std::shared_ptr<Eigen::VectorXi>, std::shared_ptr<Eigen::MatrixXd>> ActiveLearning::GetTrainingData(const mitk::LabelSetImage::Pointer annotationImage, const std::vector<mitk::Image::Pointer> featureImageVector)
 {
     // Get indices and labels
     std::vector<itk::Index<3>> indices;
     std::vector<int> labels;
     itk::Image<mitk::LabelSetImage::PixelType, 3>::Pointer annotationImageItk;
     mitk::CastToItkImage(annotationImage, annotationImageItk);
 
     itk::ImageRegionIteratorWithIndex<itk::Image<mitk::LabelSetImage::PixelType, 3>> it(annotationImageItk, annotationImageItk->GetLargestPossibleRegion());
     while (!it.IsAtEnd())
     {
         if (it.Get() != 0)
         {
             indices.push_back(it.GetIndex());
             labels.push_back(it.Get());
         }
         ++it;
     }
 
     Eigen::MatrixXd trainingData(indices.size(), featureImageVector.size());
     Eigen::VectorXi trainingLabels = Eigen::VectorXi::Map(labels.data(), labels.size());
 
     int j = 0;
     for (mitk::Image::Pointer feature : featureImageVector)
     {
         int i = 0;
         mitk::ImagePixelReadAccessor<FeaturePixelType> access(feature, feature->GetVolumeData());
         for (auto index : indices)
         {
             trainingData(i, j) = access.GetPixelByIndexSafe(index);
             i++;
         }
         j++;
     }
 
     auto trainingLabelsPtr = std::make_shared<Eigen::VectorXi>(trainingLabels);
     auto trainingDataPtr = std::make_shared<Eigen::MatrixXd>(trainingData);
     std::pair<std::shared_ptr<Eigen::VectorXi>, std::shared_ptr<Eigen::MatrixXd>> result = std::make_pair(trainingLabelsPtr, trainingDataPtr);
 
     return result;
 }
 
 const std::string ActiveLearning::VIEW_ID = "org.mitk.views.activelearning";
 
 
 /* ==================================================================
  * PROTECTED SLOTS
  * =============================================================== */
 
 
 void ActiveLearning::OnAddLabelPushButtonClicked()
 {
     QString labelName = QString("Label ") + QString::number(m_AnnotationImage->GetActiveLabelSet()->ReverseIteratorBegin()->first + 1);
     QColor labelColor = Qt::GlobalColor(m_LabelListModel->rowCount() % 12 + 7); // We only want Qt default colors 7 to 18
 
     // Create icon
     QStandardItem* colorSquare = new QStandardItem;
     colorSquare->setBackground(labelColor);
     colorSquare->setEditable(false);
     QPixmap colorPixmap(20, 20);
     colorPixmap.fill(labelColor);
     colorSquare->setIcon(QIcon(colorPixmap));
 
     // Key is the highest existing key + 1
     int value = (int)m_AnnotationImage->GetActiveLabelSet()->ReverseIteratorBegin()->first + 1;
     QStandardItem* valueItem = new QStandardItem();
     valueItem->setText(QString::number(value));
 
     // Create label item
     QStandardItem* label = new QStandardItem(labelName);
 
     // Add to image
     m_AnnotationImage->GetActiveLabelSet()->AddLabel(labelName.toStdString(), QColorToMitkColor(labelColor));
     PrintAllLabels(m_AnnotationImage);
 
     // Make list and insert
     QList<QStandardItem*> list;
     list.append(colorSquare);
     list.append(valueItem);
     list.append(label);
     m_LabelListModel->appendRow(list);
 
     // If this is the first label, we activate the paint button
     // We also have to set the data node color for this one, because for 1 values that color seems to define the rendered color
     if (m_LabelListModel->rowCount() == 1)
     {
         OnPaintToolButtonClicked();
         m_AnnotationNode->SetColor(QColorToMitkColor(labelColor));
     }
 
     // Select newly added label
     m_Controls.m_LabelTableView->selectRow(m_LabelListModel->rowCount() - 1);
 }
 
 void ActiveLearning::OnRemoveLabelPushButtonClicked()
 {
     QItemSelectionModel* selection = m_Controls.m_LabelTableView->selectionModel();
     if (selection->hasSelection())
     {
         unsigned int removeIndex = selection->selectedRows().first().row();
         QString removeMessage = QString("Remove label '") + m_LabelListModel->item(removeIndex, 2)->text() + QString("'?");
         QMessageBox::StandardButton removeReply;
         removeReply = QMessageBox::question(m_Parent,
                                             "Remove Label",
                                             removeMessage,
                                             QMessageBox::Yes | QMessageBox::No);
         if (removeReply == QMessageBox::Yes)
         {
             // if there are no annotations, reset labels
             if (m_Interactor->IsUsed())
             {
                 std::vector<mitk::LabelSetImage::PixelType> labels;
                 labels.push_back(m_LabelListModel->item(removeIndex, 1)->text().toInt());
                 m_AnnotationImage->RemoveLabels(labels);
                 m_LabelListModel->removeRow(removeIndex);
             }
             else
             {
                 m_LabelListModel->removeRow(removeIndex);
                 ResetLabels();
             }
             PrintAllLabels(m_AnnotationImage);
         }
     }
 }
 
 void ActiveLearning::OnPaintToolButtonClicked()
 {
     m_Controls.m_PaintToolButton->setChecked(true);
     QItemSelectionModel* selection = m_Controls.m_LabelTableView->selectionModel();
     int row(0);
     if (selection->hasSelection())
     {
         row = selection->selectedRows().first().row();
     }
     else
     {
         m_Controls.m_LabelTableView->selectRow(0);
     }
     m_Interactor->SetPaintingPixelValue(m_LabelListModel->item(row, 1)->text().toInt());
     m_AnnotationImage->GetActiveLabelSet()->SetActiveLabel(m_LabelListModel->item(row, 1)->text().toInt());
 }
 
 void ActiveLearning::OnEraseToolButtonClicked()
 {
     m_Controls.m_EraseToolButton->setChecked(true);
     m_Interactor->SetPaintingPixelValue(0);
     m_AnnotationImage->GetActiveLabelSet()->SetActiveLabel(0);
 }
 
 void ActiveLearning::OnActivateGuidancePushButtonToggled(bool toggled)
 {
 
 }
 
 void ActiveLearning::OnSaveSegmentationPushButtonClicked()
 {
 
 }
 
 void ActiveLearning::OnSavePredictionsPushButtonClicked()
 {
 
 }
 
 void ActiveLearning::OnExportSegmentationPushButtonClicked()
 {
 
 }
 
 void ActiveLearning::OnExportPredictionsPushButtonClicked()
 {
 
 }
 
 void ActiveLearning::OnColorIconDoubleClicked(const QModelIndex& index)
 {
     // Check if click is really from color icon
     if (index.column() != 0)
     {
         return;
     }
     else
     {
         // Color change dialog
         QColor setColor = QColorDialog::getColor(m_LabelListModel->itemFromIndex(index)->background().color(), m_Parent, "Select Label Color");
         if (setColor.isValid())
         {
             m_LabelListModel->itemFromIndex(index)->setBackground(setColor);
             QPixmap colorPixmap(20, 20);
             colorPixmap.fill(setColor);
             m_LabelListModel->itemFromIndex(index)->setIcon(QIcon(colorPixmap));
 
             // Set color on label
             m_AnnotationImage->GetActiveLabelSet()->GetLabel(m_LabelListModel->item(index.row(), 1)->text().toInt())->SetColor(QColorToMitkColor(setColor));
             m_AnnotationImage->GetActiveLabelSet()->UpdateLookupTable(m_LabelListModel->item(index.row(), 1)->text().toInt());
             PrintAllLabels(m_AnnotationImage);
 
             // If this is the label with value 1 we have to change the data node color
             if (m_LabelListModel->item(index.row(), 1)->text().toInt() == 1)
             {
                 m_AnnotationNode->SetColor(QColorToMitkColor(setColor));
             }
         }
     }
 }
 
 void ActiveLearning::OnLabelListSelectionChanged(const QItemSelection& selected, const QItemSelection& /*deselected*/)
 {
     if (selected.empty()) return;
     // This assumes that only one item can be selected (single selection table view)
     try
     {
         int labelValue = m_LabelListModel->item(selected.indexes()[0].row(), 1)->text().toInt();
         m_Interactor->SetPaintingPixelValue(labelValue);
         m_AnnotationImage->GetActiveLabelSet()->SetActiveLabel(labelValue);
     }
     catch (...)
     {
         m_Interactor->SetPaintingPixelValue(0);
         m_AnnotationImage->GetActiveLabelSet()->SetActiveLabel(0);
     }
 }
 
 void ActiveLearning::OnLabelNameChanged(const QModelIndex& topLeft, const QModelIndex& /*bottomRight*/)
 {
     auto item = m_LabelListModel->itemFromIndex(topLeft);
     if (item->column() != 2) return;
     m_AnnotationImage->GetActiveLabelSet()->GetLabel(m_LabelListModel->item(item->row(), 1)->text().toInt())->SetName(item->text().toStdString());
 }
 
 void ActiveLearning::OnInitializationFinished()
 {
     // Check if all futures are finished
     for (auto watcher : m_FeatureCalculationWatchers)
     {
         if (watcher->isFinished() == false) {return;}
     }
 
     // Empty feature vector
     m_FeatureImageVector.clear();
 
     // Insert features into feature vector and data storage
     for (unsigned int i=0; i<m_FeatureCalculationWatchers.size(); i++)
     {
         auto result = m_FeatureCalculationWatchers[i]->result();
         for (unsigned int j=0; j<result.size(); j++)
         {
             m_FeatureImageVector.push_back(result[j].first);
             auto node = mitk::DataNode::New();
             node->SetData(result[j].first);
             node->SetName(result[j].second);
             node->SetBoolProperty("helper object", true);
             node->SetVisibility(false);
             this->GetDataStorage()->Add(node, m_Nodes[i]);
         }
     }
 
     // Show controls
     m_Controls.m_LabelControlsFrame->setVisible(true);
     m_Controls.m_InitializePushButton->setHidden(true);
     MITK_INFO << "Features: " << m_FeatureImageVector.size();
 
     // Delete watchers
     for (auto watcher : m_FeatureCalculationWatchers)
     {
         delete watcher;
     }
     m_FeatureCalculationWatchers.clear();
 }
 
 void ActiveLearning::OnUpdatePredictionsPushButtonClicked()
 {
     if (m_PredictionMatrix == nullptr)
     {
         auto mat = Transform<FeaturePixelType>(m_FeatureImageVector);
         m_PredictionMatrix = std::make_shared<Eigen::MatrixXd>(mat);
     }
     MITK_INFO << "Made test data";
     auto training = GetTrainingData(m_AnnotationImage, m_FeatureImageVector);
     MITK_INFO << "Got training data";
 
     m_Classifier->Train(*training.second, *training.first);
     MITK_INFO << "Trained";
 
     Eigen::MatrixXi prediction = m_Classifier->Predict(*m_PredictionMatrix);
     MITK_INFO << "Predicted";
 
-    Transform<int, mitk::LabelSetImage::PixelType>(prediction, m_SegmentationImage->GetLayerImage(m_SegmentationImage->GetActiveLayer()));
-    m_SegmentationImage->Modified();
+    mitk::Image::Pointer referenceImage = dynamic_cast<mitk::Image*>(m_Nodes[0]->GetData());
+    mitk::Image::Pointer newImage = Transform<int, mitk::LabelSetImage::PixelType>(prediction, referenceImage);
+    auto newSegmentation = mitk::LabelSetImage::New();
+    newSegmentation->InitializeByLabeledImage(newImage);
+    m_SegmentationImage = newSegmentation;
     m_SegmentationNode->Modified();
 
+    int s = 5;
+    for (int i=0; i<s; ++i)
+    {
+        MITK_INFO << "------" << i << "------";
+        for (int j=0; j<s; ++j)
+        {
+            std::stringstream ss;
+            for (int k=0; k<s; ++k)
+            {
+                ss << m_SegmentationImage->GetLayerImage(m_SegmentationImage->GetActiveLayer())->GetPixelValueByIndex({i, j, k}) << " ";
+            }
+            MITK_INFO << ss.str();
+        }
+    }
+
     this->RequestRenderWindowUpdate();
 }
 
 /* ==================================================================
  * PROTECTED
  * =============================================================== */
 
 
 void ActiveLearning::OnSelectionChanged(berry::IWorkbenchPart::Pointer /*source*/,
                                         const QList<mitk::DataNode::Pointer>& nodes)
 {
     if (!SelectionAllImages(nodes))
     {
         SetInitializeReady(false);
         return;
     }
 
     if (nodes.length() >= 2)
     {
         // First selection is the reference (could be any other)
         mitk::Image::Pointer referenceImage = dynamic_cast<mitk::Image*>(nodes[0]->GetData());
         mitk::BaseGeometry* referenceGeometry = referenceImage->GetTimeGeometry()->GetGeometryForTimeStep(0); // Adjust for multiple timesteps
 
         for (int i=1; i<nodes.length(); ++i)
         {
             mitk::Image::Pointer currentImage = dynamic_cast<mitk::Image*>(nodes[i]->GetData());
             mitk::BaseGeometry* currentGeometry = currentImage->GetTimeGeometry()->GetGeometryForTimeStep(0); // Adjust for multiple timesteps
 
             if (!mitk::Equal(*currentGeometry, *referenceGeometry, mitk::eps, true))
             {
                 SetInitializeReady(false);
                 return;
             }
         }
     }
 
     // All nodes have the same geometry, allow init
     SetInitializeReady(true);
 }
 
 void ActiveLearning::SetFocus()
 {
 }
 
 
 /* ==================================================================
  * PRIVATE
  * =============================================================== */
 
 
 void ActiveLearning::SetInitializeReady(bool ready)
 {
     if (ready)
     {
         // get selection, check again just to be sure
         auto nodes = this->GetDataManagerSelection();
         if (!SelectionAllImages(nodes)) return;
 
         m_Controls.m_InitializePushButton->setEnabled(true);
 
         if (!m_Active)
         {
             QString nameList = QString::fromStdString(nodes[0]->GetName());
             if (nodes.length() >= 2)
             {
                 for (int i=1; i<nodes.length(); ++i)
                 {
                     nameList += QString("<br />");
                     nameList += QString::fromStdString(nodes[i]->GetName());
                 }
             }
             m_Controls.m_InitializeLabel->setText(nameList);
         }
     }
     else
     {
         m_Controls.m_InitializePushButton->setDisabled(true);
         if (!m_Active)
         {
             m_Controls.m_InitializeLabel->setText("<font color='red'>Selected images must have matching geometries</font>");
 
         }
     }
 }