diff --git a/Modules/Classification/CLActiveLearning/src/mitkActiveLearningInteractor.cpp b/Modules/Classification/CLActiveLearning/src/mitkActiveLearningInteractor.cpp
index e85b5a3a87..c74ca8be65 100644
--- a/Modules/Classification/CLActiveLearning/src/mitkActiveLearningInteractor.cpp
+++ b/Modules/Classification/CLActiveLearning/src/mitkActiveLearningInteractor.cpp
@@ -1,549 +1,549 @@
 
 /*===================================================================
 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 <mitkActiveLearningInteractor.h>
 #include <mitkImagePixelWriteAccessor.h>
 #include <mitkImage.h>
 #include <mitkInteractionPositionEvent.h>
 #include <mitkIOUtil.h>
 #include <mitkSegTool2D.h>
 #include <mitkFeedbackContourTool.h>
 #include <mitkExtractSliceFilter.h>
 #include <mitkVtkImageOverwrite.h>
 #include <mitkContourModelToSurfaceFilter.h>
 #include <mitkPlaneGeometry.h>
 #include <vtkPolyData.h>
 #include <vtkPointData.h>
 #include <vtkPolyDataToImageStencil.h>
 #include <vtkImageStencil.h>
 #include <vtkImageLogic.h>
 
 
 // Helper function to get an image from a data node.
 static mitk::Image::Pointer GetImage(mitk::DataNode::Pointer dataNode)
 {
     if (dataNode.IsNull())
         mitkThrow();
 
     mitk::Image::Pointer image = dynamic_cast<mitk::Image*>(dataNode->GetData());
 
     if (image.IsNull())
         mitkThrow();
 
     return image;
 }
 
 // Helper function to get a geometry of an image for a specific time step.
 static mitk::BaseGeometry::Pointer GetGeometry(mitk::Image* image, unsigned int timeStep)
 {
     if (image == nullptr)
         mitkThrow();
 
     mitk::TimeGeometry::Pointer timeGeometry = image->GetTimeGeometry();
 
     if (timeGeometry.IsNull())
         mitkThrow();
 
     auto geometry = timeGeometry->GetGeometryForTimeStep(timeStep);
 
     if (geometry.IsNull())
         mitkThrow();
 
     return geometry;
 }
 
 static double EuclideanDistance(const mitk::Point3D p1, const mitk::Point3D p2)
 {
     return std::sqrt( (p1[0] - p2[0])*(p1[0] - p2[0]) +
                       (p1[1] - p2[1])*(p1[1] - p2[1]) +
                       (p1[2] - p2[2])*(p1[2] - p2[2]) );
 }
 
 static double EuclideanDistance2D(const mitk::Point2D p1, const mitk::Point2D p2)
 {
     return std::sqrt( (p1[0] - p2[0])*(p1[0] - p2[0]) +
                       (p1[1] - p2[1])*(p1[1] - p2[1]) );
 }
 
 static double ScalarProduct(const mitk::Point3D p1, const mitk::Point3D p2)
 {
     return p1[0]*p2[0] + p1[1]*p2[1] + p1[2]*p2[2];
 }
 
 static double ScalarProduct2D(const mitk::Point2D p1, const mitk::Point2D p2)
 {
     return p1[0]*p2[0] + p1[1]*p2[1];
 }
 
 static std::vector<itk::Index<2>> InterpolateIndices2D(itk::Index<2> startIndex, itk::Index<2> endIndex, const mitk::PlaneGeometry* geometry, unsigned int size)
 {
     if (geometry == nullptr)
         mitkThrow();
 
     std::vector<itk::Index<2>> resultIndices;
     mitk::Point2D startPoint;
     mitk::Point2D endPoint;
     startPoint[0] = startIndex[0];
     startPoint[1] = startIndex[1];
     endPoint[0] = endIndex[0];
     endPoint[1] = endIndex[1];
     geometry->IndexToWorld(startPoint, startPoint);
     geometry->IndexToWorld(endPoint, endPoint);
 
     // Distance between end points
     double dist = EuclideanDistance2D(startPoint, endPoint);
 
     // Define region between startIndex and endIndex, padded by (size - 1)
     int regionBounds[4];
     for (int i=0; i<2; ++i)
     {
         regionBounds[2*i] = std::min(startIndex[i], endIndex[i]) - (size - 1);
         regionBounds[2*i+1] = std::max(startIndex[i], endIndex[i]) + (size - 1);
     }
 
     // We only have a threshold given in pixels (size), but image can be spaced in different units.
     // To get the corresponding distances, transform unit vectors and get their lengths.
     // The minimum spacing will be what we compare to.
     double spacings[2];
     spacings[0] = geometry->GetExtentInMM(0) / geometry->GetExtent(0);
     spacings[1] = geometry->GetExtentInMM(1) / geometry->GetExtent(1);
     double minSpacing = (spacings[0] > spacings[1]) ? spacings[0] : spacings[1];
 
     MITK_INFO << "Interpolating between " << startIndex << " and " << endIndex;
     MITK_INFO << "Distance: " << dist;
     MITK_INFO << "Minimum extent: " << minSpacing;
 
     double t = 0;
     double d = 0;
     for (int x = regionBounds[0]; x<=regionBounds[1]; ++x)
     {
         for (int y = regionBounds[2]; y<=regionBounds[3]; ++y)
         {
             mitk::Point2D p;
             mitk::Point2D index;
             index[0] = x;
             index[1] = y;
             geometry->IndexToWorld(index, p);
 
             // if there is not distance between start and end, just get distance to start
             if (dist < mitk::eps)
             {
                 d = EuclideanDistance2D(startPoint, p);
             }
             else
             {
                 t = -1./(dist*dist) * ScalarProduct2D(startPoint - p, endPoint - startPoint);
                 if (t > 0. && t < 1.)
                 {
                     d = std::sqrt( ScalarProduct2D(startPoint - p, startPoint - p) +
                                    2. * t * ScalarProduct2D(endPoint - startPoint, startPoint - p) +
                                    t * t * dist * dist );
                 }
                 else if (t <= 0.)
                 {
                     d = EuclideanDistance2D(startPoint, p);
                 }
                 else
                 {
                     d = EuclideanDistance2D(endPoint, p);
                 }
             }
 
             if (d <= minSpacing * static_cast<double>(size) / 2.)
             {
                 resultIndices.push_back(itk::Index<2>({index[0], index[1]}));
             }
         }
     }
 
     return resultIndices;
 }
 
 static std::vector<itk::Index<3>> InterpolateIndices(itk::Index<3> startIndex, itk::Index<3> endIndex, const mitk::BaseGeometry* geometry, unsigned int size)
 {
     if (geometry == nullptr)
         mitkThrow();
 
     std::vector<itk::Index<3>> resultIndices;
     mitk::Point3D startPoint;
     mitk::Point3D endPoint;
     geometry->IndexToWorld(startIndex, startPoint);
     geometry->IndexToWorld(endIndex, endPoint);
 
     // Distance between end points
     double dist = EuclideanDistance(startPoint, endPoint);
 
     // Define region between startIndex and endIndex, padded by (size - 1)
     int regionBounds[6];
     for (int i=0; i<3; ++i)
     {
         regionBounds[2*i] = std::min(startIndex[i], endIndex[i]) - (size - 1);
         regionBounds[2*i+1] = std::max(startIndex[i], endIndex[i]) + (size - 1);
     }
 
     // We only have a threshold given in pixels (size), but image can be spaced in different units.
     // To get the corresponding distances, transform unit vectors and get their lengths.
     // The minimum spacing will be what we compare to.
     // Could potentially use the spacing to normalize.
     double spacingInIndexSystem[3];
     double minSpacing = -1.;
     for (int i=0; i<3; ++i)
     {
         itk::Index<3> origin;
         origin.Fill(0);
         itk::Index<3> index;
         index.Fill(0);
         index[i] = 1;
         mitk::Point3D p_origin;
         mitk::Point3D p_index;
         geometry->IndexToWorld(origin, p_origin);
         geometry->IndexToWorld(index, p_index);
         double spacing = EuclideanDistance(p_origin, p_index);
         if ( (minSpacing > 0. && spacing < minSpacing) || minSpacing < 0. )
         {
             minSpacing = spacing;
         }
         spacingInIndexSystem[i] = spacing;
     }
 
     // Iterate over all indices in the given region and get distance to the line between startPoint and endPoint.
     // If distance is smaller than size, add to resultIndices.
     //
     // Let (x1,y1,z1) = startPoint, (x2,y2,z2) = endPoint, (x0,y0,z0) = p a point.
     //
     // Line is defined by:
     //     [x1 + (x2-x1) * t]
     // v = [y1 + (y2-y1) * t]
     //     [z1 + (z2-z1) * t]
     //
     // Then (with * dot product):
     // t(p) = - (startPoint - p) * (endPoint - startPoint) / |endPoint - startPoint|^2
     //
     // And (with x cross product):
     // d(p) = |(p - startPoint) x (p - endPoint)| / |endPoint - startPoint|
 
     double t = 0;
     double d = 0;
     for (int x = regionBounds[0]; x<=regionBounds[1]; ++x)
     {
         for (int y = regionBounds[2]; y<=regionBounds[3]; ++y)
         {
             for (int z = regionBounds[4]; z<=regionBounds[5]; ++z)
             {
                 mitk::Point3D p;
                 itk::Index<3> index = {x,y,z};
                 geometry->IndexToWorld(index, p);
 
                 // if there is not distance between start and end, just get distance to start
                 if (dist < mitk::eps)
                 {
                     d = EuclideanDistance(startPoint, p);
                 }
                 else
                 {
                     t = -1./(dist*dist) * ScalarProduct(startPoint - p, endPoint - startPoint);
                     if (t > 0. && t < 1.)
                     {
                         d = std::sqrt( ScalarProduct(startPoint - p, startPoint - p) +
                                        2. * t * ScalarProduct(endPoint - startPoint, startPoint - p) +
                                        t * t * dist * dist );
                     }
                     else if (t <= 0.)
                     {
                         d = EuclideanDistance(startPoint, p);
                     }
                     else
                     {
                         d = EuclideanDistance(endPoint, p);
                     }
                 }
 
                 if (d <= minSpacing)
                 {
                     resultIndices.push_back(index);
                 }
             }
         }
     }
 
     return resultIndices;
 
 // ======= OLD INTERPOLATION ========
 
 //    std::vector<itk::Index<3>> resultIndices;
 //    mitk::Point3D startPoint;
 //    mitk::Point3D endPoint;
 //    geometry->IndexToWorld(startIndex, startPoint);
 //    geometry->IndexToWorld(endIndex, endPoint);
 
 //    itk::Index<3> indexDelta;
 //    int indexDeltaInc[3];
 //    for (int i=0; i<3; i++)
 //    {
 //        indexDelta[i] = endIndex[i] - startIndex[i];
 //        indexDeltaInc[i] = (indexDelta[i] > 0) ? 1 : (indexDelta[i] < 0) ? -1 : 0;
 //    }
 
 //    int argm[3] = {0, 1, 2};
 //    if (abs(indexDelta[1]) > abs(indexDelta[0]))
 //    {
 //        argm[0] = 1;
 //        argm[1] = 0;
 //    }
 //    if (abs(indexDelta[2]) > abs(indexDelta[argm[1]]))
 //    {
 //        argm[2] = argm[1];
 //        argm[1] = 2;
 //    }
 //    if (abs(indexDelta[2]) > abs(indexDelta[argm[0]]))
 //    {
 //        argm[1] = argm[0];
 //        argm[0] = 2;
 //    }
 
 //    double slopes[2];
 //    slopes[0] = (endPoint[argm[1]] - startPoint[argm[1]]) / (endPoint[argm[0]] - startPoint[argm[0]]);
 //    slopes[1] = (endPoint[argm[2]] - startPoint[argm[2]]) / sqrt((endPoint[argm[1]] - startPoint[argm[1]]) * (endPoint[argm[1]] - startPoint[argm[1]]) + (endPoint[argm[0]] - startPoint[argm[0]]) * (endPoint[argm[0]] - startPoint[argm[0]]));
 //    itk::Index<3> currentIndex = startIndex;
 //    mitk::Point3D currentPoint = startPoint;
 
 //    while (currentIndex != endIndex)
 //    {
 //        currentIndex[argm[0]] += indexDeltaInc[argm[0]];
 //        geometry->IndexToWorld(currentIndex, currentPoint);
 //        currentPoint[argm[1]] = startPoint[argm[1]] + slopes[0] * (currentPoint[argm[0]] - startPoint[argm[0]]);
 //        currentPoint[argm[2]] = startPoint[argm[2]] + slopes[1] * sqrt((currentPoint[argm[1]] - startPoint[argm[1]]) * (currentPoint[argm[1]] - startPoint[argm[1]]) + (currentPoint[argm[0]] - startPoint[argm[0]]) * (currentPoint[argm[0]] - startPoint[argm[0]]));
 //        geometry->WorldToIndex(currentPoint, currentIndex);
 //        resultIndices.push_back(currentIndex);
 //    }
 
 //    return resultIndices;
 }
 
 mitk::ActiveLearningInteractor::ActiveLearningInteractor() :
     m_WorkingSlice(nullptr),
     m_WorkingPlane(nullptr),
-    m_Size(5),
+    m_Size(1),
     m_PaintingPixelValue(0)
 {
 }
 
 mitk::ActiveLearningInteractor::~ActiveLearningInteractor()
 {
 }
 
 void mitk::ActiveLearningInteractor::ConnectActionsAndFunctions()
 {
     CONNECT_FUNCTION("paint", Paint)
     CONNECT_FUNCTION("paint_interpolate", PaintInterpolate)
     CONNECT_FUNCTION("set_workingslice", SetWorkingSlice)
     CONNECT_FUNCTION("writeback_workingslice", WriteBackWorkingSlice)
 }
 
 void mitk::ActiveLearningInteractor::DataNodeChanged()
 {
     this->ResetToStartState();
 }
 
 void mitk::ActiveLearningInteractor::SetWorkingSlice(mitk::StateMachineAction* /*action*/,
                                                      mitk::InteractionEvent *event)
 {
     try
     {
         auto renderer = event->GetSender();
         auto image = GetImage(this->GetDataNode());
         auto timeStep = renderer->GetTimeStep();
         auto geometry = GetGeometry(image, timeStep);
         auto planeGeometry = renderer->GetCurrentWorldPlaneGeometry();
 
         // Check current plane geometry
         if (m_WorkingPlane.IsNotNull())
         {
             bool isSameSlice (true);
             isSameSlice &= mitk::MatrixEqualElementWise(planeGeometry->GetIndexToWorldTransform()->GetMatrix(),
                                                         m_WorkingPlane->GetIndexToWorldTransform()->GetMatrix());
             isSameSlice &= mitk::Equal(planeGeometry->GetIndexToWorldTransform()->GetOffset(),
                                        m_WorkingPlane->GetIndexToWorldTransform()->GetOffset());
             if (isSameSlice) return;
         }
         m_WorkingPlane = renderer->GetCurrentWorldPlaneGeometry()->Clone();
 
         // Extract corresponding slice
         vtkSmartPointer<mitkVtkImageOverwrite> reslice = vtkSmartPointer<mitkVtkImageOverwrite>::New();
         reslice->SetOverwriteMode(false);
         reslice->Modified();
         auto extract = mitk::ExtractSliceFilter::New(reslice);
         extract->SetInput(image);
         extract->SetTimeStep(timeStep);
         extract->SetWorldGeometry(m_WorkingPlane);
         extract->SetResliceTransformByGeometry(geometry);
         extract->SetVtkOutputRequest(false);
         extract->Modified();
         extract->Update();
         m_WorkingSlice = extract->GetOutput();
     }
     catch (itk::ExceptionObject& e)
     {
         mitkThrow() << "Could not set working slice, because:";
         mitkThrow() << e.GetDescription();
     }
 }
 
 void mitk::ActiveLearningInteractor::WriteBackWorkingSlice(mitk::StateMachineAction* /*action*/,
                                                            mitk::InteractionEvent *event)
 {
     try
     {
         auto renderer = event->GetSender();
         auto image = GetImage(this->GetDataNode());
         auto timeStep = renderer->GetTimeStep();
         auto geometry = GetGeometry(image, timeStep);
 
         // Write back
         vtkSmartPointer<mitkVtkImageOverwrite> reslice = vtkSmartPointer<mitkVtkImageOverwrite>::New();
         reslice->SetInputSlice(m_WorkingSlice->GetVtkImageData());
         reslice->SetOverwriteMode(true);
         reslice->Modified();
         mitk::ExtractSliceFilter::Pointer insert =  mitk::ExtractSliceFilter::New(reslice);
         insert->SetInput(image);
         insert->SetTimeStep(timeStep);
         insert->SetWorldGeometry(m_WorkingPlane);
         insert->SetResliceTransformByGeometry(geometry);
         insert->Modified();
         insert->Update();
 
         image->Modified();
         this->GetDataNode()->Modified();
         mitk::RenderingManager::GetInstance()->RequestUpdateAll();
     }
     catch (itk::ExceptionObject& e)
     {
         mitkThrow() << "Could not set working slice, because:";
         mitkThrow() << e.GetDescription();
     }
 }
 
 void mitk::ActiveLearningInteractor::Paint(mitk::StateMachineAction* /*action*/,
                                            mitk::InteractionEvent* event)
 {
     if (m_PaintingPixelValue == -1) mitkThrow() << "Cannot paint negative values";
 
     try
     {
         auto renderer = event->GetSender();
         auto image = GetImage(this->GetDataNode());
         auto timeStep = renderer->GetTimeStep();
         auto geometry = GetGeometry(image, timeStep);
         auto positionEvent = dynamic_cast<mitk::InteractionPositionEvent*>(event);
         auto position = positionEvent->GetPositionInWorld();
 
         if (!geometry->IsInside(position)) return;
 
         // Okay, we're safe. Convert the mouse position to the index of the pixel
         // we're pointing at.
         itk::Index<3> index;
         itk::Index<3> oldIndex;
         geometry->WorldToIndex(position, index);
         geometry->WorldToIndex(m_LastPosition, oldIndex);
 
         // We don't need to paint over and over again while moving the mouse
         // pointer inside the same pixel. That's especially relevant when operating
         // on zoomed images.
         if (index != oldIndex)
         {
             // Convert index to slice geometry
             m_WorkingSlice->GetGeometry()->WorldToIndex(position, index);
             itk::Index<2> indexInPlane2D;
             indexInPlane2D[0] = index[0];
             indexInPlane2D[1] = index[1];
 
             // Get indices
             auto indices = InterpolateIndices2D(indexInPlane2D, indexInPlane2D, m_WorkingPlane, m_Size);
 
             // Fill indices
             mitk::ImagePixelWriteAccessor<mitk::ActiveLearningInteractor::AnnotationPixelType, 2> writeAccessor(m_WorkingSlice, m_WorkingSlice->GetSliceData(0));
 
             for (auto i : indices)
             {
                 writeAccessor.SetPixelByIndexSafe(i, m_PaintingPixelValue);
             }
 
             m_LastPosition = position;
             m_Used = true;
         }
     }
     catch (itk::ExceptionObject& e)
     {
         mitkThrow() << "Could not paint, because:";
         mitkThrow() << e.GetDescription();
     }
 }
 
 void mitk::ActiveLearningInteractor::PaintInterpolate(mitk::StateMachineAction* /*action*/, mitk::InteractionEvent* event)
 {
     if (m_PaintingPixelValue == -1) mitkThrow() << "Cannot paint negative values";
 
     try
     {
         auto renderer = event->GetSender();
         auto image = GetImage(this->GetDataNode());
         auto timeStep = renderer->GetTimeStep();
         auto geometry = GetGeometry(image, timeStep);
         auto positionEvent = dynamic_cast<mitk::InteractionPositionEvent*>(event);
         auto position = positionEvent->GetPositionInWorld();
 
         if (!geometry->IsInside(position)) return;
 
         // Okay, we're safe. Convert the mouse position to the index of the pixel
         // we're pointing at.
         itk::Index<3> index;
         itk::Index<3> oldIndex;
         geometry->WorldToIndex(position, index);
         geometry->WorldToIndex(m_LastPosition, oldIndex);
 
         // We don't need to paint over and over again while moving the mouse
         // pointer inside the same pixel. That's especially relevant when operating
         // on zoomed images.
         if (index != oldIndex)
         {
             // Convert index to slice geometry
             m_WorkingSlice->GetGeometry()->WorldToIndex(position, index);
             m_WorkingSlice->GetGeometry()->WorldToIndex(m_LastPosition, oldIndex);
             itk::Index<2> indexInPlane2D;
             itk::Index<2> oldIndexInPlane2D;
             indexInPlane2D[0] = index[0];
             indexInPlane2D[1] = index[1];
             oldIndexInPlane2D[0] = oldIndex[0];
             oldIndexInPlane2D[1] = oldIndex[1];
 
             // Get indices
             auto indices = InterpolateIndices2D(oldIndexInPlane2D, indexInPlane2D, m_WorkingPlane, m_Size);
 
             // Fill indices
             mitk::ImagePixelWriteAccessor<mitk::ActiveLearningInteractor::AnnotationPixelType, 2> writeAccessor(m_WorkingSlice, m_WorkingSlice->GetSliceData(0));
             for (auto i : indices)
             {
                 writeAccessor.SetPixelByIndexSafe(i, m_PaintingPixelValue);
             }
 
             m_LastPosition = position;
             m_Used = true;
         }
     }
     catch (itk::ExceptionObject& e)
     {
         mitkThrow() << "Could not paint with interpolation, because:";
         mitkThrow() << e.GetDescription();
     }
 }
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 f76d0724d3..eeb37c2c86 100644
--- a/Plugins/org.mitk.gui.qt.activelearning/src/internal/QmitkActiveLearning.cpp
+++ b/Plugins/org.mitk.gui.qt.activelearning/src/internal/QmitkActiveLearning.cpp
@@ -1,1229 +1,1279 @@
 /*===================================================================
 
 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>
 #include <mitkLookupTable.h>
 #include <mitkLookupTableProperty.h>
 #include <mitkRenderingModeProperty.h>
 #include <mitkProperties.h>
 #include <mitkNodePredicateProperty.h>
 #include <mitkVtkResliceInterpolationProperty.h>
+#include <mitkPredictionEntropyFilter.h>
 
 // ITK/VTK
 #include <itkDiscreteGaussianImageFilter.h>
 #include <itkGradientMagnitudeRecursiveGaussianImageFilter.h>
 #include <itkLaplacianRecursiveGaussianImageFilter.h>
 #include <itkHessianMatrixEigenvalueImageFilter.h>
 #include <itkStructureTensorEigenvalueImageFilter.h>
 #include <itkHessianRecursiveGaussianImageFilter.h>
 #include <itkSymmetricEigenAnalysisImageFilter.h>
 #include <itkImageRegionIterator.h>
 #include <vtkImageReslice.h>
 #include <itkStructureTensorImageFilter.h>
 
 typedef ActiveLearning::FeaturePixelType FeaturePixelType;
 typedef ActiveLearning::AnnotationPixelType AnnotationPixelType;
 typedef ActiveLearning::LabelPixelType LabelPixelType;
 typedef ActiveLearning::FeatureMatrixType FeatureMatrixType;
 typedef ActiveLearning::LabelVectorType LabelVectorType;
 
 // 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());
 
     for (unsigned int i=0; i<images.size(); ++i)
     {
         typename itk::Image<MatrixElementType, 3>::Pointer imageItk;
         mitk::CastToItkImage(images[i], imageItk);
         outputMatrix.col(i) = Eigen::Matrix<MatrixElementType, Eigen::Dynamic, 1>::Map(imageItk->GetBufferPointer(), size);
     }
 
     return outputMatrix;
 }
 
 template<typename MatrixElementType, typename OutputPixelType>
 static mitk::Image::Pointer Transform(const Eigen::Matrix<MatrixElementType, Eigen::Dynamic, 1> &inputMatrix,
                                       const mitk::Image::Pointer referenceImage)
 {
     typename itk::Image<OutputPixelType, 3>::Pointer imageItk;
     auto outputImage = mitk::Image::New();
     outputImage->Initialize(mitk::MakeScalarPixelType<OutputPixelType>(), *(referenceImage->GetTimeGeometry()->Clone()));
     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;
     }
 
     mitk::GrabItkImageMemory(imageItk, outputImage);
     return outputImage;
 }
 
 template<typename MatrixElementType, typename OutputPixelType>
 static std::vector<mitk::Image::Pointer> Transform(const Eigen::Matrix<MatrixElementType, Eigen::Dynamic, Eigen::Dynamic> &inputMatrix,
                                                    const mitk::Image::Pointer referenceImage)
 {
     std::vector<mitk::Image::Pointer> resultVector;
 
     for (int j=0; j<inputMatrix.cols(); ++j)
     {
         typename itk::Image<OutputPixelType, 3>::Pointer imageItk;
         auto outputImage = mitk::Image::New();
         outputImage->Initialize(mitk::MakeScalarPixelType<OutputPixelType>(), *(referenceImage->GetTimeGeometry()->Clone()));
         mitk::CastToItkImage(outputImage, imageItk);
 
         auto it = itk::ImageRegionIterator<itk::Image<OutputPixelType, 3>>(imageItk, imageItk->GetLargestPossibleRegion());
         int i = 0;
         while (!it.IsAtEnd())
         {
             it.Set(static_cast<OutputPixelType>(inputMatrix(i, j)));
             ++it;
             ++i;
         }
 
         mitk::GrabItkImageMemory(imageItk, outputImage);
         resultVector.push_back(outputImage);
     }
 
     return resultVector;
 }
 
 template<typename T>
 static void PrintMatrix(const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic> dataMatrix,
                         int maxRows = 0)
 {
     if (maxRows == 0 || maxRows > dataMatrix.rows()) maxRows = dataMatrix.rows();
 
     MITK_INFO << "---------------------";
     for (int i=0; i<maxRows; ++i)
     {
         std::stringstream ss;
         for (int j=0; j<dataMatrix.cols(); ++j)
         {
             ss << std::fixed << std::setprecision(2) << dataMatrix(i, j) << " ";
         }
         MITK_INFO << ss.str();
     }
     MITK_INFO << "---------------------";
 }
 
 template<typename T1, typename T2 = T1>
 static void PrintMatrix(const Eigen::Matrix<T1, Eigen::Dynamic, Eigen::Dynamic> dataMatrix,
                         const Eigen::Matrix<T2, Eigen::Dynamic, Eigen::Dynamic> labelMatrix,
                         int maxRows = 0)
 {
     if (labelMatrix.rows() < dataMatrix.rows()) return;
 
     if (maxRows == 0 || maxRows > dataMatrix.rows()) maxRows = dataMatrix.rows();
 
     MITK_INFO << "---------------------";
     for (int i=0; i<maxRows; ++i)
     {
         std::stringstream ss;
         for (int j=0; j<labelMatrix.cols(); ++j)
         {
             ss << labelMatrix(i, j) << " ";
         }
         ss << "- ";
         for (int j=0; j<dataMatrix.cols(); ++j)
         {
             ss << std::fixed << std::setprecision(2) << dataMatrix(i, j) << " ";
         }
         MITK_INFO << ss.str();
     }
     MITK_INFO << "---------------------";
 }
 
 /* ==================================================================
  * 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]);
 ////    }
 //    typedef itk::Image<PixelType, imageDimension> ImageType;
 //    typedef itk::Image<itk::SymmetricSecondRankTensor<PixelType, imageDimension>> TensorImageType;
 //    auto filter = itk::StructureTensorImageFilter<ImageType, TensorImageType>::New();
 //    filter->SetInput(inputImage);
 //    filter->SetNoiseScale(sigma);
 //    filter->SetFeatureScale(sigma);
 //    try
 //    {
 //        filter->Update();
 //    }
 //    catch (...)
 //    {
 //        mitkThrow();
 //    }
 
 //    std::vector<typename ImageType::Pointer> eigenValueImages;
 //    std::vector<itk::ImageRegionIterator<ImageType>> eigenValueImageIterators;
 //    for (unsigned int i=0; i<imageDimension; ++i)
 //    {
 //        auto image = ImageType::New();
 //        image->SetRegions(inputImage->GetLargestPossibleRegion());
 //        image->Allocate();
 //        eigenValueImages.push_back(image);
 //        itk::ImageRegionIterator<ImageType> it(image, image->GetLargestPossibleRegion());
 //        eigenValueImageIterators.push_back(it);
 //    }
 
 //    itk::ImageRegionConstIterator<TensorImageType> tensorImageIterator(filter->GetOutput(), filter->GetOutput()->GetLargestPossibleRegion());
 //    while (!tensorImageIterator.IsAtEnd())
 //    {
 //        typename TensorImageType::PixelType::EigenValuesArrayType ev;
 //        tensorImageIterator.Get().ComputeEigenValues(ev);
 //        for (unsigned int i=0; i<imageDimension; ++i)
 //        {
 //            eigenValueImageIterators[i].Set(ev[i]);
 //            ++eigenValueImageIterators[i];
 //        }
 //        ++tensorImageIterator;
 //    }
 
 //    for (unsigned int i=0; i<imageDimension; ++i)
 //    {
 //        mitk::GrabItkImageMemory(eigenValueImages[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);
     emit SignalSetProgressMaximum(m_Nodes.length() * 6 + 1); // 6 is number of features, shouldn't be hardcoded
     emit SignalResetProgress();
 
     // 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 NODE
     // =======================================
 
 //    m_PredictionNode = mitk::DataNode::New();
 //    m_PredictionNode->SetName("Predictions");
 ////    m_PredictionNode->SetColor(1., 1. ,1.);
 ////    m_PredictionNode->SetBoolProperty("binary", false);
 //    m_PredictionNode->SetProperty("opacity", mitk::FloatProperty::New(0.0f));
 ////    m_PredictionNode->SetBoolProperty("helper object", true);
 //    this->GetDataStorage()->Add(m_PredictionNode, m_Nodes[0]);
 
     // =======================================
     // SEGMENTATION IMAGE
     // =======================================
 
     m_SegmentationImage = mitk::Image::New();
 
     try
     {
         mitk::Image::Pointer referenceImage = dynamic_cast<mitk::Image*>(m_Nodes[0]->GetData());
         m_SegmentationImage->Initialize(mitk::MakeScalarPixelType<AnnotationPixelType>(), *(referenceImage->GetTimeGeometry()->Clone()));
     }
     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(1., 1., 1.);
     m_SegmentationNode->SetBoolProperty("binary", false);
     m_SegmentationNode->SetProperty("reslice interpolation", mitk::VtkResliceInterpolationProperty::New(VTK_RESLICE_NEAREST));
 //    m_PredictionNode->SetBoolProperty("helper object", true);
     this->GetDataStorage()->Add(m_SegmentationNode, m_Nodes[0]);
 
     // =======================================
     // ANNOTATION IMAGE
     // =======================================
 
     m_AnnotationImage = mitk::Image::New();
 
     try
     {
         mitk::Image::Pointer referenceImage = dynamic_cast<mitk::Image*>(m_Nodes[0]->GetData());
         m_AnnotationImage->Initialize(mitk::MakeScalarPixelType<AnnotationPixelType>(), *(referenceImage->GetTimeGeometry()->Clone()));
     }
     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(1., 1., 1.);
     m_AnnotationNode->SetBoolProperty("binary", false);
     m_AnnotationNode->SetProperty("reslice interpolation", mitk::VtkResliceInterpolationProperty::New(VTK_RESLICE_NEAREST));
     m_AnnotationNode->SetProperty("opacity", mitk::FloatProperty::New(1.0f));
 //    m_AnnotationNode->SetBoolProperty("helper object", true);
     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);
     }
 
     emit SignalSetProgress(1);
 
     // 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);
 
     // Automatically add first label
     OnAddLabelPushButtonClicked();
 
     m_Active = true;
 }
 
 
 /* ==================================================================
  * PUBLIC
  * =============================================================== */
 
 
 ActiveLearning::ActiveLearning() :
     m_Parent(nullptr),
     m_AnnotationImage(nullptr),
     m_AnnotationNode(nullptr),
     m_SegmentationImage(nullptr),
     m_SegmentationNode(nullptr),
     m_Active(false),
     m_NumberOfTrees(50),
     m_MaximumTreeDepth(10),
     m_SamplesPerTree(0.66),
     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);
 
     // Connects
     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)));
     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_SaveSegmentationPushButton, SIGNAL(clicked()),
             this, SLOT(OnSaveSegmentationPushButtonClicked()));
     connect(m_Controls.m_SavePredictionsPushButton, SIGNAL(clicked()),
             this, SLOT(OnSavePredictionsPushButtonClicked()));
     connect(m_Controls.m_UpdatePredictionsPushButton, SIGNAL(clicked()),
             this, SLOT(OnUpdatePredictionsPushButtonClicked()));
     connect(this, SIGNAL(SignalIncrementProgress()),
             this, SLOT(OnSignalIncrementProgress()));
     connect(this, SIGNAL(SignalSetProgress(int)),
             this, SLOT(OnSignalSetProgress(int)));
     connect(this, SIGNAL(SignalResetProgress()),
             this, SLOT(OnSignalResetProgress()));
     connect(this, SIGNAL(SignalSetProgressMaximum(int)),
             this, SLOT(OnSignalSetProgressMaximum(int)));
