290 lines
12 KiB
C++
290 lines
12 KiB
C++
#include "OcrUtils.h"
|
||
|
||
#include <QCoreApplication>
|
||
#include <QDebug>
|
||
#include <QDir>
|
||
#include <QElapsedTimer>
|
||
#include <QFile>
|
||
|
||
#include <tesseract/baseapi.h>
|
||
#include <leptonica/allheaders.h>
|
||
|
||
namespace {
|
||
|
||
// Кэш байт rus.traineddata. Грузим файл через Qt (Unicode-aware пути) и
|
||
// отдаём Tesseract в виде in-memory blob через Init(data, size, lang, ...).
|
||
// Это обходит проблему, что на Windows Tesseract открывает traineddata через
|
||
// `fopen`, который использует ANSI-кодировку (cp1251 на русской системе).
|
||
// Если путь к exe содержит кириллицу (например, OneDrive\Рабочий стол) —
|
||
// fopen не находит файл и Init возвращает -1. Memory-init не зависит от пути.
|
||
const QByteArray &tessdataBytes() {
|
||
static QByteArray bytes;
|
||
static bool tried = false;
|
||
if (tried) return bytes;
|
||
tried = true;
|
||
|
||
const QString appDir = QCoreApplication::applicationDirPath();
|
||
const QStringList candidates = {
|
||
appDir + "/tessdata/rus.traineddata", // канонический путь
|
||
appDir + "/rus.traineddata", // legacy: рядом с exe
|
||
};
|
||
for (const QString &path : candidates) {
|
||
QFile f(path);
|
||
if (!f.open(QIODevice::ReadOnly)) continue;
|
||
bytes = f.readAll();
|
||
f.close();
|
||
if (!bytes.isEmpty()) {
|
||
qDebug() << "[OcrUtils] loaded rus.traineddata from" << path
|
||
<< "(" << bytes.size() << "bytes)";
|
||
return bytes;
|
||
}
|
||
}
|
||
qWarning() << "[OcrUtils] rus.traineddata NOT FOUND. Checked:" << candidates
|
||
<< "— OCR будет работать в fallback-режиме без поиска текста.";
|
||
return bytes;
|
||
}
|
||
|
||
|
||
|
||
// Расстояние Левенштейна на QString — нужен для нечёткого сравнения OCR-вывода
|
||
// с искомым словом. Для "Подтверждаю" допускаем 1-2 ошибки.
|
||
int editDistance(const QString &a, const QString &b) {
|
||
const int n = a.size();
|
||
const int m = b.size();
|
||
if (n == 0) return m;
|
||
if (m == 0) return n;
|
||
QVector<int> prev(m + 1), cur(m + 1);
|
||
for (int j = 0; j <= m; ++j) prev[j] = j;
|
||
for (int i = 1; i <= n; ++i) {
|
||
cur[0] = i;
|
||
for (int j = 1; j <= m; ++j) {
|
||
const int cost = a[i - 1] == b[j - 1] ? 0 : 1;
|
||
cur[j] = qMin(qMin(cur[j - 1] + 1, prev[j] + 1), prev[j - 1] + cost);
|
||
}
|
||
prev.swap(cur);
|
||
}
|
||
return prev[m];
|
||
}
|
||
|
||
bool fuzzyMatch(const QString &word, const QString &needle) {
|
||
if (word.isEmpty() || needle.isEmpty()) return false;
|
||
const QString w = word.toLower();
|
||
const QString n = needle.toLower();
|
||
// Substring-ветка работает только если word достаточно длинный.
|
||
// Иначе любая одиночная буква, входящая в needle (для «подтверждаю»:
|
||
// п/о/д/т/в/е/р/ж/а/ю), даст ложный матч на случайный текст экрана —
|
||
// ровно так бот тапал в (620,302) по букве «в» где-то в шапке sheet'a.
|
||
const int minLen = qMax(qMin(n.size(), 3), n.size() * 3 / 5);
|
||
if (w.size() >= minLen && (w.contains(n) || n.contains(w))) return true;
|
||
// Допускаем до 20% ошибок (но не меньше 1) — Tesseract на «Подтверждаю» иногда
|
||
// даёт «Подтверждак»/«Подтверждак)» из-за оранжевого фона и сглаживания.
|
||
const int tolerance = qMax(1, n.size() / 5);
|
||
return editDistance(w, n) <= tolerance;
|
||
}
|
||
|
||
} // namespace
|
||
|
||
namespace OcrUtils {
|
||
|
||
QPoint findTextCenter(const QByteArray &png,
|
||
const QString &needle,
|
||
int yMin,
|
||
int yMax) {
|
||
if (png.isEmpty() || needle.isEmpty()) return {-1, -1};
|
||
|
||
QElapsedTimer timer;
|
||
timer.start();
|
||
|
||
Pix *pix = pixReadMem(reinterpret_cast<const l_uint8 *>(png.constData()),
|
||
static_cast<size_t>(png.size()));
|
||
if (!pix) {
|
||
qWarning() << "[OcrUtils] pixReadMem failed";
|
||
return {-1, -1};
|
||
}
|
||
|
||
const int fullW = pixGetWidth(pix);
|
||
const int fullH = pixGetHeight(pix);
|
||
|
||
int cropY = 0;
|
||
Pix *target = pix;
|
||
Pix *cropped = nullptr;
|
||
if (yMin > 0 || (yMax > 0 && yMax < fullH)) {
|
||
const int y0 = qMax(0, yMin);
|
||
const int y1 = (yMax > 0 && yMax <= fullH) ? yMax : fullH;
|
||
if (y1 > y0) {
|
||
BOX *box = boxCreate(0, y0, fullW, y1 - y0);
|
||
cropped = pixClipRectangle(pix, box, nullptr);
|
||
boxDestroy(&box);
|
||
if (cropped) {
|
||
target = cropped;
|
||
cropY = y0;
|
||
}
|
||
}
|
||
}
|
||
|
||
// Гоняем OCR двумя проходами: оригинал + инверсия. Текст бывает разный —
|
||
// тёмный на светлом (заголовки/поля) и светлый на цветном (sticky-action
|
||
// кнопка, оранжевая/синяя). Инверсия белого-на-оранжевом превращает её в
|
||
// тёмное-на-голубом, что Tesseract разбирает на порядок надёжнее.
|
||
auto runPass = [&](Pix *src, bool invert) -> QPoint {
|
||
Pix *gray = pixConvertTo8(src, 0);
|
||
if (!gray) return {-1, -1};
|
||
Pix *prepared = gray;
|
||
Pix *inverted = nullptr;
|
||
if (invert) {
|
||
inverted = pixInvert(nullptr, gray);
|
||
if (inverted) prepared = inverted;
|
||
}
|
||
|
||
const QByteArray &td = tessdataBytes();
|
||
if (td.isEmpty()) {
|
||
if (inverted) pixDestroy(&inverted);
|
||
pixDestroy(&gray);
|
||
return {-1, -1};
|
||
}
|
||
auto *api = new tesseract::TessBaseAPI();
|
||
// Init(data, size, lang, ...) — грузим traineddata из памяти, чтобы
|
||
// обойти fopen на путях с кириллицей (Windows ANSI-API).
|
||
if (api->Init(td.constData(), td.size(), "rus",
|
||
tesseract::OEM_DEFAULT,
|
||
nullptr, 0, nullptr, nullptr, false, nullptr)) {
|
||
qWarning() << "[OcrUtils] Tesseract Init(memory) failed —"
|
||
<< "traineddata повреждён или несовместимая версия";
|
||
delete api;
|
||
if (inverted) pixDestroy(&inverted);
|
||
pixDestroy(&gray);
|
||
return {-1, -1};
|
||
}
|
||
api->SetImage(prepared);
|
||
api->Recognize(nullptr);
|
||
|
||
QPoint hit(-1, -1);
|
||
tesseract::ResultIterator *ri = api->GetIterator();
|
||
if (ri) {
|
||
do {
|
||
const char *raw = ri->GetUTF8Text(tesseract::RIL_WORD);
|
||
if (!raw) continue;
|
||
const QString word = QString::fromUtf8(raw).trimmed();
|
||
delete[] raw;
|
||
if (!fuzzyMatch(word, needle)) continue;
|
||
|
||
int x1, y1, x2, y2;
|
||
ri->BoundingBox(tesseract::RIL_WORD, &x1, &y1, &x2, &y2);
|
||
const int cx = (x1 + x2) / 2;
|
||
const int cy = (y1 + y2) / 2 + cropY;
|
||
qDebug() << "[OcrUtils] found" << needle << "→" << word
|
||
<< (invert ? "[inverted]" : "[normal]")
|
||
<< "at (" << cx << "," << cy << ")";
|
||
hit = QPoint(cx, cy);
|
||
break;
|
||
} while (ri->Next(tesseract::RIL_WORD));
|
||
}
|
||
api->End();
|
||
delete api;
|
||
if (inverted) pixDestroy(&inverted);
|
||
pixDestroy(&gray);
|
||
return hit;
|
||
};
|
||
|
||
QPoint result = runPass(target, /*invert=*/false);
|
||
if (result.x() < 0) {
|
||
result = runPass(target, /*invert=*/true);
|
||
}
|
||
|
||
if (cropped) pixDestroy(&cropped);
|
||
pixDestroy(&pix);
|
||
|
||
if (result.x() < 0) {
|
||
qDebug() << "[OcrUtils]" << needle << "not found in" << timer.elapsed() << "ms";
|
||
} else {
|
||
qDebug() << "[OcrUtils]" << needle << "located in" << timer.elapsed() << "ms";
|
||
}
|
||
return result;
|
||
}
|
||
|
||
QPoint findColorButtonCenter(const QByteArray &png,
|
||
int targetR,
|
||
int targetG,
|
||
int targetB,
|
||
int tolerance,
|
||
int yMin,
|
||
int yMax) {
|
||
if (png.isEmpty()) return {-1, -1};
|
||
|
||
QElapsedTimer timer;
|
||
timer.start();
|
||
|
||
Pix *src = pixReadMem(reinterpret_cast<const l_uint8 *>(png.constData()),
|
||
static_cast<size_t>(png.size()));
|
||
if (!src) {
|
||
qWarning() << "[OcrUtils] findColorButtonCenter: pixReadMem failed";
|
||
return {-1, -1};
|
||
}
|
||
// Приводим к 32bpp, чтобы pixGetRGBPixel работал независимо от исходной
|
||
// глубины (PNG со скриншота бывает 24/32 bpp).
|
||
Pix *pix = pixConvertTo32(src);
|
||
pixDestroy(&src);
|
||
if (!pix) return {-1, -1};
|
||
|
||
const int W = pixGetWidth(pix);
|
||
const int H = pixGetHeight(pix);
|
||
const int y0 = qBound(0, yMin, H);
|
||
const int y1 = (yMax > 0 && yMax <= H) ? yMax : H;
|
||
|
||
// Для каждой строки считаем число пикселей целевого цвета и их X-диапазон.
|
||
QVector<int> rowCount(H, 0), rowMinX(H, -1), rowMaxX(H, -1);
|
||
for (int y = y0; y < y1; ++y) {
|
||
int cnt = 0, mn = W, mx = -1;
|
||
for (int x = 0; x < W; ++x) {
|
||
l_int32 r, g, b;
|
||
pixGetRGBPixel(pix, x, y, &r, &g, &b);
|
||
if (qAbs(r - targetR) <= tolerance &&
|
||
qAbs(g - targetG) <= tolerance &&
|
||
qAbs(b - targetB) <= tolerance) {
|
||
++cnt;
|
||
if (x < mn) mn = x;
|
||
if (x > mx) mx = x;
|
||
}
|
||
}
|
||
rowCount[y] = cnt;
|
||
rowMinX[y] = mn;
|
||
rowMaxX[y] = mx;
|
||
}
|
||
|
||
// Кнопка — широкая сплошная полоса: ищем самую длинную серию строк, где
|
||
// заполнение целевым цветом ≥ 25% ширины. Мелкие цветные иконки в статус-
|
||
// баре такой порог не проходят.
|
||
const int minRowFill = qMax(1, W / 4);
|
||
int bestStart = -1, bestLen = 0, curStart = -1;
|
||
for (int y = y0; y <= y1; ++y) {
|
||
const bool filled = (y < y1) && (rowCount[y] >= minRowFill);
|
||
if (filled) {
|
||
if (curStart < 0) curStart = y;
|
||
} else if (curStart >= 0) {
|
||
const int len = y - curStart;
|
||
if (len > bestLen) { bestLen = len; bestStart = curStart; }
|
||
curStart = -1;
|
||
}
|
||
}
|
||
|
||
// Полоса должна быть не тоньше ~1% высоты — отсекаем шум/тонкие разделители.
|
||
if (bestStart < 0 || bestLen < qMax(8, H / 100)) {
|
||
pixDestroy(&pix);
|
||
qDebug() << "[OcrUtils] findColorButtonCenter: no button band in"
|
||
<< timer.elapsed() << "ms";
|
||
return {-1, -1};
|
||
}
|
||
|
||
int xmin = W, xmax = -1;
|
||
for (int y = bestStart; y < bestStart + bestLen; ++y) {
|
||
if (rowMinX[y] >= 0) { xmin = qMin(xmin, rowMinX[y]); xmax = qMax(xmax, rowMaxX[y]); }
|
||
}
|
||
pixDestroy(&pix);
|
||
|
||
const QPoint center((xmin + xmax) / 2, bestStart + bestLen / 2);
|
||
qDebug() << "[OcrUtils] findColorButtonCenter → (" << center.x() << ","
|
||
<< center.y() << ") band h=" << bestLen << "in" << timer.elapsed() << "ms";
|
||
return center;
|
||
}
|
||
|
||
} // namespace OcrUtils
|