arc/android/ocr/OcrUtils.cpp

#include "OcrUtils.h"

#include <QCoreApplication>
#include <QDebug>
#include <QElapsedTimer>

#include <tesseract/baseapi.h>
#include <leptonica/allheaders.h>

namespace {

// Расстояние Левенштейна на 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();
    if (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;
        }

        auto *api = new tesseract::TessBaseAPI();
        if (api->Init(QCoreApplication::applicationDirPath().toUtf8().constData(), "rus")) {
            qWarning() << "[OcrUtils] Tesseract init failed (tessdata in"
                       << QCoreApplication::applicationDirPath() << ")";
            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;
}

} // namespace OcrUtils