+    connect(m_Controls.m_BrushSizeSlider, SIGNAL(valueChanged(int)),
+            this, SLOT(OnBrushSizeSliderValueChanged(int)));
 
     // Set start configuration
     m_Controls.m_LabelControlsFrame->setVisible(false);
     SetInitializeReady(false);
 }
 
 void ActiveLearning::ResetLabels()
 {
     for (int i=0; i<m_LabelListModel->rowCount(); i++)
     {
         m_LabelListModel->item(i, 1)->setText(QString::number(i + 1));
     }
 }
 
 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));
         gaussImage->SetClonedTimeGeometry(inputImage->GetTimeGeometry());
         ss << "GaussianSmoothing (" << std::fixed << std::setprecision(2) << sigma << ")";
         result.push_back(std::pair<mitk::Image::Pointer, std::string>(gaussImage, ss.str()));
         ss.str("");
         emit SignalIncrementProgress();
 
         auto gradMagImage = mitk::Image::New();
         AccessByItk_n(inputImage, GaussianGradientMagnitude, (sigma, gradMagImage));
         gradMagImage->SetClonedTimeGeometry(inputImage->GetTimeGeometry());
         ss << "GaussianGradientMagnitude (" << std::fixed << std::setprecision(2) << sigma << ")";
         result.push_back(std::pair<mitk::Image::Pointer, std::string>(gradMagImage, ss.str()));
         ss.str("");
         emit SignalIncrementProgress();
 
         auto logImage = mitk::Image::New();
         AccessByItk_n(inputImage, LaplacianOfGaussian, (sigma, logImage));
         logImage->SetClonedTimeGeometry(inputImage->GetTimeGeometry());
         ss << "LaplacianOfGaussian (" << std::fixed << std::setprecision(2) << sigma << ")";
         result.push_back(std::pair<mitk::Image::Pointer, std::string>(logImage, ss.str()));
         ss.str("");
         emit SignalIncrementProgress();
 
 //        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));
 //        structImage1->SetClonedTimeGeometry(inputImage->GetTimeGeometry());
 //        structImage2->SetClonedTimeGeometry(inputImage->GetTimeGeometry());
 //        structImage3->SetClonedTimeGeometry(inputImage->GetTimeGeometry());
 //        ss << "StructureTensorEV1 (" << std::fixed << std::setprecision(2) << sigma << ")";
 //        result.push_back(std::pair<mitk::Image::Pointer, std::string>(structImage1, ss.str()));
 //        ss.str("");
 //        ss << "StructureTensorEV2 (" << 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 << "StructureTensorEV3 (" << std::fixed << std::setprecision(2) << sigma << ")";
 //            result.push_back(std::pair<mitk::Image::Pointer, std::string>(structImage3, ss.str()));
 //            ss.str("");
 //        }
 //        emit SignalIncrementProgress();
 
 //        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<mitk::Image::Pointer, std::vector<mitk::Image::Pointer>> ActiveLearning::CalculatePrediction(const mitk::Image::Pointer annotationImage,
                                                                                                        const std::vector<mitk::Image::Pointer> &featureImageVector,
                                                                                                        const mitk::Image::Pointer referenceImage,
                                                                                                        mitk::AbstractClassifier* classifier,
                                                                                                        std::shared_ptr<FeatureMatrixType> predictionMatrix)
 {
     // Create prediction matrix if necessary
     if (predictionMatrix == nullptr)
     {
         FeatureMatrixType mat = Transform<FeaturePixelType>(featureImageVector);
         predictionMatrix = std::make_shared<FeatureMatrixType>(mat);
     }
     emit SignalIncrementProgress();
 
     // Get training data and train
     auto training = GetTrainingData(annotationImage, featureImageVector);
     classifier->Train(*training.second, *training.first);
     emit SignalIncrementProgress();
 
     // Get result
     LabelVectorType segmentation = classifier->Predict(*predictionMatrix);
     emit SignalIncrementProgress();
     FeatureMatrixType prediction = classifier->GetPointWiseProbabilities();
     mitk::Image::Pointer segmentationImage = Transform<LabelPixelType, AnnotationPixelType>(segmentation, referenceImage);
     std::vector<mitk::Image::Pointer> predictionImages = Transform<FeaturePixelType, FeaturePixelType>(prediction, referenceImage);
     emit SignalIncrementProgress();
 
     std::pair<mitk::Image::Pointer, std::vector<mitk::Image::Pointer>> result = std::make_pair(segmentationImage, predictionImages);
     return result;
 }
 
