Enable use of other OCR language models

This commit is contained in:
WarmUpTill
2023-05-29 18:13:51 +02:00
committed by WarmUpTill
parent 339418b5f3
commit f506600b8c
5 changed files with 73 additions and 1 deletions

View File

@@ -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="<html><head/><body><p>Language model for \"%1\" was not found in:<br>\"%2\"<br><br>You can find additional language models <a href=\"https://github.com/tesseract-ocr/tessdata\"><span style=\" text-decoration: underline; color:#268bd2;\">here</span></a>."
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:"

View File

@@ -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<std::string, QWidget *> 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<int>(_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<MacroConditionVideo> &data)

View File

@@ -9,6 +9,7 @@
#include <screenshot-helper.hpp>
#include <slider-spinbox.hpp>
#include <variable-text-edit.hpp>
#include <variable-line-edit.hpp>
#include <QWidget>
#include <QComboBox>
@@ -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;

View File

@@ -1,5 +1,7 @@
#include "paramerter-wrappers.hpp"
#include <filesystem>
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<int>(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<tesseract::PageSegMode>(
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<tesseract::TessBaseAPI>();
@@ -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;
}

View File

@@ -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();