From c9c0e4db2e241dac43cb30862077729941dce9ba Mon Sep 17 00:00:00 2001 From: WarmUpTill Date: Tue, 27 Jun 2023 04:18:15 +0200 Subject: [PATCH] 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 --- .../video/macro-condition-video.cpp | 28 +++++++- .../video/macro-condition-video.hpp | 2 + src/macro-external/video/opencv-helpers.cpp | 68 +++++++++++-------- src/macro-external/video/opencv-helpers.hpp | 6 +- .../video/paramerter-wrappers.cpp | 28 +++++--- .../video/paramerter-wrappers.hpp | 1 + src/macro-external/video/preview-dialog.cpp | 6 +- 7 files changed, 93 insertions(+), 46 deletions(-) diff --git a/src/macro-external/video/macro-condition-video.cpp b/src/macro-external/video/macro-condition-video.cpp index a65e396a..f64b5cdc 100644 --- a/src/macro-external/video/macro-condition-video.cpp +++ b/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 &)), + this, + SLOT(ColorThresholdChanged(const NumberVariable &))); 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(_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) { diff --git a/src/macro-external/video/macro-condition-video.hpp b/src/macro-external/video/macro-condition-video.hpp index 3f48f147..5c8aac9b 100644 --- a/src/macro-external/video/macro-condition-video.hpp +++ b/src/macro-external/video/macro-condition-video.hpp @@ -117,6 +117,7 @@ public: private slots: void SelectColorClicked(); + void ColorThresholdChanged(const NumberVariable &); 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; diff --git a/src/macro-external/video/opencv-helpers.cpp b/src/macro-external/video/opencv-helpers.cpp index 59436c6c..5dadf443 100644 --- a/src/macro-external/video/opencv-helpers.cpp +++ b/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(y, x) = {0, 0, 0, 255}; + } else { + mat.at(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 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++; } } diff --git a/src/macro-external/video/opencv-helpers.hpp b/src/macro-external/video/opencv-helpers.hpp index cc709619..73c9f933 100644 --- a/src/macro-external/video/opencv-helpers.hpp +++ b/src/macro-external/video/opencv-helpers.hpp @@ -59,8 +59,10 @@ std::vector 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); diff --git a/src/macro-external/video/paramerter-wrappers.cpp b/src/macro-external/video/paramerter-wrappers.cpp index ec82d426..082c2600 100644 --- a/src/macro-external/video/paramerter-wrappers.cpp +++ b/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(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( obs_data_get_int(data, "pageSegMode")); obs_data_release(data); diff --git a/src/macro-external/video/paramerter-wrappers.hpp b/src/macro-external/video/paramerter-wrappers.hpp index 8af28a34..33642311 100644 --- a/src/macro-external/video/paramerter-wrappers.hpp +++ b/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: diff --git a/src/macro-external/video/preview-dialog.cpp b/src/macro-external/video/preview-dialog.cpp index 5a00a391..e7ed92d3 100644 --- a/src/macro-external/video/preview-dialog.cpp +++ b/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 } }