+mitk::Image::Pointer ActiveLearning::CalculateUncertainty(const std::vector<mitk::Image::Pointer> &probabilityImageVector)
+{
+    typedef itk::Image<FeaturePixelType, 3> ImageType;
+    std::vector<ImageType::Pointer> itkImages;
+    std::vector<itk::ImageRegionConstIterator<ImageType>> iterators;
+
+    for (auto image : probabilityImageVector)
+    {
+        typename ImageType::Pointer itkImage;
+        mitk::CastToItkImage(image, itkImage);
+        itkImages.push_back(itkImage);
+        itk::ImageRegionConstIterator<ImageType> it(itkImage, itkImage->GetLargestPossibleRegion());
+        iterators.push_back(it);
+    }
+
+    auto uncertaintyImage = ImageType::New();
+    uncertaintyImage->SetRegions(itkImages[0]->GetLargestPossibleRegion());
+    uncertaintyImage->Allocate();
+    itk::ImageRegionConstIterator<ImageType> uit(uncertaintyImage, uncertaintyImage->GetLargestPossibleRegion());
+
+    while (!uit.IsAtEnd())
+    {
+        FeaturePixelType value = 0;
+        for (unsigned int i=0; i<iterators.size(); ++i)
+        {
+            auto p = iterators[i].Get();
+            p = std::max(p, 0.000001);
+            value -= p * std::log(p);
+            ++iterators[i];
+        }
+        uit.Set(value);
+        ++uit;
+    }
+
+    auto outputImage = mitk::Image::New();
+    mitk::GrabItkImageMemory(uncertaintyImage, outputImage);
+    return outputImage;
+}
+
 std::pair<std::shared_ptr<LabelVectorType>, std::shared_ptr<FeatureMatrixType>> ActiveLearning::GetTrainingData(const mitk::Image::Pointer annotationImage,
                                                                                                                 const std::vector<mitk::Image::Pointer> &featureImageVector)
 {
     // Get indices and labels
     std::vector<itk::Index<3>> indices;
     std::vector<LabelPixelType> labels;
     itk::Image<LabelPixelType, 3>::Pointer annotationImageItk;
     mitk::CastToItkImage(annotationImage, annotationImageItk);
 
     itk::ImageRegionIteratorWithIndex<itk::Image<LabelPixelType, 3>> it(annotationImageItk, annotationImageItk->GetLargestPossibleRegion());
     while (!it.IsAtEnd())
     {
         if (it.Get() != 0)
         {
             indices.push_back(it.GetIndex());
             labels.push_back(it.Get());
         }
         ++it;
     }
 
     FeatureMatrixType trainingData(indices.size(), featureImageVector.size());
     LabelVectorType trainingLabels = LabelVectorType::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<LabelVectorType>(trainingLabels);
     auto trainingDataPtr = std::make_shared<FeatureMatrixType>(trainingData);
     std::pair<std::shared_ptr<LabelVectorType>, std::shared_ptr<FeatureMatrixType>> result = std::make_pair(trainingLabelsPtr, trainingDataPtr);
 
     return result;
 }
 
 void ActiveLearning::UpdateLookupTables()
 {
     // Create new lookup table from list
     // Annotation type is int, but we only use a ushort lookup table
     auto lut = vtkSmartPointer<vtkLookupTable>::New();
     int lowlim = std::numeric_limits<AnnotationPixelType>::min();
     int uplim = std::numeric_limits<AnnotationPixelType>::max();
     lut->SetNumberOfTableValues(uplim - lowlim + 1);
     lut->SetTableRange(lowlim, uplim);
     for (long i=0; i<(uplim-lowlim+1); ++i)
     {
         lut->SetTableValue(i, 0.0, 0.0, 0.0, 0.0);
     }
     for (int j=0; j<m_LabelListModel->rowCount(); ++j)
     {
         int value = m_LabelListModel->item(j, 1)->text().toInt();
         const QColor color = m_LabelListModel->item(j, 0)->background().color();
         lut->SetTableValue(value, color.redF(), color.greenF(), color.blueF(), 1.0);
     }
 
     auto lutMitk = mitk::LookupTable::New();
     lutMitk->SetVtkLookupTable(lut);
 
     // Set to annotation image and segmentation image
     auto * lut_prop = dynamic_cast<mitk::LookupTableProperty *>(m_AnnotationNode->GetProperty("LookupTable"));
     lut_prop->SetLookupTable(lutMitk);
     m_AnnotationNode->SetProperty("Image Rendering.Mode", mitk::RenderingModeProperty::New(mitk::RenderingModeProperty::LOOKUPTABLE_COLOR));
     m_AnnotationNode->Modified();
     lut_prop = dynamic_cast<mitk::LookupTableProperty *>(m_SegmentationNode->GetProperty("LookupTable"));
     lut_prop->SetLookupTable(lutMitk);
     m_SegmentationNode->SetProperty("Image Rendering.Mode", mitk::RenderingModeProperty::New(mitk::RenderingModeProperty::LOOKUPTABLE_COLOR));
     m_SegmentationNode->Modified();
 }
 
 void ActiveLearning::UpdatePredictionNodes()
 {
     if (m_LabelListModel->rowCount() == 0) return;
 
     // Build lookup table
     vtkSmartPointer<vtkLookupTable> lut = vtkSmartPointer<vtkLookupTable>::New();
     lut->SetTableRange (0, 1);
     lut->SetSaturationRange (0, 0);
     lut->SetHueRange (0, 0);
     lut->SetValueRange (0, 1);
     lut->SetAlphaRange (0, 1);
     lut->Build();
     auto lutMitk = mitk::LookupTable::New();
     lutMitk->SetVtkLookupTable(lut);
 
     for (int i=0; i<m_LabelListModel->rowCount(); ++i)
     {
         auto property = mitk::NodePredicateProperty::New("segmentation_value", mitk::IntProperty::New(m_LabelListModel->item(i, 1)->text().toInt()));
         auto nodes = this->GetDataStorage()->GetDerivations(m_Nodes[0], property);
         if (nodes->Size() == 1)
         {
             auto node = nodes->GetElement(0);
             QString name = "Prediction ";
             name += m_LabelListModel->item(i, 2)->text();
             node->SetName(name.toStdString());
             auto * lut_prop = dynamic_cast<mitk::LookupTableProperty *>(node->GetProperty("LookupTable"));
             lut_prop->SetLookupTable(lutMitk);
             node->SetProperty("Image Rendering.Mode", mitk::RenderingModeProperty::New(mitk::RenderingModeProperty::LOOKUPTABLE_COLOR));
             node->SetColor(QColorToMitkColor(m_LabelListModel->item(i, 0)->background().color()));
             node->Modified();
         }
         if (nodes->Size() > 1) mitkThrow();
     }
 
     mitk::RenderingManager::GetInstance()->RequestUpdateAll();
 }
 
 const std::string ActiveLearning::VIEW_ID = "org.mitk.views.activelearning";
 
 
 /* ==================================================================
  * PROTECTED SLOTS
  * =============================================================== */
 
 
 void ActiveLearning::OnSignalIncrementProgress()
 {
     m_Controls.m_ProgressBar->setValue(m_Controls.m_ProgressBar->value() + 1);
 }
 
 void ActiveLearning::OnSignalSetProgress(int value)
 {
     m_Controls.m_ProgressBar->setValue(value);
 }
 
 void ActiveLearning::OnSignalResetProgress()
 {
     m_Controls.m_ProgressBar->setValue(0);
 }
 
 void ActiveLearning::OnSignalSetProgressMaximum(int value)
 {
     m_Controls.m_ProgressBar->setMaximum(value);
 }
 
 void ActiveLearning::OnAddLabelPushButtonClicked()
 {
     QString labelName = QString("Label ") + QString::number(m_LabelListModel->rowCount() + 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 = 1;
     if (m_LabelListModel->rowCount() >= 1)
     {
         value = m_LabelListModel->item(m_LabelListModel->rowCount() - 1, 1)->text().toInt() + 1;
     }
     QStandardItem* valueItem = new QStandardItem;
     valueItem->setText(QString::number(value));
 
     // Create label item
     QStandardItem* label = new QStandardItem(labelName);
 
     // Make list and insert
     QList<QStandardItem*> list;
     list.append(colorSquare);
     list.append(valueItem);
     list.append(label);
     m_LabelListModel->appendRow(list);
     m_Controls.m_LabelTableView->selectRow(m_LabelListModel->rowCount() - 1);
 
     // 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();
     }
 
     // Update colors
     UpdateLookupTables();
 }
 
 void ActiveLearning::OnRemoveLabelPushButtonClicked()
 {
     // can't remove last label
     if (m_LabelListModel->rowCount() <= 1) return;
 
     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)
         {
             AnnotationPixelType removeValue = m_LabelListModel->item(removeIndex, 1)->text().toInt();
             m_LabelListModel->removeRow(removeIndex);
             if (!m_Interactor->IsUsed())
             {
                 ResetLabels();
             }
             else
             {
                 itk::Image<AnnotationPixelType, 3>::Pointer imageItk;
                 mitk::CastToItkImage(m_AnnotationImage, imageItk);
                 auto it = itk::ImageRegionIterator<itk::Image<AnnotationPixelType, 3>>(imageItk, imageItk->GetLargestPossibleRegion());
                 while (!it.IsAtEnd())
                 {
                     if (it.Get() == removeValue)
                         it.Set(0);
                     ++it;
                 }
                 mitk::GrabItkImageMemory(imageItk, m_AnnotationImage);
                 UpdateLookupTables();
                 mitk::RenderingManager::GetInstance()->RequestUpdateAll();
             }
         }
     }
 }
 
 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());
 }
 
 void ActiveLearning::OnEraseToolButtonClicked()
 {
     m_Controls.m_EraseToolButton->setChecked(true);
     m_Interactor->SetPaintingPixelValue(0);
 }
 
 void ActiveLearning::OnSaveSegmentationPushButtonClicked()
 {
     auto newNode = mitk::DataNode::New();
     newNode->SetName("Segmentation");
     newNode->SetBoolProperty("binary", false);
     newNode->SetOpacity(1.0);
     newNode->SetVisibility(false);
     newNode->SetData(m_SegmentationImage->Clone());
     this->GetDataStorage()->Add(newNode);
 }
 
 void ActiveLearning::OnSavePredictionsPushButtonClicked()
 {
     if (m_LabelListModel->rowCount() < 1) return;
 
     for (int i=0; i<m_LabelListModel->rowCount(); ++i)
     {
         auto property = mitk::NodePredicateProperty::New("segmentation_value", mitk::IntProperty::New(m_LabelListModel->item(i, 1)->text().toInt()));
         auto nodes = this->GetDataStorage()->GetDerivations(m_Nodes[0], property);
         if (nodes->Size() == 1)
         {
             auto sourceNode = nodes->GetElement(0);
             mitk::Image::Pointer sourceImage = dynamic_cast<mitk::Image*>(sourceNode->GetData());
             auto newNode = mitk::DataNode::New();
             QString name = "Prediction ";
             name += m_LabelListModel->item(i, 2)->text();
             newNode->SetName(name.toStdString());
             newNode->SetBoolProperty("binary", false);
             newNode->SetOpacity(1.0);
             newNode->SetVisibility(false);
             newNode->SetProperty("segmentation_value", mitk::IntProperty::New(m_LabelListModel->item(i, 1)->text().toInt()));
             newNode->SetData(sourceImage->Clone());
             this->GetDataStorage()->Add(newNode);
         }
         if (nodes->Size() > 1) mitkThrow();
     }
 }
 
 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));
             UpdateLookupTables();
             UpdatePredictionNodes();
         }
     }
 }
 
 void ActiveLearning::OnLabelListSelectionChanged(const QItemSelection& selected,
                                                  const QItemSelection& /*deselected*/)
 {
     if (selected.empty()) return;
     if (m_Controls.m_EraseToolButton->isChecked()) 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);
     }
     catch (...)
     {
         m_Interactor->SetPaintingPixelValue(-1);
     }
 }
 
 void ActiveLearning::OnLabelNameChanged(const QModelIndex& topLeft, const QModelIndex& /*bottomRight*/)
 {
     UpdatePredictionNodes();
 }
 
 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);
     emit SignalResetProgress();
 
     // Delete watchers
     for (auto watcher : m_FeatureCalculationWatchers)
     {
         delete watcher;
     }
     m_FeatureCalculationWatchers.clear();
 }
 
 void ActiveLearning::OnUpdatePredictionsPushButtonClicked()
 {
     if (m_LabelListModel->rowCount() < 1) return;
 
     m_Controls.m_UpdatePredictionsPushButton->setDisabled(true);
     emit SignalSetProgressMaximum(4);
 
     // Clear old predictions
     for (int i=0; i<m_LabelListModel->rowCount(); ++i)
     {
         auto property = mitk::NodePredicateProperty::New("segmentation_value", mitk::IntProperty::New(m_LabelListModel->item(i, 1)->text().toInt()));
         auto nodes = this->GetDataStorage()->GetDerivations(m_Nodes[0], property);
         if (nodes->Size() == 1)
         {
             this->GetDataStorage()->Remove(nodes);
         }
         if (nodes->Size() > 1) mitkThrow();
     }
 
     // Classifier
     auto classifier = mitk::VigraRandomForestClassifier::New();
     classifier->SetTreeCount(m_NumberOfTrees);
     classifier->SetMaximumTreeDepth(m_MaximumTreeDepth);
     classifier->SetSamplesPerTree(m_SamplesPerTree);
 
     mitk::Image::Pointer referenceImage = dynamic_cast<mitk::Image*>(m_Nodes[0]->GetData());
     QFuture<std::pair<mitk::Image::Pointer, std::vector<mitk::Image::Pointer>>> future;
     future = QtConcurrent::run(this, &ActiveLearning::CalculatePrediction, m_AnnotationImage, m_FeatureImageVector, referenceImage, classifier, m_PredictionMatrix);
     m_PredictionCalculationWatcher = new QFutureWatcher<std::pair<mitk::Image::Pointer, std::vector<mitk::Image::Pointer>>>();
     m_PredictionCalculationWatcher->setFuture(future);
 
     connect(m_PredictionCalculationWatcher, SIGNAL(finished()),
             this, SLOT(OnPredictionCalculationFinished()));
 }
 
 void ActiveLearning::OnPredictionCalculationFinished()
 {
     auto result = m_PredictionCalculationWatcher->result();
 
     m_SegmentationImage = result.first;
     m_SegmentationImage->Modified();
     m_SegmentationNode->SetData(m_SegmentationImage);
     m_SegmentationNode->Modified();
 
     for (unsigned int i=0; i<result.second.size(); ++i)
     {
         auto node = mitk::DataNode::New();
         QString name = "Prediction ";
         name += m_LabelListModel->item(i, 2)->text();
         node->SetName(name.toStdString());
         node->SetBoolProperty("binary", false);
         node->SetVisibility(false);
         node->SetOpacity(0.3);
         node->SetColor(QColorToMitkColor(m_LabelListModel->item(i, 0)->background().color()));
         node->SetProperty("segmentation_value", mitk::IntProperty::New(m_LabelListModel->item(i, 1)->text().toInt()));
         node->SetData(result.second[i]);
         this->GetDataStorage()->Add(node, m_Nodes[0]);
     }
 
     UpdateLookupTables();
     UpdatePredictionNodes();
     emit SignalResetProgress();
     m_Controls.m_UpdatePredictionsPushButton->setEnabled(true);
 }
 
