Improve OCR option of video condition

* Improve preprocessing of image by separating text color from the rest
  of the image properly
* Add option to choose how similar the matched color can be to other
  colors for it to still be counted as a match
* Show the prepreprocessed image in the "show match" preview dialog

src/macro-external/video/macro-condition-video.cpp

@@ -315,7 +315,7 @@ bool MacroConditionVideo::CheckOCR()
 	}
 
 	auto text = RunOCR(_ocrParameters.GetOCR(), _screenshotData.image,
-			   _ocrParameters.color);
+			   _ocrParameters.color, _ocrParameters.colorThreshold);
 
 	if (_ocrParameters.regex.Enabled()) {
 		auto expr = _ocrParameters.regex.GetRegularExpression(
@@ -466,6 +466,13 @@ OCREdit::OCREdit(QWidget *parent, PreviewDialog *previewDialog,
 	  _textColor(new QLabel),
 	  _selectColor(new QPushButton(obs_module_text(
 		  "AdvSceneSwitcher.condition.video.selectColor"))),
+	  _colorThreshold(new SliderSpinBox(
+		  0., 1.,
+		  obs_module_text(
+			  "AdvSceneSwitcher.condition.video.colorDeviationThreshold"),
+		  obs_module_text(
+			  "AdvSceneSwitcher.condition.video.colorDeviationThresholdDescription"),
+		  true)),
 	  _pageSegMode(new QComboBox()),
 	  _languageCode(new VariableLineEdit(this)),
 	  _previewDialog(previewDialog),
@@ -475,6 +482,11 @@ OCREdit::OCREdit(QWidget *parent, PreviewDialog *previewDialog,
 
 	QWidget::connect(_selectColor, SIGNAL(clicked()), this,
 			 SLOT(SelectColorClicked()));
+	QWidget::connect(
+		_colorThreshold,
+		SIGNAL(DoubleValueChanged(const NumberVariable<double> &)),
+		this,
+		SLOT(ColorThresholdChanged(const NumberVariable<double> &)));
 	QWidget::connect(_matchText, SIGNAL(textChanged()), this,
 			 SLOT(MatchTextChanged()));
 	QWidget::connect(_regex, SIGNAL(RegexConfigChanged(RegexConfig)), this,
@@ -516,11 +528,13 @@ OCREdit::OCREdit(QWidget *parent, PreviewDialog *previewDialog,
 			"AdvSceneSwitcher.condition.video.entry.orcColorPick"),
 		colorPickLayout, widgetPlaceholders);
 	layout->addLayout(colorPickLayout);
+	layout->addWidget(_colorThreshold);
 	setLayout(layout);
 
 	_matchText->setPlainText(_data->_ocrParameters.text);
 	_regex->SetRegexConfig(_data->_ocrParameters.regex);
 	SetupColorLabel(_data->_ocrParameters.color);
+	_colorThreshold->SetDoubleValue(_data->_ocrParameters.colorThreshold);
 	_pageSegMode->setCurrentIndex(_pageSegMode->findData(
 		static_cast<int>(_data->_ocrParameters.GetPageMode())));
 	_languageCode->setText(_data->_ocrParameters.GetLanguageCode());
@@ -556,6 +570,18 @@ void OCREdit::SelectColorClicked()
 	_previewDialog->OCRParametersChanged(_data->_ocrParameters);
 }
 
+void OCREdit::ColorThresholdChanged(const DoubleVariable &value)
+{
+	if (_loading || !_data) {
+		return;
+	}
+
+	auto lock = LockContext();
+	_data->_ocrParameters.colorThreshold = value;
+
+	_previewDialog->OCRParametersChanged(_data->_ocrParameters);
+}
+
 void OCREdit::MatchTextChanged()
 {
 	if (_loading || !_data) {

src/macro-external/video/macro-condition-video.hpp

@@ -117,6 +117,7 @@ public:
 
 private slots:
 	void SelectColorClicked();
+	void ColorThresholdChanged(const NumberVariable<double> &);
 	void MatchTextChanged();
 	void RegexChanged(RegexConfig conf);
 	void PageSegModeChanged(int);
@@ -129,6 +130,7 @@ private:
 	RegexConfigWidget *_regex;
 	QLabel *_textColor;
 	QPushButton *_selectColor;
+	SliderSpinBox *_colorThreshold;
 	QComboBox *_pageSegMode;
 	VariableLineEdit *_languageCode;
 

src/macro-external/video/opencv-helpers.cpp

@@ -124,21 +124,40 @@ uchar GetAvgBrightness(QImage &img)
 	return brightnessSum / (hsvImage.rows * hsvImage.cols);
 }
 
-cv::UMat PreprocessForOCR(const QImage &image, const QColor &color)
+static bool colorIsSimilar(const QColor &color1, const QColor &color2,
+			   int maxDiff)
 {
-	auto mat = QImageToMat(image);
+	const int diffRed = std::abs(color1.red() - color2.red());
+	const int diffGreen = std::abs(color1.green() - color2.green());
+	const int diffBlue = std::abs(color1.blue() - color2.blue());
 
-	// Only keep the desired color
-	cv::cvtColor(mat, mat, cv::COLOR_RGBA2RGB);
-	cv::cvtColor(mat, mat, cv::COLOR_RGB2HSV);
-	cv::inRange(mat, cv::Scalar(0, 0, 0),
-		    cv::Scalar(color.red(), color.green(), color.blue()), mat);
+	return diffRed <= maxDiff && diffGreen <= maxDiff &&
+	       diffBlue <= maxDiff;
+}
 
-	// Invert to improve ORC detection
-	cv::bitwise_not(mat, mat);
+cv::Mat PreprocessForOCR(const QImage &image, const QColor &textColor,
+			 double colorDiff)
+{
+	auto umat = QImageToMat(image);
+	auto mat = umat.getMat(cv::ACCESS_RW);
 
-	// Scale image up if selected area is too small
-	// Results will probably still be unsatisfying
+	// Tesseract works best when matching black text on a white background,
+	// so everything that matches the text color will be displayed black
+	// while the rest of the image should be white.
+	const int diff = colorDiff * 255;
+	for (int y = 0; y < image.height(); y++) {
+		for (int x = 0; x < image.width(); x++) {
+			if (colorIsSimilar(image.pixelColor(x, y), textColor,
+					   diff)) {
+				mat.at<cv::Vec4b>(y, x) = {0, 0, 0, 255};
+			} else {
+				mat.at<cv::Vec4b>(y, x) = {255, 255, 255, 255};
+			}
+		}
+	}
+
+	// Scale image up if selected area is very small.
+	// Results will probably still be unsatisfying.
 	if (mat.rows <= 300 || mat.cols <= 300) {
 		double scale = 0.;
 		if (mat.rows < mat.cols) {
@@ -146,26 +165,28 @@ cv::UMat PreprocessForOCR(const QImage &image, const QColor &color)
 		} else {
 			scale = 300. / mat.cols;
 		}
-
 		cv::resize(mat, mat,
 			   cv::Size(mat.cols * scale, mat.rows * scale),
 			   cv::INTER_CUBIC);
 	}
 
-	return mat;
+	cv::Mat result;
+	mat.copyTo(result);
+	return result;
 }
 
 std::string RunOCR(tesseract::TessBaseAPI *ocr, const QImage &image,
-		   const QColor &color)
+		   const QColor &color, double colorDiff)
 {
 	if (image.isNull()) {
 		return "";
 	}
 
 #ifdef OCR_SUPPORT
-	auto mat = PreprocessForOCR(image, color);
-	ocr->SetImage(mat.getMat(cv::ACCESS_READ).data, mat.cols, mat.rows, 1,
-		      mat.step);
+	auto mat = PreprocessForOCR(image, color, colorDiff);
+	cv::Mat gray;
+	cv::cvtColor(mat, gray, cv::COLOR_RGBA2GRAY);
+	ocr->SetImage(gray.data, gray.cols, gray.rows, 1, gray.step);
 	ocr->Recognize(0);
 	std::unique_ptr<char[]> detectedText(ocr->GetUTF8Text());
 
@@ -189,17 +210,8 @@ bool ContainsPixelsInColorRange(const QImage &image, const QColor &color,
 
 	for (int y = 0; y < image.height(); y++) {
 		for (int x = 0; x < image.width(); x++) {
-			const auto pixelColor = image.pixelColor(x, y);
-			const int diffRed =
-				std::abs(pixelColor.red() - color.red());
-			const int diffGreen =
-				std::abs(pixelColor.green() - color.green());
-			const int diffBlue =
-				std::abs(pixelColor.blue() - color.blue());
-
-			if (diffRed <= maxColorDiff &&
-			    diffGreen <= maxColorDiff &&
-			    diffBlue <= maxColorDiff) {
+			if (colorIsSimilar(image.pixelColor(x, y), color,
+					   maxColorDiff)) {
 				matchingPixels++;
 			}
 		}

src/macro-external/video/opencv-helpers.hpp

@@ -59,8 +59,10 @@ std::vector<cv::Rect> MatchObject(QImage &img, cv::CascadeClassifier &cascade,
 				  const cv::Size &minSize,
 				  const cv::Size &maxSize);
 uchar GetAvgBrightness(QImage &img);
-cv::UMat PreprocessForOCR(const QImage &image, const QColor &color);
-std::string RunOCR(tesseract::TessBaseAPI *, const QImage &, const QColor &);
+cv::Mat PreprocessForOCR(const QImage &image, const QColor &color,
+			 double colorDiff);
+std::string RunOCR(tesseract::TessBaseAPI *, const QImage &, const QColor &,
+		   double colorDiff);
 bool ContainsPixelsInColorRange(const QImage &image, const QColor &color,
 				double colorDeviationThreshold,
 				double totalPixelMatchThreshold);

src/macro-external/video/paramerter-wrappers.cpp

@@ -221,6 +221,17 @@ static void SaveColor(obs_data_t *obj, const char *name, const QColor &color)
 	obs_data_release(data);
 }
 
+static QColor LoadColor(obs_data_t *obj, const char *name)
+{
+	QColor color = Qt::black;
+	auto data = obs_data_get_obj(obj, name);
+	color.setRed(obs_data_get_int(data, "red"));
+	color.setGreen(obs_data_get_int(data, "green"));
+	color.setBlue(obs_data_get_int(data, "blue"));
+	obs_data_release(data);
+	return color;
+}
+
 OCRParameters::OCRParameters()
 {
 	Setup();
@@ -238,6 +249,7 @@ OCRParameters::OCRParameters(const OCRParameters &other)
 	: text(other.text),
 	  regex(other.regex),
 	  color(other.color),
+	  colorThreshold(other.colorThreshold),
 	  pageSegMode(other.pageSegMode)
 {
 	Setup();
@@ -251,6 +263,7 @@ OCRParameters &OCRParameters::operator=(const OCRParameters &other)
 	text = other.text;
 	regex = other.regex;
 	color = other.color;
+	colorThreshold = other.colorThreshold;
 	pageSegMode = other.pageSegMode;
 	ocr->SetPageSegMode(pageSegMode);
 	return *this;
@@ -263,6 +276,7 @@ bool OCRParameters::Save(obs_data_t *obj) const
 	regex.Save(data);
 	languageCode.Save(data, "language");
 	SaveColor(data, "textColor", color);
+	colorThreshold.Save(data, "colorThreshold");
 	obs_data_set_int(data, "pageSegMode", static_cast<int>(pageSegMode));
 	obs_data_set_int(data, "version", 1);
 	obs_data_set_obj(obj, "ocrData", data);
@@ -270,17 +284,6 @@ bool OCRParameters::Save(obs_data_t *obj) const
 	return true;
 }
 
-static QColor LoadColor(obs_data_t *obj, const char *name)
-{
-	QColor color = Qt::black;
-	auto data = obs_data_get_obj(obj, name);
-	color.setRed(obs_data_get_int(data, "red"));
-	color.setGreen(obs_data_get_int(data, "green"));
-	color.setBlue(obs_data_get_int(data, "blue"));
-	obs_data_release(data);
-	return color;
-}
-
 bool OCRParameters::Load(obs_data_t *obj)
 {
 	auto data = obs_data_get_obj(obj, "ocrData");
@@ -289,6 +292,9 @@ bool OCRParameters::Load(obs_data_t *obj)
 	obs_data_set_default_string(data, "language", "eng");
 	languageCode.Load(data, "language");
 	color = LoadColor(data, "textColor");
+	if (obs_data_has_user_value(data, "version")) {
+		colorThreshold.Load(data, "colorThreshold");
+	}
 	pageSegMode = static_cast<tesseract::PageSegMode>(
 		obs_data_get_int(data, "pageSegMode"));
 	obs_data_release(data);

src/macro-external/video/paramerter-wrappers.hpp

@@ -98,6 +98,7 @@ public:
 	StringVariable text = obs_module_text("AdvSceneSwitcher.enterText");
 	RegexConfig regex = RegexConfig::PartialMatchRegexConfig();
 	QColor color = Qt::black;
+	DoubleVariable colorThreshold = 0.3;
 	StringVariable languageCode = "eng";
 
 private:

src/macro-external/video/preview-dialog.cpp

@@ -331,13 +331,11 @@ void PreviewImage::MarkMatch(QImage &screenshot,
 			markObjects(screenshot, objects);
 		}
 	} else if (condition == VideoCondition::OCR) {
-		auto text =
-			RunOCR(ocrParams.GetOCR(), screenshot, ocrParams.color);
+		auto text = RunOCR(ocrParams.GetOCR(), screenshot,
+				   ocrParams.color, ocrParams.colorThreshold);
 		QString status(obs_module_text(
 			"AdvSceneSwitcher.condition.video.ocrMatchSuccess"));
 		emit StatusUpdate(status.arg(QString::fromStdString(text)));
-
-		// TODO: show preprocessed image
 	}
 }