wangjue/llm-ocr
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add --screenshot flag for interactive screen region selection

Browse Source
This commit is contained in:
wangjue 2026-05-29 15:38:05 +08:00
parent 7e9a014f4c
commit 14bd455bde
2 changed files with 64 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
.gitignore vendored
View File

@ -1,8 +1,10 @@
.idea
.vscode
build
cmake-build-*
DESIGN.md
deps
model
build
model

64
main.cpp
View File

@ -14,6 +14,7 @@
#include <getopt.h>
#include <signal.h>
#include <inttypes.h>
#include <sys/stat.h>
static volatile bool g_interrupted = false;
@ -96,6 +97,46 @@ static std::vector<unsigned char> read_clipboard() {
return data;
}
static std::string capture_screenshot() {
const char * tmpdir = getenv("TMPDIR");
if (!tmpdir) tmpdir = "/tmp";
std::string path = std::string(tmpdir) + "/ocr_screenshot.png";
fprintf(stderr, "Select a screen region to OCR...\n");
auto check = [&]() -> bool {
struct stat st;
return stat(path.c_str(), &st) == 0 && st.st_size > 0;
};
// 1. gnome-screenshot (blocks until selection complete)
unlink(path.c_str());
if (system(("gnome-screenshot --area -f " + path + " >/dev/null 2>&1").c_str()) == 0 && check())
return path;
// 2. spectacle (KDE, blocks)
unlink(path.c_str());
if (system(("spectacle --region -b -o " + path + " >/dev/null 2>&1").c_str()) == 0 && check())
return path;
// 3. flameshot (writes raw PNG to stdout, blocks)
unlink(path.c_str());
if (system(("flameshot gui -r > " + path + " 2>/dev/null").c_str()) == 0 && check())
return path;
// 4. maim (lightweight X11, blocks)
unlink(path.c_str());
if (system(("maim -s " + path + " 2>/dev/null").c_str()) == 0 && check())
return path;
// 5. slurp + grim (Sway/Wayland, blocks)
unlink(path.c_str());
if (system(("slurp | grim -g - " + path + " 2>/dev/null").c_str()) == 0 && check())
return path;
return "";
}
static void print_usage(const char * prog) {
fprintf(stderr,
"Usage: %s -m <model> --mmproj <mmproj> [options] [<image_path>]\n"
@ -111,6 +152,7 @@ static void print_usage(const char * prog) {
" -s <seed> random seed\n"
" --json output in JSON format\n"
" --chat-template <name> chat template (default: auto from model)\n"
" --screenshot interactively select screen region to OCR\n"
" -h show this help\n"
"\n"
"If <image_path> is omitted, reads image from clipboard.\n"
@ -118,8 +160,9 @@ static void print_usage(const char * prog) {
"\n"
"Examples:\n"
" %s -m model.gguf --mmproj mmproj.gguf -p \"OCR\" image.png\n"
" %s -m model.gguf --mmproj mmproj.gguf # use clipboard image\n",
prog, prog, prog);
" %s -m model.gguf --mmproj mmproj.gguf # use clipboard image\n"
" %s -m model.gguf --mmproj mmproj.gguf --screenshot # select region\n",
prog, prog, prog, prog);
}
int main(int argc, char ** argv) {
@ -128,6 +171,7 @@ int main(int argc, char ** argv) {
std::string image_path;
std::string prompt = "OCR";
std::string chat_template;
bool screenshot_mode = false;
int n_threads = (int)std::thread::hardware_concurrency();
int n_gpu_layers = -1;
int n_ctx = 8192;
@ -151,6 +195,7 @@ int main(int argc, char ** argv) {
{"ngl", required_argument, nullptr, 3},
{"model", required_argument, nullptr, 'm'},
{"chat-template", required_argument, nullptr, 4},
{"screenshot", no_argument, nullptr, 5},
{"help", no_argument, nullptr, 'h'},
{nullptr, 0, nullptr, 0}
};
@ -162,7 +207,8 @@ int main(int argc, char ** argv) {
case 1: json_output = true; break;
case 2: mmproj_path = optarg; break;
case 3: n_gpu_layers = std::stoi(optarg); break;
case 4: chat_template = optarg; break;
case 4: chat_template = optarg; break;
case 5: screenshot_mode = true; break;
case 'm': model_path = optarg; break;
case 'p': prompt = optarg; break;
case 't': n_threads = std::stoi(optarg); break;
@ -178,7 +224,13 @@ int main(int argc, char ** argv) {
return 1;
}
if (optind < argc) {
if (screenshot_mode) {
image_path = capture_screenshot();
if (image_path.empty()) {
fprintf(stderr, "ERROR: no screenshot tool found. Install gnome-screenshot, flameshot, spectacle, maim, or slurp+grim.\n");
return 1;
}
} else if (optind < argc) {
image_path = argv[optind];
}
@ -232,6 +284,7 @@ int main(int argc, char ** argv) {
bmp = mtmd_helper_bitmap_init_from_file(ctx_vision, image_path.c_str());
if (!bmp) {
fprintf(stderr, "ERROR: failed to load image from %s\n", image_path.c_str());
if (screenshot_mode) unlink(image_path.c_str());
mtmd_free(ctx_vision);
llama_free(lctx);
llama_model_free(model);
@ -296,6 +349,7 @@ int main(int argc, char ** argv) {
fprintf(stderr, "ERROR: mtmd_tokenize failed (%d)\n", res);
mtmd_input_chunks_free(chunks);
mtmd_bitmap_free(bmp);
if (screenshot_mode) unlink(image_path.c_str());
mtmd_free(ctx_vision);
llama_free(lctx);
llama_model_free(model);
@ -308,6 +362,7 @@ int main(int argc, char ** argv) {
fprintf(stderr, "ERROR: mtmd_helper_eval_chunks failed (%d)\n", res);
mtmd_input_chunks_free(chunks);
mtmd_bitmap_free(bmp);
if (screenshot_mode) unlink(image_path.c_str());
mtmd_free(ctx_vision);
llama_free(lctx);
llama_model_free(model);
@ -316,6 +371,7 @@ int main(int argc, char ** argv) {
mtmd_input_chunks_free(chunks);
mtmd_bitmap_free(bmp);
if (screenshot_mode) unlink(image_path.c_str());
struct llama_sampler * smpl = nullptr;
if (temp <= 0.0f) {