diff --git a/GUI/Model/Data/Data2DItem.cpp b/GUI/Model/Data/Data2DItem.cpp
index b0d00cf5b26c412c4121c2ea0f65cb58fb38b02d..f8033c8676d6f12088b69396a79f38e035fa425f 100644
--- a/GUI/Model/Data/Data2DItem.cpp
+++ b/GUI/Model/Data/Data2DItem.cpp
@@ -21,6 +21,7 @@
 #include "Device/Mask/MaskStack.h"
 #include "GUI/Model/Axis/AmplitudeAxisItem.h"
 #include "GUI/Model/Axis/BasicAxisItem.h"
+#include "GUI/Model/Mask/MasksQModel.h"
 #include "GUI/Model/Mask/MasksSet.h"
 #include "GUI/Support/Data/ComboProperty.h"
 #include "GUI/Support/Style/QCP_Util.h"
@@ -365,18 +366,18 @@ Datafield* Data2DItem::createMaskedField() const
 
 MasksSet* Data2DItem::projectionsSet()
 {
-    return m_proj_model.get();
+    return m_proj_model ? m_proj_model->set() : nullptr;
 }
 
 const MasksSet* Data2DItem::projectionsSet() const
 {
-    return m_proj_model.get();
+    return m_proj_model ? m_proj_model->set() : nullptr;
 }
 
-MasksSet* Data2DItem::getOrCreateProjectionModel()
+MasksQModel* Data2DItem::getOrCreateProjectionModel()
 {
     if (!m_proj_model)
-        m_proj_model = std::make_unique<MasksSet>();
+        m_proj_model = std::make_unique<MasksQModel>(new MasksSet);
 
     return m_proj_model.get();
 }
diff --git a/GUI/Model/Data/Data2DItem.h b/GUI/Model/Data/Data2DItem.h
index 50523cc459fd61bc62bf92f1feeabb86d4180fe0..1294b41f5d329d136c26502e97b8f394f78e28d9 100644
--- a/GUI/Model/Data/Data2DItem.h
+++ b/GUI/Model/Data/Data2DItem.h
@@ -83,7 +83,7 @@ public:
 
     MasksSet* projectionsSet();
     const MasksSet* projectionsSet() const;
-    MasksSet* getOrCreateProjectionModel();
+    MasksQModel* getOrCreateProjectionModel();
 
     size_t rank() const override { return 2; }
     //! Returns dimensions and axesbins of data
@@ -115,7 +115,7 @@ private:
     std::unique_ptr<ComboProperty> m_gradient;
     std::unique_ptr<AmplitudeAxisItem> m_z_axis;
     std::unique_ptr<MasksSet> m_model;
-    std::unique_ptr<MasksSet> m_proj_model;
+    std::unique_ptr<MasksQModel> m_proj_model;
 };
 
 #endif // BORNAGAIN_GUI_MODEL_DATA_DATA2DITEM_H