From 14bd455bde349fc4a56f7517f4546e508b5fd85f Mon Sep 17 00:00:00 2001 From: wangjue Date: Fri, 29 May 2026 15:38:05 +0800 Subject: [PATCH] Add --screenshot flag for interactive screen region selection --- .gitignore | 6 +++-- main.cpp | 64 ++++++++++++++++++++++++++++++++++++++++++++++++++---- 2 files changed, 64 insertions(+), 6 deletions(-) diff --git a/.gitignore b/.gitignore index 5553f05..2c3fef9 100644 --- a/.gitignore +++ b/.gitignore @@ -1,8 +1,10 @@ .idea .vscode +build +cmake-build-* + DESIGN.md deps -model -build \ No newline at end of file +model \ No newline at end of file diff --git a/main.cpp b/main.cpp index f0496a5..01984e1 100644 --- a/main.cpp +++ b/main.cpp @@ -14,6 +14,7 @@ #include #include #include +#include static volatile bool g_interrupted = false; @@ -96,6 +97,46 @@ static std::vector 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 --mmproj [options] []\n" @@ -111,6 +152,7 @@ static void print_usage(const char * prog) { " -s random seed\n" " --json output in JSON format\n" " --chat-template chat template (default: auto from model)\n" + " --screenshot interactively select screen region to OCR\n" " -h show this help\n" "\n" "If 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) {