From f506600b8cddc1b0981b59a071ccac44e015dc37 Mon Sep 17 00:00:00 2001 From: WarmUpTill Date: Mon, 29 May 2023 18:13:51 +0200 Subject: [PATCH] Enable use of other OCR language models --- data/locale/en-US.ini | 2 + .../video/macro-condition-video.cpp | 39 +++++++++++++++++++ .../video/macro-condition-video.hpp | 4 ++ .../video/paramerter-wrappers.cpp | 26 ++++++++++++- .../video/paramerter-wrappers.hpp | 3 ++ 5 files changed, 73 insertions(+), 1 deletion(-) diff --git a/data/locale/en-US.ini b/data/locale/en-US.ini index 22aaf580..d45487c3 100644 --- a/data/locale/en-US.ini +++ b/data/locale/en-US.ini @@ -246,6 +246,7 @@ AdvSceneSwitcher.condition.video.ocrMode.circleWord="Single word in a circle" AdvSceneSwitcher.condition.video.ocrMode.singleChar="Single character" AdvSceneSwitcher.condition.video.ocrMode.sparseText="Find text in no particular order" AdvSceneSwitcher.condition.video.ocrMode.sparseTextOSD="Find text in no particular order (OSD)" +AdvSceneSwitcher.condition.video.ocrLanguageNotFound="

Language model for \"%1\" was not found in:
\"%2\"

