From cdc5d16e95a8e0857460ee7e839587079e4afabc Mon Sep 17 00:00:00 2001 From: warmuptill <19472752+WarmUpTill@users.noreply.github.com> Date: Mon, 12 May 2025 13:21:33 +0200 Subject: [PATCH] Allow changing Tesseract model directory This will prevent custom models being deleted when installing a new version of the plugin, as the plugin's data directory might get wiped on some operating systems. --- data/locale/en-US.ini | 1 + plugins/video/macro-condition-video.cpp | 50 ++++++++++++++++++++--- plugins/video/macro-condition-video.hpp | 5 ++- plugins/video/parameter-wrappers.cpp | 53 +++++++++++++++++++++---- plugins/video/parameter-wrappers.hpp | 7 +++- 5 files changed, 102 insertions(+), 14 deletions(-) diff --git a/data/locale/en-US.ini b/data/locale/en-US.ini index 8f3a7cd0..35e53d84 100644 --- a/data/locale/en-US.ini +++ b/data/locale/en-US.ini @@ -419,6 +419,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.orcBaseDir="Tesseract base directory:{{tesseractBaseDir}}" 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:" diff --git a/plugins/video/macro-condition-video.cpp b/plugins/video/macro-condition-video.cpp index f0dc3730..826002b7 100644 --- a/plugins/video/macro-condition-video.cpp +++ b/plugins/video/macro-condition-video.cpp @@ -281,11 +281,16 @@ void MacroConditionVideo::SetPageSegMode(tesseract::PageSegMode mode) _ocrParameters.SetPageMode(mode); } -bool MacroConditionVideo::SetLanguage(const std::string &language) +bool MacroConditionVideo::SetLanguageCode(const std::string &language) { return _ocrParameters.SetLanguageCode(language); } +bool MacroConditionVideo::SetTesseractBaseDir(const std::string &dir) +{ + return _ocrParameters.SetTesseractBasePath(dir); +} + void MacroConditionVideo::SetCondition(VideoCondition condition) { _condition = condition; @@ -570,6 +575,7 @@ OCREdit::OCREdit(QWidget *parent, PreviewDialog *previewDialog, "AdvSceneSwitcher.condition.video.colorDeviationThresholdDescription"), true)), _pageSegMode(new QComboBox()), + _tesseractBaseDir(new FileSelection(FileSelection::Type::FOLDER)), _languageCode(new VariableLineEdit(this)), _previewDialog(previewDialog), _data(data) @@ -590,13 +596,17 @@ OCREdit::OCREdit(QWidget *parent, PreviewDialog *previewDialog, SLOT(RegexChanged(const RegexConfig &))); QWidget::connect(_pageSegMode, SIGNAL(currentIndexChanged(int)), this, SLOT(PageSegModeChanged(int))); + QWidget::connect(_tesseractBaseDir, + SIGNAL(PathChanged(const QString &)), this, + SLOT(TesseractBaseDirChanged(const QString &))); QWidget::connect(_languageCode, SIGNAL(editingFinished()), this, SLOT(LanguageChanged())); - std::unordered_map widgetPlaceholders = { + const std::unordered_map widgetPlaceholders = { {"{{textColor}}", _textColor}, {"{{selectColor}}", _selectColor}, {"{{textType}}", _pageSegMode}, + {"{{tesseractBaseDir}}", _tesseractBaseDir}, {"{{languageCode}}", _languageCode}, }; @@ -613,6 +623,12 @@ OCREdit::OCREdit(QWidget *parent, PreviewDialog *previewDialog, "AdvSceneSwitcher.condition.video.entry.orcTextType"), pageModeSegLayout, widgetPlaceholders); layout->addLayout(pageModeSegLayout); + auto baseDirLayout = new QHBoxLayout(); + PlaceWidgets( + obs_module_text( + "AdvSceneSwitcher.condition.video.entry.orcBaseDir"), + baseDirLayout, widgetPlaceholders, false); + layout->addLayout(baseDirLayout); auto languageLayout = new QHBoxLayout(); PlaceWidgets( obs_module_text( @@ -634,6 +650,8 @@ OCREdit::OCREdit(QWidget *parent, PreviewDialog *previewDialog, _colorThreshold->SetDoubleValue(_data->_ocrParameters.colorThreshold); _pageSegMode->setCurrentIndex(_pageSegMode->findData( static_cast(_data->_ocrParameters.GetPageMode()))); + _tesseractBaseDir->SetPath( + _data->_ocrParameters.GetTesseractBasePath()); _languageCode->setText(_data->_ocrParameters.GetLanguageCode()); _loading = false; } @@ -722,6 +740,28 @@ void OCREdit::PageSegModeChanged(int idx) _previewDialog->OCRParametersChanged(_data->_ocrParameters); } +void OCREdit::TesseractBaseDirChanged(const QString &path){ + if (_loading || !_data) { + return; + } + + auto lock = LockContext(); + if (!_data->SetTesseractBaseDir(path.toStdString())) { + const QString message(obs_module_text( + "AdvSceneSwitcher.condition.video.ocrLanguageNotFound")); + const QDir dataDir(path); + const QString fileName(_languageCode->text() + ".traineddata"); + DisplayMessage(message.arg(fileName, dataDir.absolutePath())); + + // Reset to previous value + const QSignalBlocker b(this); + _tesseractBaseDir->SetPath(_data->_ocrParameters.GetTesseractBasePath()); + return; + } + _previewDialog->OCRParametersChanged(_data->_ocrParameters); +} + + void OCREdit::LanguageChanged() { if (_loading || !_data) { @@ -729,11 +769,11 @@ void OCREdit::LanguageChanged() } auto lock = LockContext(); - if (!_data->SetLanguage(_languageCode->text().toStdString())) { + if (!_data->SetLanguageCode(_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 QDir dataDir(QString::fromStdString( + _data->_ocrParameters.GetTesseractBasePath())); const QString fileName(_languageCode->text() + ".traineddata"); DisplayMessage(message.arg(fileName, dataDir.absolutePath())); diff --git a/plugins/video/macro-condition-video.hpp b/plugins/video/macro-condition-video.hpp index 925a3831..c6533495 100644 --- a/plugins/video/macro-condition-video.hpp +++ b/plugins/video/macro-condition-video.hpp @@ -47,7 +47,8 @@ public: void ResetLastMatch() { _lastMatchResult = false; } double GetCurrentBrightness() const { return _currentBrightness; } void SetPageSegMode(tesseract::PageSegMode); - bool SetLanguage(const std::string &); + bool SetLanguageCode(const std::string &); + bool SetTesseractBaseDir(const std::string &); void SetCondition(VideoCondition); VideoCondition GetCondition() const { return _condition; } @@ -144,6 +145,7 @@ private slots: void MatchTextChanged(); void RegexChanged(const RegexConfig &conf); void PageSegModeChanged(int); + void TesseractBaseDirChanged(const QString &); void LanguageChanged(); private: @@ -155,6 +157,7 @@ private: QPushButton *_selectColor; SliderSpinBox *_colorThreshold; QComboBox *_pageSegMode; + FileSelection *_tesseractBaseDir; VariableLineEdit *_languageCode; PreviewDialog *_previewDialog; diff --git a/plugins/video/parameter-wrappers.cpp b/plugins/video/parameter-wrappers.cpp index b7cac4b3..8a7c3ec2 100644 --- a/plugins/video/parameter-wrappers.cpp +++ b/plugins/video/parameter-wrappers.cpp @@ -252,6 +252,7 @@ OCRParameters::OCRParameters(const OCRParameters &other) color(other.color), colorThreshold(other.colorThreshold), pageSegMode(other.pageSegMode), + tesseractBasePath(other.tesseractBasePath), languageCode(other.languageCode) { if (!initDone) { @@ -269,6 +270,7 @@ OCRParameters &OCRParameters::operator=(const OCRParameters &other) color = other.color; colorThreshold = other.colorThreshold; pageSegMode = other.pageSegMode; + tesseractBasePath = other.tesseractBasePath; languageCode = other.languageCode; if (!initDone) { Setup(); @@ -284,11 +286,12 @@ bool OCRParameters::Save(obs_data_t *obj) const auto data = obs_data_create(); text.Save(data, "pattern"); regex.Save(data); + tesseractBasePath.Save(data, "tesseractBasePath"); 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_int(data, "version", 2); obs_data_set_obj(obj, "ocrData", data); obs_data_release(data); return true; @@ -301,6 +304,14 @@ bool OCRParameters::Load(obs_data_t *obj) regex.Load(data); obs_data_set_default_string(data, "language", "eng"); languageCode.Load(data, "language"); + tesseractBasePath.Load(data, "tesseractBasePath"); + if (!obs_data_has_user_value(data, "version") || + obs_data_get_int(data, "version") < 2) { + tesseractBasePath = std::string(obs_get_module_data_path( + obs_current_module())) + + "/res/ocr/"; + } + color = LoadColor(data, "textColor"); if (obs_data_has_user_value(data, "version")) { colorThreshold.Load(data, "colorThreshold"); @@ -324,10 +335,8 @@ void OCRParameters::SetPageMode(tesseract::PageSegMode mode) bool OCRParameters::SetLanguageCode(const std::string &value) { - const auto dataPath = - QString(obs_get_module_data_path(obs_current_module())) + - QString("/res/ocr") + "/" + QString::fromStdString(value) + - ".traineddata"; + const auto dataPath = QString::fromStdString(tesseractBasePath) + "/" + + QString::fromStdString(value) + ".traineddata"; QFileInfo fileInfo(dataPath); if (!fileInfo.exists(dataPath)) { return false; @@ -343,6 +352,26 @@ std::string OCRParameters::GetLanguageCode() const return languageCode; } +bool OCRParameters::SetTesseractBasePath(const std::string &value) +{ + const auto dataPath = QString::fromStdString(value) + "/" + + QString::fromStdString(languageCode) + + ".traineddata"; + QFileInfo fileInfo(dataPath); + if (!fileInfo.exists(dataPath)) { + return false; + } + tesseractBasePath = value; + Setup(); + ocr->SetPageSegMode(pageSegMode); + return true; +} + +std::string OCRParameters::GetTesseractBasePath() const +{ + return tesseractBasePath; +} + void OCRParameters::Setup() { ocr = std::make_unique(); @@ -350,8 +379,18 @@ void OCRParameters::Setup() initDone = false; return; } - std::string dataPath = obs_get_module_data_path(obs_current_module()) + - std::string("/res/ocr"); + + const std::string dataPath = std::string(tesseractBasePath) + "/"; + const std::string modelFile = + std::string(languageCode) + ".traineddata"; + const auto fullPath = QString::fromStdString(dataPath) + + QString::fromStdString(modelFile); + QFileInfo fileInfo(fullPath); + if (!fileInfo.exists(fullPath)) { + initDone = false; + return; + } + if (ocr->Init(dataPath.c_str(), languageCode.c_str()) != 0) { initDone = false; return; diff --git a/plugins/video/parameter-wrappers.hpp b/plugins/video/parameter-wrappers.hpp index 6fa0e531..f905e6e9 100644 --- a/plugins/video/parameter-wrappers.hpp +++ b/plugins/video/parameter-wrappers.hpp @@ -93,6 +93,8 @@ public: void SetPageMode(tesseract::PageSegMode); bool SetLanguageCode(const std::string &); std::string GetLanguageCode() const; + bool SetTesseractBasePath(const std::string &); + std::string GetTesseractBasePath() const; tesseract::PageSegMode GetPageMode() const { return pageSegMode; } tesseract::TessBaseAPI *GetOCR() const { return ocr.get(); } @@ -100,12 +102,15 @@ public: RegexConfig regex = RegexConfig::PartialMatchRegexConfig(); QColor color = Qt::black; DoubleVariable colorThreshold = 0.3; - StringVariable languageCode = "eng"; private: void Setup(); tesseract::PageSegMode pageSegMode = tesseract::PSM_SINGLE_BLOCK; + StringVariable tesseractBasePath = + obs_get_module_data_path(obs_current_module()) + + std::string("/res/ocr"); + StringVariable languageCode = "eng"; std::unique_ptr ocr; bool initDone = false; };