LightGBMWrapper - Type of return value of classify methods for binary classification was changed to two-element vector of probabilities (instead of one-element vector)

Author: hudlijac

src/wif/ml/lightGBMWrapper.cpp

@@ -53,9 +53,11 @@ ClfResult LightGBMWrapper::classify(const FlowFeatures& flowFeatures)
 	std::vector<double> dataToClassify; // classified features from flowfeatures are extracted here
 	int64_t outLen; // length of output result
 	int numOfClasses; // number of classes
+
 	LGBM_BoosterGetNumClasses(m_booster, &numOfClasses);
 
 	std::vector<double> pred(numOfClasses); // vector with predictions
+
 	for (const auto& featureID : m_featureIDs) {
 		double value = flowFeatures.get<double>(featureID);
 		dataToClassify.push_back(value);
@@ -74,6 +76,12 @@ ClfResult LightGBMWrapper::classify(const FlowFeatures& flowFeatures)
 		&outLen,
 		pred.data());
 
+	if (numOfClasses == 1) {
+		double tmp = pred[0];
+
+		pred.insert(pred.begin(), (1.0 - tmp));
+	}
+
 	return ClfResult(pred);
 }
 
@@ -114,6 +122,13 @@ std::vector<ClfResult> LightGBMWrapper::classify(const std::vector<FlowFeatures>
 		std::vector<double> probabilities(
 			pred.begin() + idx * numOfClasses,
 			pred.begin() + (idx + 1) * numOfClasses);
+
+		if (numOfClasses == 1) {
+			double tmp = probabilities[0];
+
+			probabilities.insert(probabilities.begin(), (1.0 - tmp));
+		}
+
 		burstResults.emplace_back(probabilities);
 	}