New training dialog + selection of the most important features

Train3DMASCDialog.ui

@@ -6,8 +6,8 @@
    <rect>
     <x>0</x>
     <y>0</y>
-    <width>397</width>
+    <width>400</width>
-    <height>476</height>
+    <height>500</height>
    </rect>
   </property>
   <property name="windowTitle">
@@ -103,7 +103,7 @@
    <item>
     <widget class="QGroupBox" name="otherGroupBox">
      <property name="title">
-      <string>Other</string>
+      <string>Other parameters</string>
      </property>
      <layout class="QFormLayout" name="formLayout_2">
       <item row="0" column="0">
@@ -136,26 +136,96 @@
     </widget>
    </item>
    <item>
-    <spacer name="verticalSpacer">
+    <widget class="QFrame" name="frame">
-     <property name="orientation">
+     <property name="frameShape">
-      <enum>Qt::Vertical</enum>
+      <enum>QFrame::StyledPanel</enum>
      </property>
-     <property name="sizeHint" stdset="0">
+     <property name="frameShadow">
-      <size>
+      <enum>QFrame::Raised</enum>
-       <width>20</width>
-       <height>40</height>
-      </size>
      </property>
-    </spacer>
+     <layout class="QHBoxLayout" name="horizontalLayout">
+      <item>
+       <widget class="QPushButton" name="runPushButton">
+        <property name="text">
+         <string>Run</string>
+        </property>
+       </widget>
+      </item>
+      <item>
+       <widget class="QPushButton" name="savePushButton">
+        <property name="enabled">
+         <bool>false</bool>
+        </property>
+        <property name="text">
+         <string>Save</string>
+        </property>
+       </widget>
+      </item>
+      <item>
+       <widget class="QPushButton" name="closePushButton">
+        <property name="text">
+         <string>Close</string>
+        </property>
+       </widget>
+      </item>
+     </layout>
+     <zorder>runPushButton</zorder>
+     <zorder>closePushButton</zorder>
+     <zorder>savePushButton</zorder>
+    </widget>
    </item>
    <item>
-    <widget class="QDialogButtonBox" name="buttonBox">
+    <widget class="QGroupBox" name="resultGroupBox">
-     <property name="orientation">
+     <property name="sizePolicy">
-      <enum>Qt::Horizontal</enum>
+      <sizepolicy hsizetype="Preferred" vsizetype="Expanding">
+       <horstretch>0</horstretch>
+       <verstretch>0</verstretch>
+      </sizepolicy>
      </property>
-     <property name="standardButtons">
+     <property name="title">
-      <set>QDialogButtonBox::Cancel|QDialogButtonBox::Ok</set>
+      <string>Results</string>
      </property>
+     <layout class="QVBoxLayout" name="verticalLayout_2">
+      <item>
+       <widget class="QLabel" name="resultLabel">
+        <property name="text">
+         <string>correct = 0/0 - accuracy = 0.0</string>
+        </property>
+       </widget>
+      </item>
+      <item>
+       <widget class="QTableWidget" name="tableWidget">
+        <property name="verticalScrollBarPolicy">
+         <enum>Qt::ScrollBarAlwaysOn</enum>
+        </property>
+        <property name="showDropIndicator" stdset="0">
+         <bool>false</bool>
+        </property>
+        <property name="dragDropOverwriteMode">
+         <bool>false</bool>
+        </property>
+        <property name="selectionMode">
+         <enum>QAbstractItemView::NoSelection</enum>
+        </property>
+        <attribute name="horizontalHeaderDefaultSectionSize">
+         <number>170</number>
+        </attribute>
+        <attribute name="verticalHeaderVisible">
+         <bool>false</bool>
+        </attribute>
+        <column>
+         <property name="text">
+          <string>Feature</string>
+         </property>
+        </column>
+        <column>
+         <property name="text">
+          <string>Importance</string>
+         </property>
+        </column>
+       </widget>
+      </item>
+     </layout>
     </widget>
    </item>
   </layout>