+void ActiveLearning::OnBrushSizeSliderValueChanged(int value)
+{
+    QString labelText = "Size ";
+    labelText += QString::number(value);
+    m_Controls.m_BrushSizeLabel->setText(labelText);
+    m_Interactor->SetSize(value);
+}
+
 /* ==================================================================
  * 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>");
 
         }
     }
 }
diff --git a/Plugins/org.mitk.gui.qt.activelearning/src/internal/QmitkActiveLearning.h b/Plugins/org.mitk.gui.qt.activelearning/src/internal/QmitkActiveLearning.h
index 041c21f715..f9fbf6828b 100644
--- a/Plugins/org.mitk.gui.qt.activelearning/src/internal/QmitkActiveLearning.h
+++ b/Plugins/org.mitk.gui.qt.activelearning/src/internal/QmitkActiveLearning.h
@@ -1,186 +1,190 @@
 /*===================================================================
 
 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 ActiveLearning_h
 #define ActiveLearning_h
 
 #include <berryISelectionListener.h>
 #include <QmitkAbstractView.h>
 #include "ui_QmitkActiveLearningControls.h"
 
 // Qt
 #include <QItemDelegate>
 #include <QStandardItemModel>
 #include <QWidget>
 #include <QFutureWatcher>
 
 // MITK
 #include <mitkLabelSetImage.h>
 #include <mitkActiveLearningInteractor.h>
 #include <mitkVigraRandomForestClassifier.h>
 
 #include <Eigen/Dense>
 
 /**
 \brief ActiveLearning
 
 \warning  This class is not yet documented. Use "git blame" and ask the author to provide basic documentation.
 
 \sa QmitkAbstractView
 \ingroup ${plugin_target}_internal
 */
 
 
 
 // Just a helper class
 class NotEditableDelegate : public QItemDelegate
 {
     Q_OBJECT
 public:
     explicit NotEditableDelegate(QObject* parent = nullptr) : QItemDelegate(parent) {}
 protected:
     QWidget* createEditor(QWidget*, const QStyleOptionViewItem&, const QModelIndex&) const {return Q_NULLPTR;}
 };
 
 
 
 class ActiveLearning : public QmitkAbstractView
 {
 
     Q_OBJECT
 
 public slots:
 
     void Initialize();
 
 public:
 
     typedef unsigned short AnnotationPixelType;
     typedef double FeaturePixelType;
     typedef int LabelPixelType;
     typedef Eigen::Matrix<FeaturePixelType, Eigen::Dynamic, Eigen::Dynamic> FeatureMatrixType;
     typedef Eigen::Matrix<LabelPixelType, Eigen::Dynamic, 1> LabelVectorType;
 
     ActiveLearning();
     ~ActiveLearning();
 
     void CreateQtPartControl(QWidget *parent) override;
 
     std::vector<std::pair<mitk::Image::Pointer, std::string>> CalculateFeatures(const mitk::Image::Pointer inputImage);
 
     std::pair<mitk::Image::Pointer, std::vector<mitk::Image::Pointer> > CalculatePrediction(const mitk::Image::Pointer annotationImage,
                                                                                             const std::vector<mitk::Image::Pointer> &featureImageVector,
                                                                                             const mitk::Image::Pointer referenceImage,
                                                                                             mitk::AbstractClassifier *classifier,
                                                                                             std::shared_ptr<FeatureMatrixType> predictionMatrix);
 
     std::pair<std::shared_ptr<LabelVectorType>, std::shared_ptr<FeatureMatrixType> > GetTrainingData(const mitk::Image::Pointer annotationImage,
                                                                                                      const std::vector<mitk::Image::Pointer> &featureImageVector);
 
+    mitk::Image::Pointer CalculateUncertainty(const std::vector<mitk::Image::Pointer> &probabilityImageVector);
+
     void UpdateLookupTables();
 
     void UpdatePredictionNodes();
 
     static const std::string VIEW_ID;
 
 signals:
 
     void SignalIncrementProgress();
 
     void SignalSetProgress(int value);
 
     void SignalResetProgress();
 
     void SignalSetProgressMaximum(int value);
 
 protected slots:
 
     void OnSignalIncrementProgress();
 
     void OnSignalSetProgress(int value);
 
     void OnSignalResetProgress();
 
     void OnSignalSetProgressMaximum(int value);
 
     void OnAddLabelPushButtonClicked();
 
     void OnRemoveLabelPushButtonClicked();
 
     void OnPaintToolButtonClicked();
 
     void OnEraseToolButtonClicked();
 
     void OnSaveSegmentationPushButtonClicked();
 
     void OnSavePredictionsPushButtonClicked();
 
     void OnColorIconDoubleClicked(const QModelIndex& index);
 
     void OnLabelListSelectionChanged(const QItemSelection& selected, const QItemSelection& /*deselected*/);
 
     void OnLabelNameChanged(const QModelIndex& topLeft, const QModelIndex& /*bottomRight*/);
 
+    void OnBrushSizeSliderValueChanged(int value);
+
     void OnUpdatePredictionsPushButtonClicked();
 
     void OnPredictionCalculationFinished();
 
     void OnInitializationFinished();
 
 protected:
 
     void OnSelectionChanged(berry::IWorkbenchPart::Pointer /*source*/,
                             const QList<mitk::DataNode::Pointer>& nodes) override;
 
     void SetFocus() override;
 
     void ResetLabels();
 
     void SetInitializeReady(bool ready);
 
     Ui::ActiveLearningControls m_Controls;
     QWidget* m_Parent;
 
     mitk::Image::Pointer m_AnnotationImage;
     mitk::DataNode::Pointer m_AnnotationNode;
 
     mitk::Image::Pointer m_SegmentationImage;
     mitk::DataNode::Pointer m_SegmentationNode;
 
     std::vector<mitk::Image::Pointer> m_FeatureImageVector;
     std::vector<QFutureWatcher<std::vector<std::pair<mitk::Image::Pointer, std::string>>>*> m_FeatureCalculationWatchers;
 
     QFutureWatcher<std::pair<mitk::Image::Pointer, std::vector<mitk::Image::Pointer>>>* m_PredictionCalculationWatcher;
 
     QStandardItemModel* m_LabelListModel;
 
     mitk::ActiveLearningInteractor::Pointer m_Interactor;
 
     QList<mitk::DataNode::Pointer> m_Nodes;
 
 private:
 
     bool m_Active;
 
     int m_NumberOfTrees;
     int m_MaximumTreeDepth;
     float m_SamplesPerTree;
 
     std::shared_ptr<FeatureMatrixType> m_PredictionMatrix;
 
 };
 
 #endif // ActiveLearning_h
