427 lines
18 KiB
C++
427 lines
18 KiB
C++
#include "ocr.h"
|
|
#include <memory>
|
|
#include <roapi.h>
|
|
#include <wincodec.h>
|
|
#include <winstring.h>
|
|
|
|
namespace ks
|
|
{
|
|
namespace
|
|
{
|
|
template <class T> struct Com
|
|
{
|
|
T *p{};
|
|
~Com()
|
|
{
|
|
if (p)
|
|
p->Release();
|
|
}
|
|
T *operator->() const
|
|
{
|
|
return p;
|
|
}
|
|
T **out()
|
|
{
|
|
return &p;
|
|
}
|
|
Com() = default;
|
|
Com(const Com &) = delete;
|
|
Com &operator=(const Com &) = delete;
|
|
};
|
|
void check(HRESULT h)
|
|
{
|
|
if (FAILED(h))
|
|
{
|
|
char text[80];
|
|
snprintf(text, sizeof(text), "Windows image/OCR service failed (0x%08lx)", (unsigned long)h);
|
|
throw std::runtime_error(text);
|
|
}
|
|
}
|
|
struct String
|
|
{
|
|
HSTRING value{};
|
|
explicit String(const wchar_t *s)
|
|
{
|
|
check(WindowsCreateString(s, (UINT32)wcslen(s), &value));
|
|
}
|
|
~String()
|
|
{
|
|
WindowsDeleteString(value);
|
|
}
|
|
};
|
|
GUID guid(const wchar_t *text)
|
|
{
|
|
GUID result{};
|
|
check(IIDFromString(text, &result));
|
|
return result;
|
|
}
|
|
template <class T> void factory(const wchar_t *name, const wchar_t *iid, Com<T> &out)
|
|
{
|
|
String s(name);
|
|
check(RoGetActivationFactory(s.value, guid(iid), (void **)out.out()));
|
|
}
|
|
} // namespace
|
|
// Minimal, strongly typed Windows Runtime ABI declarations. No CLR or C++/WinRT runtime is needed.
|
|
// Interface identities/method order follow Microsoft's Windows SDK metadata (Windows.Media.Ocr,
|
|
// Windows.Graphics.Imaging, Windows.Storage.Streams, Windows.Foundation).
|
|
struct Buffer : IInspectable
|
|
{
|
|
virtual HRESULT STDMETHODCALLTYPE Capacity(UINT32 *) = 0;
|
|
virtual HRESULT STDMETHODCALLTYPE Length(UINT32 *) = 0;
|
|
virtual HRESULT STDMETHODCALLTYPE SetLength(UINT32) = 0;
|
|
};
|
|
struct BufferFactory : IInspectable
|
|
{
|
|
virtual HRESULT STDMETHODCALLTYPE Create(UINT32, Buffer **) = 0;
|
|
};
|
|
struct BufferBytes : IUnknown
|
|
{
|
|
virtual HRESULT STDMETHODCALLTYPE Data(BYTE **) = 0;
|
|
};
|
|
struct BitmapStatics : IInspectable
|
|
{
|
|
virtual HRESULT STDMETHODCALLTYPE Copy(IInspectable *, IInspectable **) = 0;
|
|
virtual HRESULT STDMETHODCALLTYPE Convert(IInspectable *, INT32, IInspectable **) = 0;
|
|
virtual HRESULT STDMETHODCALLTYPE ConvertAlpha(IInspectable *, INT32, INT32, IInspectable **) = 0;
|
|
virtual HRESULT STDMETHODCALLTYPE FromBuffer(Buffer *, INT32, INT32, INT32, IInspectable **) = 0;
|
|
virtual HRESULT STDMETHODCALLTYPE FromBufferAlpha(Buffer *, INT32, INT32, INT32, INT32,
|
|
IInspectable **) = 0;
|
|
};
|
|
struct OcrResult : IInspectable
|
|
{
|
|
virtual HRESULT STDMETHODCALLTYPE Lines(IInspectable **) = 0;
|
|
virtual HRESULT STDMETHODCALLTYPE Angle(IInspectable **) = 0;
|
|
virtual HRESULT STDMETHODCALLTYPE Text(HSTRING *) = 0;
|
|
};
|
|
struct OcrOperation : IInspectable
|
|
{
|
|
virtual HRESULT STDMETHODCALLTYPE PutCompleted(IUnknown *) = 0;
|
|
virtual HRESULT STDMETHODCALLTYPE GetCompleted(IUnknown **) = 0;
|
|
virtual HRESULT STDMETHODCALLTYPE Results(OcrResult **) = 0;
|
|
};
|
|
struct AsyncInfo : IInspectable
|
|
{
|
|
virtual HRESULT STDMETHODCALLTYPE Id(UINT32 *) = 0;
|
|
virtual HRESULT STDMETHODCALLTYPE Status(INT32 *) = 0;
|
|
virtual HRESULT STDMETHODCALLTYPE ErrorCode(HRESULT *) = 0;
|
|
virtual HRESULT STDMETHODCALLTYPE Cancel() = 0;
|
|
virtual HRESULT STDMETHODCALLTYPE Close() = 0;
|
|
};
|
|
struct OcrEngine : IInspectable
|
|
{
|
|
virtual HRESULT STDMETHODCALLTYPE Recognize(IInspectable *, OcrOperation **) = 0;
|
|
virtual HRESULT STDMETHODCALLTYPE Language(IInspectable **) = 0;
|
|
};
|
|
struct OcrStatics : IInspectable
|
|
{
|
|
virtual HRESULT STDMETHODCALLTYPE MaxDimension(UINT32 *) = 0;
|
|
virtual HRESULT STDMETHODCALLTYPE Languages(IInspectable **) = 0;
|
|
virtual HRESULT STDMETHODCALLTYPE Supported(IInspectable *, boolean *) = 0;
|
|
virtual HRESULT STDMETHODCALLTYPE FromLanguage(IInspectable *, OcrEngine **) = 0;
|
|
virtual HRESULT STDMETHODCALLTYPE FromProfile(OcrEngine **) = 0;
|
|
};
|
|
Pixels decodeImage(const Clip &clip, UINT maxWidth, UINT maxHeight)
|
|
{
|
|
Bytes encoded;
|
|
for (const auto &f : clip.formats)
|
|
if (f.name == L"PNG")
|
|
{
|
|
encoded = f.data;
|
|
break;
|
|
}
|
|
if (encoded.empty())
|
|
for (const auto &f : clip.formats)
|
|
if ((f.id == CF_DIB || f.id == CF_DIBV5) && f.data.size() >= 40)
|
|
{
|
|
BITMAPINFOHEADER h;
|
|
memcpy(&h, f.data.data(), 40);
|
|
if (h.biSize < 40 || h.biSize > f.data.size())
|
|
continue;
|
|
uint64_t colors = h.biClrUsed ? h.biClrUsed : (h.biBitCount <= 8 ? 1ull << h.biBitCount : 0);
|
|
uint64_t offset = sizeof(BITMAPFILEHEADER) + h.biSize + 4 * colors +
|
|
(h.biSize == 40 && h.biCompression == BI_BITFIELDS ? 12 : 0);
|
|
if (offset > sizeof(BITMAPFILEHEADER) + f.data.size())
|
|
throw std::runtime_error("Invalid bitmap header");
|
|
BITMAPFILEHEADER file{};
|
|
file.bfType = 0x4d42;
|
|
file.bfSize = (DWORD)(sizeof(file) + f.data.size());
|
|
file.bfOffBits = (DWORD)offset;
|
|
encoded.resize(sizeof(file));
|
|
memcpy(encoded.data(), &file, sizeof(file));
|
|
encoded.insert(encoded.end(), f.data.begin(), f.data.end());
|
|
break;
|
|
}
|
|
if (encoded.empty())
|
|
throw std::runtime_error("This clipboard item has no supported image");
|
|
Com<IWICImagingFactory> wic;
|
|
check(CoCreateInstance(CLSID_WICImagingFactory, nullptr, CLSCTX_INPROC_SERVER, IID_IWICImagingFactory,
|
|
(void **)wic.out()));
|
|
Com<IWICStream> stream;
|
|
check(wic->CreateStream(stream.out()));
|
|
check(stream->InitializeFromMemory(encoded.data(), (DWORD)encoded.size()));
|
|
Com<IWICBitmapDecoder> decoder;
|
|
check(wic->CreateDecoderFromStream(stream.p, nullptr, WICDecodeMetadataCacheOnLoad, decoder.out()));
|
|
Com<IWICBitmapFrameDecode> frame;
|
|
check(decoder->GetFrame(0, frame.out()));
|
|
Pixels pixels;
|
|
check(frame->GetSize(&pixels.width, &pixels.height));
|
|
if (!pixels.width || !pixels.height || uint64_t(pixels.width) * pixels.height > 32 * 1024 * 1024)
|
|
throw std::runtime_error("Image is too large to decode safely");
|
|
Com<IWICFormatConverter> converter;
|
|
Com<IWICBitmapScaler> scaler;
|
|
IWICBitmapSource *source = frame.p;
|
|
if (maxWidth && maxHeight && (pixels.width > maxWidth || pixels.height > maxHeight))
|
|
{
|
|
const double ratio = std::min(double(maxWidth) / pixels.width, double(maxHeight) / pixels.height);
|
|
pixels.width = std::max(1u, UINT(pixels.width * ratio));
|
|
pixels.height = std::max(1u, UINT(pixels.height * ratio));
|
|
check(wic->CreateBitmapScaler(scaler.out()));
|
|
check(scaler->Initialize(frame.p, pixels.width, pixels.height, WICBitmapInterpolationModeFant));
|
|
source = scaler.p;
|
|
}
|
|
check(wic->CreateFormatConverter(converter.out()));
|
|
check(converter->Initialize(source, GUID_WICPixelFormat32bppBGRA, WICBitmapDitherTypeNone, nullptr, 0,
|
|
WICBitmapPaletteTypeCustom));
|
|
pixels.bgra.resize(size_t(pixels.width) * pixels.height * 4);
|
|
check(converter->CopyPixels(nullptr, pixels.width * 4, (UINT)pixels.bgra.size(), pixels.bgra.data()));
|
|
return pixels;
|
|
}
|
|
Bytes encodePng(const Pixels &p)
|
|
{
|
|
if (!p.width || !p.height || uint64_t(p.width) * p.height * 4 != p.bgra.size() ||
|
|
p.bgra.size() > UINT_MAX || p.width > UINT_MAX / 4)
|
|
throw std::runtime_error("Invalid PNG image dimensions");
|
|
Com<IWICImagingFactory> factory;
|
|
check(CoCreateInstance(CLSID_WICImagingFactory, nullptr, CLSCTX_INPROC_SERVER,
|
|
IID_PPV_ARGS(factory.out())));
|
|
Com<IStream> stream;
|
|
check(CreateStreamOnHGlobal(nullptr, TRUE, stream.out()));
|
|
Com<IWICBitmapEncoder> encoder;
|
|
check(factory->CreateEncoder(GUID_ContainerFormatPng, nullptr, encoder.out()));
|
|
check(encoder->Initialize(stream.p, WICBitmapEncoderNoCache));
|
|
Com<IWICBitmapFrameEncode> frame;
|
|
check(encoder->CreateNewFrame(frame.out(), nullptr));
|
|
check(frame->Initialize(nullptr));
|
|
check(frame->SetSize(p.width, p.height));
|
|
auto format = GUID_WICPixelFormat32bppBGRA;
|
|
check(frame->SetPixelFormat(&format));
|
|
if (format != GUID_WICPixelFormat32bppBGRA)
|
|
throw std::runtime_error("Unsupported PNG pixel format");
|
|
check(frame->WritePixels(p.height, p.width * 4, (UINT)p.bgra.size(),
|
|
const_cast<BYTE *>(p.bgra.data())));
|
|
check(frame->Commit());
|
|
check(encoder->Commit());
|
|
STATSTG stat{};
|
|
check(stream->Stat(&stat, STATFLAG_NONAME));
|
|
if (stat.cbSize.QuadPart > UINT_MAX)
|
|
throw std::runtime_error("PNG image is too large");
|
|
check(stream->Seek({}, STREAM_SEEK_SET, nullptr));
|
|
Bytes bytes((size_t)stat.cbSize.QuadPart);
|
|
ULONG read{};
|
|
check(stream->Read(bytes.data(), (ULONG)bytes.size(), &read));
|
|
if (read != bytes.size())
|
|
throw std::runtime_error("Cannot read encoded PNG image");
|
|
return bytes;
|
|
}
|
|
Format imageDib(const Pixels &p)
|
|
{
|
|
BITMAPINFOHEADER h{};
|
|
h.biSize = sizeof(h);
|
|
h.biWidth = (LONG)p.width;
|
|
h.biHeight = -(LONG)p.height;
|
|
h.biPlanes = 1;
|
|
h.biBitCount = 32;
|
|
Format format;
|
|
format.id = CF_DIB;
|
|
format.data.resize(sizeof(h));
|
|
memcpy(format.data.data(), &h, sizeof(h));
|
|
format.data.insert(format.data.end(), p.bgra.begin(), p.bgra.end());
|
|
return format;
|
|
}
|
|
std::wstring recognizeText(const Clip &clip)
|
|
{
|
|
check(RoInitialize(RO_INIT_MULTITHREADED));
|
|
struct Apartment
|
|
{
|
|
~Apartment()
|
|
{
|
|
RoUninitialize();
|
|
}
|
|
} apartment;
|
|
auto pixels = decodeImage(clip);
|
|
Com<OcrStatics> ocr;
|
|
factory(L"Windows.Media.Ocr.OcrEngine", L"{5bffa85a-3384-3540-9940-699120d428a8}", ocr);
|
|
UINT32 max{};
|
|
check(ocr->MaxDimension(&max));
|
|
if (pixels.width > max || pixels.height > max)
|
|
throw std::runtime_error("Image exceeds Windows OCR's maximum dimensions");
|
|
Com<OcrEngine> engine;
|
|
check(ocr->FromProfile(engine.out()));
|
|
if (!engine.p)
|
|
throw std::runtime_error("No Windows OCR language is installed for your profile");
|
|
Com<BufferFactory> buffers;
|
|
factory(L"Windows.Storage.Streams.Buffer", L"{71af914d-c10f-484b-bc50-14bc623b3a27}", buffers);
|
|
Com<Buffer> buffer;
|
|
check(buffers->Create((UINT32)pixels.bgra.size(), buffer.out()));
|
|
check(buffer->SetLength((UINT32)pixels.bgra.size()));
|
|
Com<BufferBytes> bytes;
|
|
check(buffer->QueryInterface(guid(L"{905a0fef-bc53-11df-8c49-001e4fc686da}"), (void **)bytes.out()));
|
|
BYTE *memory{};
|
|
check(bytes->Data(&memory));
|
|
memcpy(memory, pixels.bgra.data(), pixels.bgra.size());
|
|
Com<BitmapStatics> bitmaps;
|
|
factory(L"Windows.Graphics.Imaging.SoftwareBitmap", L"{df0385db-672f-4a9d-806e-c2442f343e86}", bitmaps);
|
|
Com<IInspectable> bitmap;
|
|
check(bitmaps->FromBufferAlpha(buffer.p, 87, (INT32)pixels.width, (INT32)pixels.height, 2, bitmap.out()));
|
|
Com<OcrOperation> operation;
|
|
check(engine->Recognize(bitmap.p, operation.out()));
|
|
Com<AsyncInfo> info;
|
|
check(operation->QueryInterface(guid(L"{00000036-0000-0000-c000-000000000046}"), (void **)info.out()));
|
|
auto deadline = GetTickCount64() + 30000;
|
|
INT32 status{};
|
|
do
|
|
{
|
|
check(info->Status(&status));
|
|
if (status == 0)
|
|
{
|
|
if (GetTickCount64() > deadline)
|
|
{
|
|
info->Cancel();
|
|
throw std::runtime_error("Windows OCR timed out");
|
|
}
|
|
Sleep(20);
|
|
}
|
|
} while (status == 0);
|
|
if (status != 1)
|
|
{
|
|
HRESULT failure{};
|
|
check(info->ErrorCode(&failure));
|
|
check(FAILED(failure) ? failure : E_ABORT);
|
|
}
|
|
Com<OcrResult> result;
|
|
check(operation->Results(result.out()));
|
|
HSTRING text{};
|
|
check(result->Text(&text));
|
|
UINT32 length{};
|
|
const wchar_t *raw = WindowsGetStringRawBuffer(text, &length);
|
|
std::wstring out(raw, length);
|
|
WindowsDeleteString(text);
|
|
return out;
|
|
}
|
|
std::filesystem::path tesseractExecutable(const Settings &s)
|
|
{
|
|
if (!s.tesseractPath.empty())
|
|
{
|
|
std::filesystem::path path(s.tesseractPath);
|
|
if (!path.is_absolute()) throw std::runtime_error("Select an absolute path to tesseract.exe");
|
|
return path;
|
|
}
|
|
for (auto variable : {L"ProgramW6432", L"ProgramFiles", L"LOCALAPPDATA"})
|
|
{
|
|
wchar_t base[32768]{};
|
|
auto length = GetEnvironmentVariableW(variable, base, std::size(base));
|
|
if (!length || length >= std::size(base)) continue;
|
|
auto root = std::filesystem::path(base);
|
|
for (auto relative : {L"Tesseract-OCR/tesseract.exe", L"Programs/Tesseract-OCR/tesseract.exe"})
|
|
if (std::filesystem::is_regular_file(root / relative)) return root / relative;
|
|
}
|
|
throw std::runtime_error("Tesseract was not found. Install the optional 64-bit Tesseract 5 package or select tesseract.exe in OCR options.");
|
|
}
|
|
std::wstring ocrProfile(const Settings &s)
|
|
{
|
|
return s.useTesseract ? L"tesseract:" + lower(s.tesseractPath) + L":" + s.tesseractLanguages : L"windows";
|
|
}
|
|
namespace
|
|
{
|
|
struct TessLibrary
|
|
{
|
|
HMODULE module{};
|
|
explicit TessLibrary(const Settings &s)
|
|
{
|
|
auto exe = tesseractExecutable(s);
|
|
if (!std::filesystem::is_regular_file(exe)) throw std::runtime_error("Tesseract executable does not exist");
|
|
auto library = exe.parent_path() / L"libtesseract-5.dll";
|
|
module = LoadLibraryExW(library.c_str(), nullptr,
|
|
LOAD_LIBRARY_SEARCH_DLL_LOAD_DIR | LOAD_LIBRARY_SEARCH_DEFAULT_DIRS);
|
|
if (!module) throw std::runtime_error("Cannot load Tesseract 5. Select a complete 64-bit installation with libtesseract-5.dll and its dependencies.");
|
|
}
|
|
~TessLibrary() { if (module) FreeLibrary(module); }
|
|
template<class T> T symbol(const char *name)
|
|
{
|
|
auto address = GetProcAddress(module, name);
|
|
if (!address) throw std::runtime_error(std::string("Unsupported Tesseract library: missing ") + name);
|
|
T result; static_assert(sizeof(result) == sizeof(address));
|
|
memcpy(&result, &address, sizeof(result));
|
|
return result;
|
|
}
|
|
};
|
|
struct TessEngine
|
|
{
|
|
TessLibrary library;
|
|
void *handle{};
|
|
void (*destroy)(void *){};
|
|
explicit TessEngine(const Settings &s) : library(s)
|
|
{
|
|
auto folder = tesseractExecutable(s).parent_path() / L"tessdata";
|
|
size_t start = 0;
|
|
do
|
|
{
|
|
auto end = s.tesseractLanguages.find(L'+', start);
|
|
auto language = s.tesseractLanguages.substr(start, end == std::wstring::npos ? end : end - start);
|
|
if (!std::filesystem::is_regular_file(folder / (language + L".traineddata")))
|
|
throw std::runtime_error("Tesseract language is not installed: " + utf8(language));
|
|
if (end == std::wstring::npos) break;
|
|
start = end + 1;
|
|
} while (true);
|
|
auto create = library.symbol<void *(*)()>("TessBaseAPICreate");
|
|
destroy = library.symbol<void (*)(void *)>("TessBaseAPIDelete");
|
|
auto init = library.symbol<int (*)(void *, const char *, const char *)>("TessBaseAPIInit3");
|
|
handle = create();
|
|
if (!handle) throw std::runtime_error("Cannot create Tesseract engine");
|
|
if (init(handle, utf8(folder.wstring()).c_str(), utf8(s.tesseractLanguages).c_str()) != 0)
|
|
{
|
|
destroy(handle); handle = nullptr;
|
|
throw std::runtime_error("Cannot initialize the selected Tesseract languages");
|
|
}
|
|
}
|
|
~TessEngine() { if (handle) destroy(handle); }
|
|
};
|
|
}
|
|
void validateTesseract(const Settings &s)
|
|
{
|
|
if (s.useTesseract) { TessEngine engine(s); }
|
|
}
|
|
std::wstring recognizeText(const Clip &clip, const Settings &s)
|
|
{
|
|
if (!s.useTesseract) return recognizeText(clip);
|
|
auto apartmentResult = CoInitializeEx(nullptr, COINIT_MULTITHREADED);
|
|
if (FAILED(apartmentResult) && apartmentResult != RPC_E_CHANGED_MODE) check(apartmentResult);
|
|
struct Apartment { bool owned; ~Apartment() { if (owned) CoUninitialize(); } } apartment{SUCCEEDED(apartmentResult)};
|
|
// Load the optional C API dynamically: Windows OCR users need no Tesseract DLLs.
|
|
TessEngine engine(s);
|
|
auto &lib = engine.library;
|
|
auto image = lib.symbol<void (*)(void *, const unsigned char *, int, int, int, int)>("TessBaseAPISetImage");
|
|
auto mode = lib.symbol<void (*)(void *, int)>("TessBaseAPISetPageSegMode");
|
|
auto recognize = lib.symbol<int (*)(void *, void *)>("TessBaseAPIRecognize");
|
|
auto getText = lib.symbol<char *(*)(void *)>("TessBaseAPIGetUTF8Text");
|
|
auto freeText = lib.symbol<void (*)(const char *)>("TessDeleteText");
|
|
auto createMonitor = lib.symbol<void *(*)()>("TessMonitorCreate");
|
|
auto freeMonitor = lib.symbol<void (*)(void *)>("TessMonitorDelete");
|
|
auto deadline = lib.symbol<void (*)(void *, int)>("TessMonitorSetDeadlineMSecs");
|
|
std::unique_ptr<void, void (*)(void *)> monitor(createMonitor(), freeMonitor);
|
|
if (!monitor) throw std::runtime_error("Cannot create Tesseract progress monitor");
|
|
auto pixels = decodeImage(clip);
|
|
Bytes gray(size_t(pixels.width) * pixels.height);
|
|
for (size_t i = 0; i < gray.size(); ++i)
|
|
gray[i] = BYTE((29 * pixels.bgra[4*i] + 150 * pixels.bgra[4*i+1] + 77 * pixels.bgra[4*i+2]) >> 8);
|
|
image(engine.handle, gray.data(), pixels.width, pixels.height, 1, pixels.width);
|
|
mode(engine.handle, 11); // Sparse text: screenshots often contain separate labels and paragraphs.
|
|
deadline(monitor.get(), 30000);
|
|
if (recognize(engine.handle, monitor.get()) != 0)
|
|
throw std::runtime_error("Tesseract recognition failed or exceeded 30 seconds");
|
|
std::unique_ptr<char, void (*)(const char *)> text(getText(engine.handle), freeText);
|
|
if (!text) throw std::runtime_error("Tesseract returned no result");
|
|
return wide(text.get());
|
|
}
|
|
} // namespace ks
|