You can find additional language models here." AdvSceneSwitcher.condition.video.colorMatchThreshold="Match threshold:" AdvSceneSwitcher.condition.video.colorMatchThresholdDescription="How much of the image has to match the given color?\nA value of 1 requires every pixel of the input image to match the given color." AdvSceneSwitcher.condition.video.colorDeviationThreshold="Color deviation:" @@ -261,6 +262,7 @@ AdvSceneSwitcher.condition.video.entry.checkAreaEnable="Perform check only in ar AdvSceneSwitcher.condition.video.entry.checkArea="{{checkAreaEnable}}{{checkArea}}{{selectArea}}" AdvSceneSwitcher.condition.video.entry.orcColorPick="Check for text color:{{textColor}}{{selectColor}}" AdvSceneSwitcher.condition.video.entry.orcTextType="Check for text type:{{textType}}" +AdvSceneSwitcher.condition.video.entry.orcLanguage="Check for language:{{languageCode}}" AdvSceneSwitcher.condition.video.entry.color="Check for color:{{color}}{{selectColor}}" AdvSceneSwitcher.condition.video.minSize="Minimum size:" AdvSceneSwitcher.condition.video.maxSize="Maximum size:" diff --git a/src/macro-external/video/macro-condition-video.cpp b/src/macro-external/video/macro-condition-video.cpp index efb7df73..6d3b8fed 100644 --- a/src/macro-external/video/macro-condition-video.cpp +++ b/src/macro-external/video/macro-condition-video.cpp @@ -254,6 +254,11 @@ void MacroConditionVideo::SetPageSegMode(tesseract::PageSegMode mode) _ocrParameters.SetPageMode(mode); } +bool MacroConditionVideo::SetLanguage(const std::string &language) +{ + return _ocrParameters.SetLanguageCode(language); +} + bool MacroConditionVideo::ScreenshotContainsPattern() { cv::Mat result; @@ -454,6 +459,7 @@ OCREdit::OCREdit(QWidget *parent, PreviewDialog *previewDialog, _selectColor(new QPushButton(obs_module_text( "AdvSceneSwitcher.condition.video.selectColor"))), _pageSegMode(new QComboBox()), + _languageCode(new VariableLineEdit(this)), _previewDialog(previewDialog), _data(data) { @@ -467,11 +473,14 @@ OCREdit::OCREdit(QWidget *parent, PreviewDialog *previewDialog, SLOT(RegexChanged(RegexConfig))); QWidget::connect(_pageSegMode, SIGNAL(currentIndexChanged(int)), this, SLOT(PageSegModeChanged(int))); + QWidget::connect(_languageCode, SIGNAL(editingFinished()), this, + SLOT(LanguageChanged())); std::unordered_map widgetPlaceholders = { {"{{textColor}}", _textColor}, {"{{selectColor}}", _selectColor}, {"{{textType}}", _pageSegMode}, + {"{{languageCode}}", _languageCode}, }; auto layout = new QVBoxLayout(); @@ -487,6 +496,12 @@ OCREdit::OCREdit(QWidget *parent, PreviewDialog *previewDialog, "AdvSceneSwitcher.condition.video.entry.orcTextType"), pageModeSegLayout, widgetPlaceholders); layout->addLayout(pageModeSegLayout); + auto languageLayout = new QHBoxLayout(); + PlaceWidgets( + obs_module_text( + "AdvSceneSwitcher.condition.video.entry.orcLanguage"), + languageLayout, widgetPlaceholders); + layout->addLayout(languageLayout); auto colorPickLayout = new QHBoxLayout(); PlaceWidgets( obs_module_text( @@ -500,6 +515,7 @@ OCREdit::OCREdit(QWidget *parent, PreviewDialog *previewDialog, SetupColorLabel(_data->_ocrParameters.color); _pageSegMode->setCurrentIndex(_pageSegMode->findData( static_cast(_data->_ocrParameters.GetPageMode()))); + _languageCode->setText(_data->_ocrParameters.GetLanguageCode()); _loading = false; } @@ -575,6 +591,29 @@ void OCREdit::PageSegModeChanged(int idx) _previewDialog->OCRParametersChanged(_data->_ocrParameters); } +void OCREdit::LanguageChanged() +{ + if (_loading || !_data) { + return; + } + + auto lock = LockContext(); + if (!_data->SetLanguage(_languageCode->text().toStdString())) { + const QString message(obs_module_text( + "AdvSceneSwitcher.condition.video.ocrLanguageNotFound")); + const QDir dataDir( + obs_get_module_data_path(obs_current_module())); + const QString fileName(_languageCode->text() + ".traineddata"); + DisplayMessage(message.arg(fileName, dataDir.absolutePath())); + + // Reset to previous value + const QSignalBlocker b(this); + _languageCode->setText(_data->_ocrParameters.GetLanguageCode()); + return; + } + _previewDialog->OCRParametersChanged(_data->_ocrParameters); +} + ObjectDetectEdit::ObjectDetectEdit( QWidget *parent, PreviewDialog *previewDialog, const std::shared_ptr &data) diff --git a/src/macro-external/video/macro-condition-video.hpp b/src/macro-external/video/macro-condition-video.hpp index 3f65d296..6f9ef508 100644 --- a/src/macro-external/video/macro-condition-video.hpp +++ b/src/macro-external/video/macro-condition-video.hpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include @@ -42,6 +43,7 @@ public: void ResetLastMatch() { _lastMatchResult = false; } double GetCurrentBrightness() const { return _currentBrightness; } void SetPageSegMode(tesseract::PageSegMode); + bool SetLanguage(const std::string &); VideoInput _video; VideoCondition _condition = VideoCondition::MATCH; @@ -118,6 +120,7 @@ private slots: void MatchTextChanged(); void RegexChanged(RegexConfig conf); void PageSegModeChanged(int); + void LanguageChanged(); private: void SetupColorLabel(const QColor &); @@ -127,6 +130,7 @@ private: QLabel *_textColor; QPushButton *_selectColor; QComboBox *_pageSegMode; + VariableLineEdit *_languageCode; PreviewDialog *_previewDialog; diff --git a/src/macro-external/video/paramerter-wrappers.cpp b/src/macro-external/video/paramerter-wrappers.cpp index b68451e3..ec82d426 100644 --- a/src/macro-external/video/paramerter-wrappers.cpp +++ b/src/macro-external/video/paramerter-wrappers.cpp @@ -1,5 +1,7 @@ #include "paramerter-wrappers.hpp" +#include + namespace advss { bool PatternMatchParameters::Save(obs_data_t *obj) const @@ -259,8 +261,10 @@ bool OCRParameters::Save(obs_data_t *obj) const auto data = obs_data_create(); text.Save(data, "pattern"); regex.Save(data); + languageCode.Save(data, "language"); SaveColor(data, "textColor", color); obs_data_set_int(data, "pageSegMode", static_cast(pageSegMode)); + obs_data_set_int(data, "version", 1); obs_data_set_obj(obj, "ocrData", data); obs_data_release(data); return true; @@ -282,6 +286,8 @@ bool OCRParameters::Load(obs_data_t *obj) auto data = obs_data_get_obj(obj, "ocrData"); text.Load(data, "pattern"); regex.Load(data); + obs_data_set_default_string(data, "language", "eng"); + languageCode.Load(data, "language"); color = LoadColor(data, "textColor"); pageSegMode = static_cast( obs_data_get_int(data, "pageSegMode")); @@ -299,6 +305,24 @@ void OCRParameters::SetPageMode(tesseract::PageSegMode mode) ocr->SetPageSegMode(mode); } +bool OCRParameters::SetLanguageCode(const std::string &value) +{ + std::string dataPath = obs_get_module_data_path(obs_current_module()) + + std::string("/res/ocr") + "/" + value + + ".traineddata"; + if (!std::filesystem::exists(dataPath)) { + return false; + } + Setup(); + ocr->SetPageSegMode(pageSegMode); + return true; +} + +std::string OCRParameters::GetLanguageCode() const +{ + return languageCode; +} + void OCRParameters::Setup() { ocr = std::make_unique(); @@ -308,7 +332,7 @@ void OCRParameters::Setup() } std::string dataPath = obs_get_module_data_path(obs_current_module()) + std::string("/res/ocr"); - if (ocr->Init(dataPath.c_str(), "eng") != 0) { + if (ocr->Init(dataPath.c_str(), languageCode.c_str()) != 0) { initDone = false; return; } diff --git a/src/macro-external/video/paramerter-wrappers.hpp b/src/macro-external/video/paramerter-wrappers.hpp index 4894c3e6..8af28a34 100644 --- a/src/macro-external/video/paramerter-wrappers.hpp +++ b/src/macro-external/video/paramerter-wrappers.hpp @@ -90,12 +90,15 @@ public: bool Initialized() const { return initDone; } void SetPageMode(tesseract::PageSegMode); + bool SetLanguageCode(const std::string &); + std::string GetLanguageCode() const; tesseract::PageSegMode GetPageMode() const { return pageSegMode; } tesseract::TessBaseAPI *GetOCR() const { return ocr.get(); } StringVariable text = obs_module_text("AdvSceneSwitcher.enterText"); RegexConfig regex = RegexConfig::PartialMatchRegexConfig(); QColor color = Qt::black; + StringVariable languageCode = "eng"; private: void Setup();