diff --git a/Plugins/org.mitk.gui.qt.activelearning/src/internal/QmitkActiveLearningControls.ui b/Plugins/org.mitk.gui.qt.activelearning/src/internal/QmitkActiveLearningControls.ui
index 23ba48a271..38308ae44a 100644
--- a/Plugins/org.mitk.gui.qt.activelearning/src/internal/QmitkActiveLearningControls.ui
+++ b/Plugins/org.mitk.gui.qt.activelearning/src/internal/QmitkActiveLearningControls.ui
@@ -1,491 +1,547 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <ui version="4.0">
  <class>ActiveLearningControls</class>
  <widget class="QWidget" name="ActiveLearningControls">
   <property name="geometry">
    <rect>
     <x>0</x>
     <y>0</y>
     <width>352</width>
     <height>581</height>
    </rect>
   </property>
   <property name="minimumSize">
    <size>
     <width>0</width>
     <height>0</height>
    </size>
   </property>
   <property name="windowTitle">
    <string>QmitkTemplate</string>
   </property>
   <layout class="QVBoxLayout" name="verticalLayout">
    <item>
     <widget class="QPushButton" name="m_InitializePushButton">
      <property name="text">
       <string>Initialize with selection</string>
      </property>
      <property name="checkable">
       <bool>false</bool>
      </property>
      <property name="autoDefault">
       <bool>false</bool>
      </property>
      <property name="default">
       <bool>false</bool>
      </property>
      <property name="flat">
       <bool>false</bool>
      </property>
     </widget>
    </item>
    <item>
     <widget class="QLabel" name="m_InitializeLabel">
      <property name="text">
       <string>&lt;font color=&quot;red&quot;&gt;Selection must have matching dimensions and transforms&lt;/font&gt;</string>
      </property>
      <property name="wordWrap">
       <bool>true</bool>
      </property>
     </widget>
    </item>
    <item>
     <widget class="QFrame" name="m_LabelControlsFrame">
      <property name="frameShape">
       <enum>QFrame::NoFrame</enum>
      </property>
      <property name="frameShadow">
       <enum>QFrame::Raised</enum>
      </property>
      <layout class="QVBoxLayout" name="verticalLayout_2">
       <property name="leftMargin">
        <number>0</number>
       </property>
       <property name="topMargin">
        <number>0</number>
       </property>
       <property name="rightMargin">
        <number>0</number>
       </property>
       <property name="bottomMargin">
        <number>0</number>
       </property>
       <item>
        <spacer name="verticalSpacer_2">
         <property name="orientation">
          <enum>Qt::Vertical</enum>
         </property>
         <property name="sizeType">
          <enum>QSizePolicy::Fixed</enum>
         </property>
         <property name="sizeHint" stdset="0">
          <size>
           <width>20</width>
           <height>10</height>
          </size>
         </property>
        </spacer>
       </item>
       <item>
        <widget class="Line" name="m_UpperSeparator">
         <property name="frameShadow">
          <enum>QFrame::Sunken</enum>
         </property>
         <property name="lineWidth">
          <number>1</number>
         </property>
         <property name="orientation">
          <enum>Qt::Horizontal</enum>
         </property>
        </widget>
       </item>
       <item>
        <spacer name="verticalSpacer_3">
         <property name="orientation">
          <enum>Qt::Vertical</enum>
         </property>
         <property name="sizeType">
          <enum>QSizePolicy::Fixed</enum>
         </property>
         <property name="sizeHint" stdset="0">
          <size>
           <width>20</width>
           <height>10</height>
          </size>
         </property>
        </spacer>
       </item>
       <item>
        <widget class="QFrame" name="m_LabelControlsAddRemoveFrame">
         <property name="frameShape">
          <enum>QFrame::NoFrame</enum>
         </property>
         <property name="frameShadow">
          <enum>QFrame::Raised</enum>
         </property>
         <layout class="QHBoxLayout" name="horizontalLayout">
          <property name="leftMargin">
           <number>0</number>
          </property>
          <property name="topMargin">
           <number>0</number>
          </property>
          <property name="rightMargin">
           <number>0</number>
          </property>
          <property name="bottomMargin">
           <number>0</number>
          </property>
          <item>
           <widget class="QPushButton" name="m_AddLabelPushButton">
            <property name="text">
             <string> Add Label </string>
            </property>
           </widget>
          </item>
          <item>
           <widget class="QPushButton" name="m_RemoveLabelPushButton">
            <property name="text">
             <string> Remove Label </string>
            </property>
           </widget>
          </item>
          <item>
           <spacer name="horizontalSpacer">
            <property name="orientation">
             <enum>Qt::Horizontal</enum>
            </property>
            <property name="sizeHint" stdset="0">
             <size>
              <width>40</width>
              <height>20</height>
             </size>
            </property>
           </spacer>
          </item>
         </layout>
        </widget>
       </item>
       <item>
        <widget class="QTableView" name="m_LabelTableView">
         <property name="frameShadow">
          <enum>QFrame::Sunken</enum>
         </property>
         <property name="alternatingRowColors">
          <bool>false</bool>
         </property>
         <property name="selectionMode">
          <enum>QAbstractItemView::SingleSelection</enum>
         </property>
         <property name="selectionBehavior">
          <enum>QAbstractItemView::SelectRows</enum>
         </property>
         <attribute name="horizontalHeaderVisible">
          <bool>false</bool>
         </attribute>
         <attribute name="horizontalHeaderStretchLastSection">
          <bool>true</bool>
         </attribute>
         <attribute name="verticalHeaderVisible">
          <bool>false</bool>
         </attribute>
        </widget>
       </item>
       <item>
        <widget class="QFrame" name="m_LabelControlsPaintEraseFrame">
         <property name="frameShape">
          <enum>QFrame::NoFrame</enum>
         </property>
         <property name="frameShadow">
          <enum>QFrame::Raised</enum>
         </property>
         <layout class="QHBoxLayout" name="horizontalLayout_2">
          <property name="leftMargin">
           <number>0</number>
          </property>
          <property name="topMargin">
           <number>0</number>
          </property>
          <property name="rightMargin">
           <number>0</number>
          </property>
          <property name="bottomMargin">
           <number>0</number>
          </property>
          <item>
           <widget class="QToolButton" name="m_PaintToolButton">
            <property name="sizePolicy">
             <sizepolicy hsizetype="Expanding" vsizetype="Fixed">
              <horstretch>0</horstretch>
              <verstretch>0</verstretch>
             </sizepolicy>
            </property>
            <property name="minimumSize">
             <size>
              <width>0</width>
              <height>40</height>
             </size>
            </property>
            <property name="text">
             <string/>
            </property>
            <property name="icon">
             <iconset resource="../../../../../MITK-al-superbuild/MITK-build/Plugins/org.mitk.gui.qt.activelearning/org_mitk_gui_qt_activelearning_cached.qrc">
              <normaloff>:/org.mitk.gui.qt.activelearning/resources/paint_icon_200.png</normaloff>:/org.mitk.gui.qt.activelearning/resources/paint_icon_200.png</iconset>
            </property>
            <property name="iconSize">
             <size>
              <width>25</width>
              <height>25</height>
             </size>
            </property>
            <property name="checkable">
             <bool>true</bool>
            </property>
            <property name="autoExclusive">
             <bool>true</bool>
            </property>
           </widget>
          </item>
          <item>
           <widget class="QToolButton" name="m_EraseToolButton">
            <property name="sizePolicy">
             <sizepolicy hsizetype="Expanding" vsizetype="Fixed">
              <horstretch>0</horstretch>
              <verstretch>0</verstretch>
             </sizepolicy>
            </property>
            <property name="minimumSize">
             <size>
              <width>0</width>
              <height>40</height>
             </size>
            </property>
            <property name="text">
             <string>...</string>
            </property>
            <property name="icon">
             <iconset resource="../../../../../MITK-al-superbuild/MITK-build/Plugins/org.mitk.gui.qt.activelearning/org_mitk_gui_qt_activelearning_cached.qrc">
              <normaloff>:/org.mitk.gui.qt.activelearning/resources/erase_icon_100.png</normaloff>:/org.mitk.gui.qt.activelearning/resources/erase_icon_100.png</iconset>
            </property>
            <property name="iconSize">
             <size>
              <width>25</width>
              <height>25</height>
             </size>
            </property>
            <property name="checkable">
             <bool>true</bool>
            </property>
            <property name="autoExclusive">
             <bool>true</bool>
            </property>
           </widget>
          </item>
         </layout>
        </widget>
       </item>