@@ -163,34 +233,18 @@
  <resources/>
  <connections>
   <connection>
-   <sender>buttonBox</sender>
+   <sender>runPushButton</sender>
-   <signal>accepted()</signal>
+   <signal>clicked()</signal>
    <receiver>Train3DMASCDialog</receiver>
    <slot>accept()</slot>
    <hints>
     <hint type="sourcelabel">
-     <x>248</x>
+     <x>76</x>
-     <y>254</y>
+     <y>296</y>
     </hint>
     <hint type="destinationlabel">
-     <x>157</x>
+     <x>198</x>
-     <y>274</y>
+     <y>237</y>
-    </hint>
-   </hints>
-  </connection>
-  <connection>
-   <sender>buttonBox</sender>
-   <signal>rejected()</signal>
-   <receiver>Train3DMASCDialog</receiver>
-   <slot>reject()</slot>
-   <hints>
-    <hint type="sourcelabel">
-     <x>316</x>
-     <y>260</y>
-    </hint>
-    <hint type="destinationlabel">
-     <x>286</x>
-     <y>274</y>
     </hint>
    </hints>
   </connection>

q3DMASC.cpp

@@ -162,19 +162,19 @@ void q3DMASCPlugin::doClassifyAction()
 	corePoints.origin = corePoints.cloud = clouds[mainCloudLabel];
 
 	//prepare the main cloud
-	ccProgressDialog pDlg(true, m_app->getMainWindow());
+	ccProgressDialog progressDlg(true, m_app->getMainWindow());
-	pDlg.setAutoClose(false); //we don't want the progress dialog to 'pop' for each feature
+	progressDlg.setAutoClose(false); //we don't want the progress dialog to 'pop' for each feature
 	QString error;
 	SFCollector generatedScalarFields;
-	if (!masc::Tools::PrepareFeatures(corePoints, features, error, &pDlg, &generatedScalarFields))
+	if (!masc::Tools::PrepareFeatures(corePoints, features, error, &progressDlg, &generatedScalarFields))
 	{
 		m_app->dispToConsole(error, ccMainAppInterface::ERR_CONSOLE_MESSAGE);
 		generatedScalarFields.clear();
 		return;
 	}
-	pDlg.close();
+	progressDlg.close();
 	QCoreApplication::processEvents();
-	pDlg.setAutoClose(true); //restore the default behavior of the progress dialog
+	progressDlg.setAutoClose(true); //restore the default behavior of the progress dialog
 
 	//apply classifier
 	{
@@ -193,6 +193,15 @@ void q3DMASCPlugin::doClassifyAction()
 	}
 }
 
+struct FeatureSelection
+{
+	FeatureSelection(masc::Feature::Shared f = masc::Feature::Shared(nullptr)) : feature(f) {}
+	masc::Feature::Shared feature;
+	bool selected = true;
+	bool prepared = false;
+	float importance = std::numeric_limits<float>::quiet_NaN();
+};
+
 void q3DMASCPlugin::doTrainAction()
 {
 	//disclaimer accepted?
@@ -236,8 +245,32 @@ void q3DMASCPlugin::doTrainAction()
 		group->addChild(pc);
 	}
 