+      <item>
+       <widget class="QFrame" name="m_BrushSizeFrame">
+        <property name="frameShape">
+         <enum>QFrame::NoFrame</enum>
+        </property>
+        <property name="frameShadow">
+         <enum>QFrame::Raised</enum>
+        </property>
+        <layout class="QHBoxLayout" name="horizontalLayout_5">
+         <property name="leftMargin">
+          <number>0</number>
+         </property>
+         <property name="topMargin">
+          <number>0</number>
+         </property>
+         <property name="rightMargin">
+          <number>0</number>
+         </property>
+         <property name="bottomMargin">
+          <number>0</number>
+         </property>
+         <item>
+          <widget class="QLabel" name="m_BrushSizeLabel">
+           <property name="sizePolicy">
+            <sizepolicy hsizetype="Preferred" vsizetype="Preferred">
+             <horstretch>0</horstretch>
+             <verstretch>0</verstretch>
+            </sizepolicy>
+           </property>
+           <property name="minimumSize">
+            <size>
+             <width>52</width>
+             <height>0</height>
+            </size>
+           </property>
+           <property name="text">
+            <string>Size 1</string>
+           </property>
+          </widget>
+         </item>
+         <item>
+          <widget class="QSlider" name="m_BrushSizeSlider">
+           <property name="minimum">
+            <number>1</number>
+           </property>
+           <property name="maximum">
+            <number>20</number>
+           </property>
+           <property name="orientation">
+            <enum>Qt::Horizontal</enum>
+           </property>
+          </widget>
+         </item>
+        </layout>
+       </widget>
+      </item>
       <item>
        <widget class="QCheckBox" name="m_ActiveGuidanceCheckbox">
         <property name="text">
          <string>Active Guidance</string>
         </property>
        </widget>
       </item>
       <item>
        <widget class="QPushButton" name="m_UpdatePredictionsPushButton">
         <property name="minimumSize">
          <size>
           <width>0</width>
           <height>40</height>
          </size>
         </property>
         <property name="text">
          <string>Update</string>
         </property>
        </widget>
       </item>
       <item>
        <spacer name="verticalSpacer_4">
         <property name="orientation">
          <enum>Qt::Vertical</enum>
         </property>
         <property name="sizeType">
          <enum>QSizePolicy::Fixed</enum>
         </property>
         <property name="sizeHint" stdset="0">
          <size>
           <width>20</width>
           <height>10</height>
          </size>
         </property>
        </spacer>
       </item>
       <item>
        <widget class="Line" name="m_MidSeparator">
         <property name="orientation">
          <enum>Qt::Horizontal</enum>
         </property>
        </widget>
       </item>
       <item>
        <spacer name="verticalSpacer_5">
         <property name="orientation">
          <enum>Qt::Vertical</enum>
         </property>
         <property name="sizeType">
          <enum>QSizePolicy::Fixed</enum>
         </property>
         <property name="sizeHint" stdset="0">
          <size>
           <width>20</width>
           <height>10</height>
          </size>
         </property>
        </spacer>
       </item>
       <item>
        <widget class="QLabel" name="m_SaveDataManagerLabel">
         <property name="text">
          <string>Save to Data Manager</string>
         </property>
        </widget>
       </item>
       <item>
        <widget class="QFrame" name="m_SaveDataManagerFrame">
         <property name="frameShape">
          <enum>QFrame::NoFrame</enum>
         </property>
         <property name="frameShadow">
          <enum>QFrame::Raised</enum>
         </property>
         <layout class="QHBoxLayout" name="horizontalLayout_3">
          <property name="leftMargin">
           <number>0</number>
          </property>
          <property name="topMargin">
           <number>0</number>
          </property>
          <property name="rightMargin">
           <number>0</number>
          </property>
          <property name="bottomMargin">
           <number>0</number>
          </property>
          <item>
           <widget class="QPushButton" name="m_SaveSegmentationPushButton">
            <property name="text">
             <string>Segmentation</string>
            </property>
           </widget>
          </item>
          <item>
           <widget class="QPushButton" name="m_SavePredictionsPushButton">
            <property name="text">
             <string>Predictions</string>
            </property>
           </widget>
          </item>
         </layout>
        </widget>
       </item>
       <item>
        <widget class="QFrame" name="m_ExportFrame">
         <property name="frameShape">
          <enum>QFrame::NoFrame</enum>
         </property>
         <property name="frameShadow">
          <enum>QFrame::Raised</enum>
         </property>
         <layout class="QHBoxLayout" name="horizontalLayout_4">
          <property name="leftMargin">
           <number>0</number>
          </property>
          <property name="topMargin">
           <number>0</number>
          </property>
          <property name="rightMargin">
           <number>0</number>
          </property>
          <property name="bottomMargin">
           <number>0</number>
          </property>
         </layout>
        </widget>
       </item>
       <item>
        <spacer name="verticalSpacer_7">
         <property name="orientation">
          <enum>Qt::Vertical</enum>
         </property>
         <property name="sizeType">
          <enum>QSizePolicy::Fixed</enum>
         </property>
         <property name="sizeHint" stdset="0">
          <size>
           <width>20</width>
           <height>10</height>
          </size>
         </property>
        </spacer>
       </item>
       <item>
        <widget class="Line" name="m_LowerSeparator">
         <property name="orientation">
          <enum>Qt::Horizontal</enum>
         </property>
        </widget>
       </item>
      </layout>
     </widget>
    </item>
    <item>
     <spacer name="verticalSpacer_6">
      <property name="orientation">
       <enum>Qt::Vertical</enum>
      </property>
      <property name="sizeType">
       <enum>QSizePolicy::Fixed</enum>
      </property>
      <property name="sizeHint" stdset="0">
       <size>
        <width>20</width>
        <height>10</height>
       </size>
      </property>
     </spacer>
    </item>
    <item>
     <widget class="QProgressBar" name="m_ProgressBar">
      <property name="value">
       <number>0</number>
      </property>
      <property name="textVisible">
       <bool>false</bool>
      </property>
      <property name="invertedAppearance">
       <bool>false</bool>
      </property>
     </widget>
    </item>
    <item>
     <spacer name="verticalSpacer">
      <property name="orientation">
       <enum>Qt::Vertical</enum>
      </property>
      <property name="sizeHint" stdset="0">
       <size>
        <width>20</width>
        <height>40</height>
       </size>
      </property>
     </spacer>
    </item>
   </layout>
  </widget>
  <layoutdefault spacing="6" margin="11"/>
  <resources>
   <include location="../../../../../MITK-al-superbuild/MITK-build/Plugins/org.mitk.gui.qt.activelearning/org_mitk_gui_qt_activelearning_cached.qrc"/>
  </resources>
  <connections/>
 </ui>