-	ccProgressDialog pDlg(true, m_app->getMainWindow());
+	//show the training dialog for the first time
-	if (!corePoints.prepare(&pDlg))
+	Train3DMASCDialog trainDlg(m_app->getMainWindow());
+	trainDlg.maxDepthSpinBox->setValue(s_params.rt.maxDepth);
+	trainDlg.maxTreeCountSpinBox->setValue(s_params.rt.maxTreeCount);
+	trainDlg.activeVarCountSpinBox->setValue(s_params.rt.activeVarCount);
+	trainDlg.minSampleCountSpinBox->setValue(s_params.rt.minSampleCount);
+	trainDlg.testDataRatioSpinBox->setValue(static_cast<int>(s_params.testDataRatio * 100));
+
+	//display the loaded features and let the user select the ones to use
+	trainDlg.setResultText("Select features and press 'Run'");
+	std::vector<FeatureSelection> originalFeatures;
+	originalFeatures.reserve(features.size());
+	for (const masc::Feature::Shared& f : features)
+	{
+		originalFeatures.push_back(FeatureSelection(f));
+		trainDlg.addFeature(f->toString(), originalFeatures.back().importance, originalFeatures.back().selected);
+	}
+	if (!trainDlg.exec())
+	{
+		return;
+	}
+	assert(!trainDlg.shouldSaveClassifier()); //the save button should be disabled at this point
+
+	//compute the core points (if necessary)
+	ccProgressDialog progressDlg(true, m_app->getMainWindow());
+	if (!corePoints.prepare(&progressDlg))
 	{
 		m_app->dispToConsole("Failed to compute/prepare the core points!", ccMainAppInterface::ERR_CONSOLE_MESSAGE);
 		delete group;
@@ -271,33 +304,64 @@ void q3DMASCPlugin::doTrainAction()
 	m_app->addToDB(group);
 	QCoreApplication::processEvents();
 
-	pDlg.setAutoClose(false); //we don't want the progress dialog to 'pop' for each feature
+	//train / test subsets
-	QString error;
+	QScopedPointer<CCLib::ReferenceCloud> trainSubset(new CCLib::ReferenceCloud(corePoints.cloud));
-	if (!masc::Tools::PrepareFeatures(corePoints, features, error, &pDlg))
+	QScopedPointer<CCLib::ReferenceCloud> testSubset(new CCLib::ReferenceCloud(corePoints.cloud));
-	{
+	float previousTrainSubsetRatio = -1.0f;
-		m_app->dispToConsole(error, ccMainAppInterface::ERR_CONSOLE_MESSAGE);
-		delete group;
-		return;
-	}
-	pDlg.setAutoClose(true); //restore the default behavior of the progress dialog
-	pDlg.close();
-	QCoreApplication::processEvents();
 
-	m_app->redrawAll();
+	for (int iteration = 0; ; ++iteration)
-
-	while (true)
 	{
-		Train3DMASCDialog trainDlg(m_app->getMainWindow());
+		//look for selected features
-		trainDlg.maxDepthSpinBox->setValue(s_params.rt.maxDepth);
+		features.clear();
-		trainDlg.maxTreeCountSpinBox->setValue(s_params.rt.maxTreeCount);
+		masc::Feature::Set toPrepare;
-		trainDlg.activeVarCountSpinBox->setValue(s_params.rt.activeVarCount);
+		for (size_t i = 0; i < originalFeatures.size(); ++i)
-		trainDlg.minSampleCountSpinBox->setValue(s_params.rt.minSampleCount);
-		trainDlg.testDataRatioSpinBox->setValue(static_cast<int>(s_params.testDataRatio * 100));
-		if (!trainDlg.exec())
 		{
-			return;
+			originalFeatures[i].selected = trainDlg.isFeatureSelected(i);
+
+			//if the feature is selected
+			if (originalFeatures[i].selected)
+			{
+				if (!originalFeatures[i].prepared)
+				{
+					//we should prepare it first!
+					toPrepare.push_back(originalFeatures[i].feature);
+				}
+				features.push_back(originalFeatures[i].feature);
+			}
 		}
 
+		if (features.empty())
+		{
+			m_app->dispToConsole("No feature selected!", ccMainAppInterface::ERR_CONSOLE_MESSAGE);
+			continue;
+		}
+
+		//prepare the features
+		if (!toPrepare.empty())
+		{
+			progressDlg.setAutoClose(false); //we don't want the progress dialog to 'pop' for each feature
+			QString error;
+			if (!masc::Tools::PrepareFeatures(corePoints, toPrepare, error, &progressDlg))
+			{
+				m_app->dispToConsole(error, ccMainAppInterface::ERR_CONSOLE_MESSAGE);
+				return;
+			}
+			progressDlg.setAutoClose(true); //restore the default behavior of the progress dialog
+			progressDlg.close();
+			QCoreApplication::processEvents();
+			m_app->redrawAll();
+
+			//flag the prepared features as 'prepared' ;)
+			for (FeatureSelection& fs : originalFeatures)
+			{
+				if (fs.selected && !fs.prepared)
+					fs.prepared = true;
+			}
+		}
+
+		masc::Classifier classifier;
+
+		//retrieve parameters
 		s_params.rt.maxDepth = trainDlg.maxDepthSpinBox->value();
 		s_params.rt.maxTreeCount = trainDlg.maxTreeCountSpinBox->value();
 		s_params.rt.activeVarCount = trainDlg.activeVarCountSpinBox->value();
@@ -307,64 +371,116 @@ void q3DMASCPlugin::doTrainAction()
 		{
 			assert(false);
 			m_app->dispToConsole("Invalid test data ratio", ccMainAppInterface::ERR_CONSOLE_MESSAGE);
-			return;
 		}
-
+		else
-		//randomly select the training points
-		QScopedPointer<CCLib::ReferenceCloud> trainSubset(new CCLib::ReferenceCloud(corePoints.cloud));
-		QScopedPointer<CCLib::ReferenceCloud> testSubset(new CCLib::ReferenceCloud(corePoints.cloud));
-		if (!masc::Tools::RandomSubset(corePoints.cloud, s_params.testDataRatio, testSubset.data(), trainSubset.data()))
 		{
-			m_app->dispToConsole("Not enough memory", ccMainAppInterface::ERR_CONSOLE_MESSAGE);
+			if (previousTrainSubsetRatio != s_params.testDataRatio)
-			return;
-		}
-
-		//train the classifier
-		masc::Classifier classifier;
-		{
-			QString errorMessage;
-			if (!classifier.train(corePoints.cloud, s_params.rt, features, errorMessage, trainSubset.data(), m_app->getMainWindow()))
 			{
-				m_app->dispToConsole(errorMessage, ccMainAppInterface::ERR_CONSOLE_MESSAGE);
+				//randomly select the training points
-				return;
+				testSubset->clear();
-			}
+				trainSubset->clear();
-
+				if (!masc::Tools::RandomSubset(corePoints.cloud, s_params.testDataRatio, testSubset.data(), trainSubset.data()))
-			QString outputFilename;
-			{
-				QSettings settings;
-				settings.beginGroup("3DMASC");
-				QString outputPath = settings.value("FilePath", QCoreApplication::applicationDirPath()).toString();
-				outputFilename = QFileDialog::getSaveFileName(m_app->getMainWindow(), "Save 3DMASC classifier", outputPath, "*.txt");
-				if (outputFilename.isNull())
 				{
-					//process cancelled by the user
+					m_app->dispToConsole("Not enough memory", ccMainAppInterface::ERR_CONSOLE_MESSAGE);
 					return;
 				}
-				settings.setValue("FilePath", QFileInfo(outputFilename).absolutePath());
+				previousTrainSubsetRatio = s_params.testDataRatio;
-				settings.endGroup();
 			}
 
-			//save the classifier
+			//train the classifier
-			if (masc::Tools::SaveClassifier(outputFilename, features, classifier, m_app->getMainWindow()))
 			{
-				m_app->dispToConsole("Classifier succesfully saved to " + outputFilename, ccMainAppInterface::STD_CONSOLE_MESSAGE);
+				QString errorMessage;
+				if (!classifier.train(corePoints.cloud, s_params.rt, features, errorMessage, trainSubset.data(), m_app, m_app->getMainWindow()))
+				{
+					m_app->dispToConsole(errorMessage, ccMainAppInterface::ERR_CONSOLE_MESSAGE);
+					return;
+				}
+				trainDlg.setFirstRunDone();
+				trainDlg.shouldSaveClassifier();
+			}
+
+			//test the trained classifier
+			{
+				masc::Classifier::AccuracyMetrics metrics;
+				QString errorMessage;
+				if (!classifier.evaluate(features, testSubset.data(), metrics, errorMessage, m_app->getMainWindow()))
+				{
+					m_app->dispToConsole(errorMessage, ccMainAppInterface::ERR_CONSOLE_MESSAGE);
+					return;
+				}
+
+				QString resultText = QString("Correct guess = %1 / %2 --> accuracy = %3").arg(metrics.goodGuess).arg(metrics.sampleCount).arg(metrics.ratio);
+				m_app->dispToConsole(resultText, ccMainAppInterface::STD_CONSOLE_MESSAGE);
+				trainDlg.setResultText(resultText);
+
+				cv::Mat importanceMat = classifier.getVarImportance();
+				//m_app->dispToConsole(QString("Var importance size = %1 x %2").arg(importanceMat.rows).arg(importanceMat.cols));
+				assert(static_cast<int>(features.size()) == importanceMat.rows);
+				int selectedFeatureIndex = 0;
+				for (size_t i = 0; i < originalFeatures.size(); ++i)
+				{
+					if (originalFeatures[i].selected)
+					{
+						//m_app->dispToConsole(QString("Feature #%1 importance = %2").arg(i + 1).arg(importanceMat.at<float>(i, 0)));
+						assert(selectedFeatureIndex < importanceMat.rows);
+						originalFeatures[i].importance = importanceMat.at<float>(selectedFeatureIndex, 0);
+						++selectedFeatureIndex;
+					}
+					else
+					{
+						originalFeatures[i].importance = std::numeric_limits<float>::quiet_NaN();
+					}
+					trainDlg.setFeatureImportance(i, originalFeatures[i].importance);
+				}
 			}
 		}
 
-		//test classifier
+		while (true)
 		{
-			masc::Classifier::AccuracyMetrics metrics;
+			if (!trainDlg.exec())
-			QString errorMessage;
-			if (!classifier.evaluate(features, testSubset.data(), metrics, errorMessage, m_app->getMainWindow()))
 			{
-				m_app->dispToConsole(errorMessage, ccMainAppInterface::ERR_CONSOLE_MESSAGE);
+				//the dialog can be closed
 				return;
 			}
 
-			m_app->dispToConsole(QString("Correct = %1 / %2 --> accuracy = %3").arg(metrics.goodGuess).arg(metrics.sampleCount).arg(metrics.ratio), ccMainAppInterface::STD_CONSOLE_MESSAGE);
+			//if the save button has been clicked
+			if (trainDlg.shouldSaveClassifier())
+			{
+				//ask for the output filename
+				QString outputFilename;
+				{
+					QSettings settings;
+					settings.beginGroup("3DMASC");
+					QString outputPath = settings.value("FilePath", QCoreApplication::applicationDirPath()).toString();
+					outputFilename = QFileDialog::getSaveFileName(m_app->getMainWindow(), "Save 3DMASC classifier", outputPath, "*.txt");
+					if (outputFilename.isNull())
+					{
+						//process cancelled by the user
+						continue;
+					}
+					settings.setValue("FilePath", QFileInfo(outputFilename).absolutePath());
+					settings.endGroup();
+				}
+
+				//save the classifier
+				if (masc::Tools::SaveClassifier(outputFilename, features, classifier, m_app->getMainWindow()))
+				{
+					m_app->dispToConsole("Classifier succesfully saved to " + outputFilename, ccMainAppInterface::STD_CONSOLE_MESSAGE);
+					trainDlg.setClassifierSaved();
+				}
+				else
+				{
+					m_app->dispToConsole("Failed to save classifier file");
+				}
+			}
+			else //we will run the classifier another time
+			{
+				//stop the local loop
+				break;
+			}
 		}
 
-		break;
+		//we are going to restart the classification process
 	}
 }
 

q3DMASCClassifier.cpp

@@ -29,6 +29,9 @@
 //qCC_io
 #include <LASFields.h>
 
+//qCC_plugins
+#include <ccMainAppInterface.h>
+
 //Qt
 #include <QCoreApplication>
 #include <QProgressDialog>
@@ -360,6 +363,7 @@ bool Classifier::train(	const ccPointCloud* cloud,
 						const Feature::Set& features,
 						QString& errorMessage,
 						CCLib::ReferenceCloud* trainSubset/*=nullptr*/,
+						ccMainAppInterface* app/*=nullptr*/,
 						QWidget* parentWidget/*=nullptr*/)
 {
 	if (features.empty())
@@ -397,7 +401,10 @@ bool Classifier::train(	const ccPointCloud* cloud,
 	int sampleCount = static_cast<int>(trainSubset ? trainSubset->size() : cloud->size());
 	int attributesPerSample = static_cast<int>(features.size());
 
-	ccLog::Print(QString("[3DMASC] Training data: %1 samples with %2 feature(s)").arg(sampleCount).arg(attributesPerSample));
+	if (app)
+	{
+		app->dispToConsole(QString("[3DMASC] Training data: %1 samples with %2 feature(s)").arg(sampleCount).arg(attributesPerSample));
+	}
 
 	cv::Mat training_data, train_labels;
 	try
@@ -530,6 +537,17 @@ bool Classifier::train(	const ccPointCloud* cloud,
 		return false;
 	}
 
+	if (app)
+	{
+		cv::Mat mat = m_rtrees->getVarImportance();
+		app->dispToConsole(QString("Var importance size = %1 x %2").arg(mat.rows).arg(mat.cols));
+		assert(attributesPerSample == mat.rows);
+		for (int i = 0; i < mat.rows; ++i)
+		{
+			app->dispToConsole(QString("Feature #%1 importance = %2").arg(i + 1).arg(mat.at<float>(i, 0)));
+		}
+	}
+
 	return true;
 }
 

q3DMASCClassifier.h

@@ -31,6 +31,7 @@
 #include <opencv2/ml.hpp>
 
 class QWidget;
+class ccMainAppInterface;
 
 //! 3DMASC classifier
 namespace masc
@@ -43,7 +44,13 @@ namespace masc
 		Classifier();
 
 		//! Train the classifier
-		bool train(const ccPointCloud* cloud, const RandomTreesParams& params, const Feature::Set& features, QString& errorMessage, CCLib::ReferenceCloud* trainSubset = nullptr, QWidget* parentWidget = nullptr);
+		bool train(	const ccPointCloud* cloud,
+					const RandomTreesParams& params,
+					const Feature::Set& features,
+					QString& errorMessage,
+					CCLib::ReferenceCloud* trainSubset = nullptr,
+					ccMainAppInterface* app = nullptr,
+					QWidget* parentWidget = nullptr);
 
 		//! Classifier accuracy metrics
 		struct AccuracyMetrics
@@ -67,6 +74,8 @@ namespace masc
 		//! Loads the classifier from file
 		bool fromFile(QString filename, QWidget* parentWidget = nullptr);
 
+		inline cv::Mat getVarImportance() const { return m_rtrees->getVarImportance(); }
+
 	protected:
 
 		//! Random trees (OpenCV)

q3DMASCTools.cpp

@@ -632,7 +632,7 @@ static bool LoadFileCommon(	const QString& filename,
 		std::vector<double> scales;
 
 		QTextStream stream(&file);
-		for (int lineNumber = 0; ; ++lineNumber)
+		for (int lineNumber = 1; ; ++lineNumber)
 		{
 			QString line = stream.readLine();
 			if (line.isNull())
@@ -640,7 +640,6 @@ static bool LoadFileCommon(	const QString& filename,
 				//eof
 				break;
 			}
-			++lineNumber;
 
 			if (line.startsWith("#"))
 			{

qTrain3DMASCDialog.cpp

@@ -17,9 +17,96 @@
 
 #include "qTrain3DMASCDialog.h"
 
+//Qt
+#include <QTableWidgetItem>
+#include <QMessageBox>
+
+//System
+#include <assert.h>
+
 Train3DMASCDialog::Train3DMASCDialog(QWidget* parent/*=nullptr*/)
 	: QDialog(parent)
 	, Ui::Train3DMASCDialog()
+	, classifierSaved(false)
+	, saveRequested(false)
 {
 	setupUi(this);
+
+	connect(closePushButton, SIGNAL(clicked()), this, SLOT(onClose()));
+	connect(savePushButton, SIGNAL(clicked()), this, SLOT(onSave()));
+}
+
+void Train3DMASCDialog::clearResults()
+{
+	resultLabel->clear();
+	tableWidget->clear();
+}
+
+int Train3DMASCDialog::addFeature(QString name, float importance, bool isChecked/*=true*/)
+{
+	int index = tableWidget->rowCount();
+	tableWidget->setRowCount(index + 1);
+
+	QTableWidgetItem* nameItem = new QTableWidgetItem(name);
+	nameItem->setCheckState(isChecked ? Qt::Checked : Qt::Unchecked);
+	tableWidget->setItem(index, 0, nameItem);
+
+	QTableWidgetItem* importanceItem = new QTableWidgetItem(std::isnan(importance) ? QString() : QString::number(importance));
+	tableWidget->setItem(index, 1, importanceItem);
+
+	return index;
+}
+
+void Train3DMASCDialog::setResultText(QString text)
+{
+	resultLabel->setText(text);
+}
+
+void Train3DMASCDialog::setFirstRunDone()
+{
+	runPushButton->setText(tr("Retry"));
+	savePushButton->setEnabled(true);
+}
+
+bool Train3DMASCDialog::isFeatureSelected(size_t index) const
+{
+	if (static_cast<int>(index) >= tableWidget->rowCount())
+	{
+		assert(false);
+		return false;
+	}
+
+	return (tableWidget->item(static_cast<int>(index), 0)->checkState() == Qt::Checked);
+}
+
+void Train3DMASCDialog::setFeatureImportance(size_t index, float importance)
+{
+	if (static_cast<int>(index) >= tableWidget->rowCount())
+	{
+		assert(false);
+		return;
+	}
+
+	if (!std::isnan(importance))
+	{
+		tableWidget->item(static_cast<int>(index), 1)->setText(QString::number(importance));
+	}
+	else
+	{
+		tableWidget->item(static_cast<int>(index), 1)->setText(QString());
+	}
+}
+
+void Train3DMASCDialog::onClose()
+{
+	if (!classifierSaved && QMessageBox::question(this, "Classifier not saved", "Classifier not saved. Do you confirm you want to close the tool?", QMessageBox::Yes, QMessageBox::No) == QMessageBox::No)
+		return;
+
+	reject();
+}
+
+void Train3DMASCDialog::onSave()
+{
+	saveRequested = true;
+	accept();
 }

qTrain3DMASCDialog.h

@@ -32,4 +32,29 @@ public:
 	//! Default constructor
 	Train3DMASCDialog(QWidget* parent = nullptr);
 
+	void clearResults();
+
+	//! Adds a feature (entry) to the results table
+	/** \return the row index
+	**/
+	int addFeature(QString name, float importance, bool isChecked = true);
+
+	void setResultText(QString text);
+	void setFirstRunDone();
+	inline void setClassifierSaved() { classifierSaved = true; saveRequested = false; }
+
+	bool isFeatureSelected(size_t index) const;
+	void setFeatureImportance(size_t index, float importance);
+	
+	inline bool shouldSaveClassifier() const { return saveRequested; }
+
+protected slots:
+
+	void onClose();
+	void onSave();
+
+protected: //members
+
+	bool classifierSaved;
+	bool saveRequested;
 };