diff --git a/.env.example b/.env.example index 33d7c57..01f8079 100644 --- a/.env.example +++ b/.env.example @@ -277,6 +277,19 @@ CRAWLER_TZ=UTC # ANALYSIS_ENABLED Turn the worker on at first boot. Toggleable live in # the dashboard. Default `false`. ANALYSIS_ENABLED=false +# ANALYSIS_BACKEND Which engine the worker runs. env-ONLY (deploy-time). +# `ocr` (default) = the in-process ocrs engine: fast, +# CPU-only, text-only, ideal for a Pi. `vision` = the local +# LLM at ANALYSIS_VISION_URL (full OCR + tags + scene + +# safety, but heavy). The ANALYSIS_VISION_* / ANALYSIS_API_KEY +# knobs below only apply to `vision`. +ANALYSIS_BACKEND=ocr +# OCRS_DETECTION_MODEL / OCRS_RECOGNITION_MODEL Paths to the ocrs `.rten` +# text-detection / -recognition models (only read when ANALYSIS_BACKEND=ocr). +# The backend image bakes both into /models, so the defaults work unchanged; +# override only to point at custom-trained models. +OCRS_DETECTION_MODEL=/models/text-detection.rten +OCRS_RECOGNITION_MODEL=/models/text-recognition.rten # ANALYSIS_VISION_URL /v1/chat/completions endpoint. Required when enabled. # For the bundled vision container, use # http://mangalord-vision:8000/v1/chat/completions. diff --git a/backend/Cargo.lock b/backend/Cargo.lock index 77e9e38..34c09f0 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -214,6 +214,12 @@ version = "1.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + [[package]] name = "bitflags" version = "2.11.1" @@ -514,6 +520,25 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "crossbeam-deque" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +dependencies = [ + "crossbeam-utils", +] + [[package]] name = "crossbeam-queue" version = "0.3.12" @@ -638,7 +663,7 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e0e367e4e7da84520dedcac1901e4da967309406d1e51017ae1abfb97adbd38" dependencies = [ - "bitflags", + "bitflags 2.11.1", "objc2", ] @@ -772,6 +797,16 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" +[[package]] +name = "flatbuffers" +version = "24.12.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4f1baf0dbf96932ec9a3038d57900329c015b0bfb7b63d904f3bc27e2b02a096" +dependencies = [ + "bitflags 1.3.2", + "rustc_version", +] + [[package]] name = "flate2" version = "1.1.9" @@ -1059,6 +1094,12 @@ version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" +[[package]] +name = "hermit-abi" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" + [[package]] name = "hex" version = "0.4.3" @@ -1448,7 +1489,7 @@ version = "0.1.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e02f3bb43d335493c96bf3fd3a321600bf6bd07ed34bc64118e9293bdffea46c" dependencies = [ - "bitflags", + "bitflags 2.11.1", "libc", "plain", "redox_syscall 0.7.5", @@ -1517,7 +1558,7 @@ checksum = "c41e0c4fef86961ac6d6f8a82609f55f31b05e4fce149ac5710e439df7619ba4" [[package]] name = "mangalord" -version = "0.89.0" +version = "0.90.0" dependencies = [ "anyhow", "argon2", @@ -1537,8 +1578,10 @@ dependencies = [ "infer", "mime", "nix 0.29.0", + "ocrs", "rand 0.8.6", "reqwest", + "rten", "scraper", "serde", "serde_json", @@ -1669,7 +1712,7 @@ version = "0.29.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "71e2746dc3a24dd78b3cfcb7be93368c6de9963d30f43a6a73998a9cf4b17b46" dependencies = [ - "bitflags", + "bitflags 2.11.1", "cfg-if", "cfg_aliases", "libc", @@ -1681,7 +1724,7 @@ version = "0.31.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf20d2fde8ff38632c426f1165ed7436270b44f199fc55284c38276f9db47c3d" dependencies = [ - "bitflags", + "bitflags 2.11.1", "cfg-if", "cfg_aliases", "libc", @@ -1757,6 +1800,16 @@ dependencies = [ "libm", ] +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + [[package]] name = "objc2" version = "0.6.4" @@ -1772,7 +1825,7 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "73ad74d880bb43877038da939b7427bba67e9dd42004a18b809ba7d87cee241c" dependencies = [ - "bitflags", + "bitflags 2.11.1", "objc2", "objc2-foundation", ] @@ -1793,7 +1846,7 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" dependencies = [ - "bitflags", + "bitflags 2.11.1", "dispatch2", "objc2", ] @@ -1804,7 +1857,7 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e022c9d066895efa1345f8e33e584b9f958da2fd4cd116792e15e07e4720a807" dependencies = [ - "bitflags", + "bitflags 2.11.1", "dispatch2", "objc2", "objc2-core-foundation", @@ -1837,7 +1890,7 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0cde0dfb48d25d2b4862161a4d5fcc0e3c24367869ad306b0c9ec0073bfed92d" dependencies = [ - "bitflags", + "bitflags 2.11.1", "objc2", "objc2-core-foundation", "objc2-core-graphics", @@ -1855,7 +1908,7 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3e0adef53c21f888deb4fa59fc59f7eb17404926ee8a6f59f5df0fd7f9f3272" dependencies = [ - "bitflags", + "bitflags 2.11.1", "block2", "libc", "objc2", @@ -1868,7 +1921,7 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "180788110936d59bab6bd83b6060ffdfffb3b922ba1396b312ae795e1de9d81d" dependencies = [ - "bitflags", + "bitflags 2.11.1", "objc2", "objc2-core-foundation", ] @@ -1879,7 +1932,7 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "96c1358452b371bf9f104e21ec536d37a650eb10f7ee379fff67d2e08d537f1f" dependencies = [ - "bitflags", + "bitflags 2.11.1", "objc2", "objc2-core-foundation", "objc2-foundation", @@ -1891,7 +1944,7 @@ version = "0.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d87d638e33c06f577498cbcc50491496a3ed4246998a7fbba7ccb98b1e7eab22" dependencies = [ - "bitflags", + "bitflags 2.11.1", "block2", "objc2", "objc2-cloud-kit", @@ -1916,6 +1969,21 @@ dependencies = [ "objc2-foundation", ] +[[package]] +name = "ocrs" +version = "0.12.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5379fdd3f11522b5a2ff53017a189463dabf5d0a9c915cb3eb97fabec4ea11c" +dependencies = [ + "anyhow", + "rayon", + "rten", + "rten-imageproc", + "rten-tensor", + "thiserror 2.0.18", + "wasm-bindgen", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -2142,7 +2210,7 @@ version = "0.18.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "60769b8b31b2a9f263dae2776c37b1b28ae246943cf719eb6946a1db05128a61" dependencies = [ - "bitflags", + "bitflags 2.11.1", "crc32fast", "fdeflate", "flate2", @@ -2361,13 +2429,33 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + [[package]] name = "redox_syscall" version = "0.5.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" dependencies = [ - "bitflags", + "bitflags 2.11.1", ] [[package]] @@ -2376,7 +2464,7 @@ version = "0.7.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4666a1a60d8412eab19d94f6d13dcc9cea0a5ef4fdf6a5db306537413c661b1b" dependencies = [ - "bitflags", + "bitflags 2.11.1", ] [[package]] @@ -2496,19 +2584,135 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rten" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43c230fa4ade87c913f61dbd911b7eb0d49460ceff3f1e4fabc837fac191137c" +dependencies = [ + "flatbuffers", + "num_cpus", + "rayon", + "rten-base", + "rten-gemm", + "rten-model-file", + "rten-onnx", + "rten-shape-inference", + "rten-simd", + "rten-tensor", + "rten-vecmath", + "rustc-hash", + "smallvec", + "typeid", + "wasm-bindgen", +] + +[[package]] +name = "rten-base" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2738cf8bb4c27f828ac788d01ccf4e367e8e773cfec6851f81851b5211de6a79" +dependencies = [ + "rayon", +] + +[[package]] +name = "rten-gemm" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a81a0ca209fb5ce21bd17efa0bd287d5881c6cebfbff0b21c4294a1a14a9e" +dependencies = [ + "rayon", + "rten-base", + "rten-simd", + "rten-tensor", +] + +[[package]] +name = "rten-imageproc" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d5f148e7e941fb5727b9046a5fa1b45525543d5105f14b384fd9261df0ee49bc" +dependencies = [ + "rten-tensor", +] + +[[package]] +name = "rten-model-file" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2f8d270f07ab1bbfff47250c6039f6caa5da59d6da7d74f66aa48559aa6fea" +dependencies = [ + "flatbuffers", + "rten-base", +] + +[[package]] +name = "rten-onnx" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23086eef75bfb55278cb0b45cf9f5a877d466d914914aafebee4ffca9b24d20c" + +[[package]] +name = "rten-shape-inference" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e8a913c7ca40e2bfbb2a0cd447cce56b33ab19435f56693271a2ef37cf58984" +dependencies = [ + "rten-tensor", + "smallvec", +] + +[[package]] +name = "rten-simd" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b19a0032dfcb70dd20960c1c51a37674b237586cbc1ce586f45b46605d108e82" + +[[package]] +name = "rten-tensor" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05dc744a270aa32d154f1a3df8e48740ccc1be9dfbcf23295ada66d83aa98de6" +dependencies = [ + "rayon", + "rten-base", + "smallvec", + "typeid", +] + +[[package]] +name = "rten-vecmath" +version = "0.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9574ddebf5671bc08ceb76e2e1638fadc57fdeff318634eab2c29e9a803cff64" +dependencies = [ + "rten-base", + "rten-simd", +] + [[package]] name = "rustc-hash" version = "2.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" +[[package]] +name = "rustc_version" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" +dependencies = [ + "semver", +] + [[package]] name = "rustix" version = "0.38.44" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" dependencies = [ - "bitflags", + "bitflags 2.11.1", "errno", "libc", "linux-raw-sys 0.4.15", @@ -2521,7 +2725,7 @@ version = "1.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" dependencies = [ - "bitflags", + "bitflags 2.11.1", "errno", "libc", "linux-raw-sys 0.12.1", @@ -2603,7 +2807,7 @@ version = "0.25.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4eb30575f3638fc8f6815f448d50cb1a2e255b0897985c8c59f4d37b72a07b06" dependencies = [ - "bitflags", + "bitflags 2.11.1", "cssparser", "derive_more", "fxhash", @@ -2911,7 +3115,7 @@ checksum = "aa003f0038df784eb8fecbbac13affe3da23b45194bd57dba231c8f48199c526" dependencies = [ "atoi", "base64", - "bitflags", + "bitflags 2.11.1", "byteorder", "bytes", "chrono", @@ -2955,7 +3159,7 @@ checksum = "db58fcd5a53cf07c184b154801ff91347e4c30d17a3562a635ff028ad5deda46" dependencies = [ "atoi", "base64", - "bitflags", + "bitflags 2.11.1", "byteorder", "chrono", "crc", @@ -3317,7 +3521,7 @@ version = "0.6.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68d6fdd9f81c2819c9a8b0e0cd91660e7746a8e6ea2ba7c6b2b057985f6bcb51" dependencies = [ - "bitflags", + "bitflags 2.11.1", "bytes", "futures-util", "http", @@ -3428,6 +3632,12 @@ dependencies = [ "utf-8", ] +[[package]] +name = "typeid" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc7d623258602320d5c55d1bc22793b57daff0ec7efc270ea7d55ce1d5f5471c" + [[package]] name = "typenum" version = "1.20.0" @@ -3668,7 +3878,7 @@ version = "0.244.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "47b807c72e1bac69382b3a6fb3dbe8ea4c0ed87ff5629b8685ae6b9a611028fe" dependencies = [ - "bitflags", + "bitflags 2.11.1", "hashbrown 0.15.5", "indexmap", "semver", @@ -4096,7 +4306,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9d66ea20e9553b30172b5e831994e35fbde2d165325bec84fc43dbf6f4eb9cb2" dependencies = [ "anyhow", - "bitflags", + "bitflags 2.11.1", "indexmap", "log", "serde", diff --git a/backend/Cargo.toml b/backend/Cargo.toml index 765d388..81ba63c 100644 --- a/backend/Cargo.toml +++ b/backend/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "mangalord" -version = "0.89.0" +version = "0.90.0" edition = "2021" default-run = "mangalord" @@ -52,6 +52,8 @@ sysinfo = { version = "0.32", default-features = false, features = ["system", "c nix = { version = "0.29", features = ["fs"] } scraper = "0.20" reqwest = { version = "0.12", default-features = false, features = ["rustls-tls", "socks", "cookies", "stream", "json"] } +ocrs = "0.12" +rten = "0.24" [dev-dependencies] tempfile = "3" diff --git a/backend/Dockerfile b/backend/Dockerfile index 065269e..7bb0e19 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -58,6 +58,20 @@ WORKDIR /app COPY --from=builder /app/target/release/mangalord /usr/local/bin/mangalord COPY --from=builder /app/migrations /app/migrations +# OCR models for the default `ANALYSIS_BACKEND=ocr` (ocrs) engine. The two +# `.rten` files are pulled at build time into /models, where the runtime's +# `OCRS_DETECTION_MODEL` / `OCRS_RECOGNITION_MODEL` defaults point. They're a +# few MB each and pure data (no native code), so they bake cleanly into the +# image and need no manual setup on the Pi. Set INSTALL_OCR_MODELS=false to +# skip (e.g. for a vision-only deploy that never runs ocrs). +ARG INSTALL_OCR_MODELS=true +ARG OCRS_MODEL_BASE_URL=https://ocrs-models.s3-accelerate.amazonaws.com +RUN if [ "$INSTALL_OCR_MODELS" = "true" ]; then \ + mkdir -p /models \ + && curl -fsSL "${OCRS_MODEL_BASE_URL}/text-detection.rten" -o /models/text-detection.rten \ + && curl -fsSL "${OCRS_MODEL_BASE_URL}/text-recognition.rten" -o /models/text-recognition.rten; \ + fi + ENV STORAGE_DIR=/var/lib/mangalord/storage # Pre-create the storage dir so the entrypoint doesn't need to # mkdir-as-root and so the named volume mount inherits the right diff --git a/backend/src/analysis/mod.rs b/backend/src/analysis/mod.rs index 5ea2d7c..2c909ee 100644 --- a/backend/src/analysis/mod.rs +++ b/backend/src/analysis/mod.rs @@ -9,5 +9,6 @@ pub mod daemon; pub mod events; +pub mod ocr; pub mod prompt; pub mod vision; diff --git a/backend/src/analysis/ocr.rs b/backend/src/analysis/ocr.rs new file mode 100644 index 0000000..e60556b --- /dev/null +++ b/backend/src/analysis/ocr.rs @@ -0,0 +1,201 @@ +//! The in-process OCR analysis backend (the `ocrs` engine). +//! +//! A lightweight alternative to [`crate::analysis::vision`]: instead of a slow +//! local LLM, each page is run through `ocrs` — a pure-Rust detect→recognize +//! OCR pipeline on the `rten` runtime. It extracts **text only** (no tags, +//! scene description or NSFW flags — those stay the vision backend's job), then +//! reuses [`repo::page_analysis::persist_analysis`] so the OCR lines land in +//! `page_ocr_text` and the weighted `search_doc` tsvector exactly as the vision +//! path produces them. That makes the existing text-search surfaces +//! (`/v1/me/page-search` and the tag aggregations) work with no further wiring. +//! +//! The engine is split behind the [`OcrEngine`] trait so the dispatcher is +//! unit-testable without shipping the (multi-MB) `.rten` model files: tests use +//! [`test_support::StubOcrEngine`], production uses [`OcrsEngine`]. + +use std::sync::Arc; + +use async_trait::async_trait; +use sqlx::PgPool; +use uuid::Uuid; + +use crate::analysis::daemon::AnalyzeDispatcher; +use crate::domain::page_analysis::{OcrResult, SafetyFlag, VisionAnalysis}; +use crate::repo; +use crate::storage::Storage; + +/// The `model` label stamped onto `page_analysis` rows written by this backend. +pub const OCR_MODEL_LABEL: &str = "ocrs"; + +/// Extracts text lines from a decoded-or-encoded page image. The production +/// impl ([`OcrsEngine`]) decodes the bytes itself; the trait takes the raw +/// stored image bytes so the dispatcher stays engine-agnostic. +pub trait OcrEngine: Send + Sync { + /// Run OCR over one page image (the bytes as stored, e.g. PNG/JPEG/WebP). + /// Returns the recognized text lines in reading order (top→bottom). + fn recognize(&self, image: &[u8]) -> anyhow::Result>; +} + +/// Turn the OCR engine's ordered text lines into the [`VisionAnalysis`] shape +/// that [`repo::page_analysis::persist_analysis`] consumes. OCR-only: the tag, +/// scene and safety fields are left empty/default. The line `kind` is left +/// blank — `persist_analysis` maps an empty kind to the neutral mid-weight +/// `OcrKind::Narration` bucket (classifying speech/sfx/… is the deferred +/// vision backend's job). +pub fn lines_to_analysis(lines: Vec) -> VisionAnalysis { + let ocr_results = lines + .into_iter() + .map(|text| OcrResult { text, kind: String::new(), y: None }) + .collect(); + VisionAnalysis { + ocr_results, + tagging_results: Vec::new(), + scene_description: String::new(), + safety_flag: SafetyFlag::default(), + } +} + +/// Production OCR engine: an `ocrs` detect+recognize pipeline with the two +/// `.rten` models loaded once at startup. Cheap to share across workers — the +/// recognize path borrows `&self`. +pub struct OcrsEngine { + engine: ocrs::OcrEngine, +} + +impl OcrsEngine { + /// Load the detection + recognition models from disk and build the engine. + /// Fails (at startup) if either model file is missing or unreadable, so a + /// misconfigured path is a loud boot error rather than a per-page failure. + pub fn from_model_paths(detection: &str, recognition: &str) -> anyhow::Result { + use anyhow::Context; + let detection_model = rten::Model::load_file(detection) + .with_context(|| format!("load ocrs detection model {detection}"))?; + let recognition_model = rten::Model::load_file(recognition) + .with_context(|| format!("load ocrs recognition model {recognition}"))?; + let engine = ocrs::OcrEngine::new(ocrs::OcrEngineParams { + detection_model: Some(detection_model), + recognition_model: Some(recognition_model), + ..Default::default() + }) + .context("construct ocrs engine")?; + Ok(Self { engine }) + } +} + +impl OcrEngine for OcrsEngine { + fn recognize(&self, image: &[u8]) -> anyhow::Result> { + use anyhow::Context; + // Decode to RGB8 so `ImageSource` gets a known channel layout. + let rgb = image::load_from_memory(image) + .context("decode page image for OCR")? + .into_rgb8(); + let source = ocrs::ImageSource::from_bytes(rgb.as_raw(), rgb.dimensions()) + .map_err(|e| anyhow::anyhow!("build OCR image source: {e}"))?; + let input = self.engine.prepare_input(source)?; + // detect words → group into lines → recognize each line. Mirrors + // `OcrEngine::get_text`, but keeps the lines as a Vec instead of + // joining them, so each becomes its own `page_ocr_text` row. + let words = self.engine.detect_words(&input)?; + let line_rects = self.engine.find_text_lines(&input, &words); + let lines = self + .engine + .recognize_text(&input, &line_rects)? + .into_iter() + .filter_map(|line| line.map(|l| l.to_string())) + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()) + .collect(); + Ok(lines) + } +} + +/// Production dispatcher for the OCR backend: load the page, read its image +/// from storage, run OCR on the blocking pool, and persist the lines. Mirrors +/// [`crate::analysis::daemon::RealAnalyzeDispatcher`] but with no network I/O. +pub struct OcrAnalyzeDispatcher { + pub db: PgPool, + pub storage: Arc, + pub engine: Arc, + pub max_image_bytes: usize, +} + +#[async_trait] +impl AnalyzeDispatcher for OcrAnalyzeDispatcher { + async fn dispatch(&self, page_id: Uuid) -> anyhow::Result<()> { + let Some(page) = repo::page::find_by_id(&self.db, page_id).await? else { + // Page was deleted between enqueue and dispatch — nothing to do. + return Ok(()); + }; + let bytes = self + .storage + .get(&page.storage_key) + .await + .map_err(|e| anyhow::anyhow!("read page image {}: {e}", page.storage_key))?; + if bytes.len() > self.max_image_bytes { + anyhow::bail!( + "page image {} is {} bytes, over the {} cap", + page.storage_key, + bytes.len(), + self.max_image_bytes + ); + } + // OCR inference is CPU-bound and synchronous — keep it off the async + // worker's runtime thread. + let engine = Arc::clone(&self.engine); + let lines = tokio::task::spawn_blocking(move || engine.recognize(&bytes)) + .await + .map_err(|e| anyhow::anyhow!("OCR task join error: {e}"))??; + let analysis = lines_to_analysis(lines); + repo::page_analysis::persist_analysis(&self.db, page_id, &analysis, OCR_MODEL_LABEL).await?; + Ok(()) + } +} + +/// Stubs for the OCR dispatcher's integration tests. Public because the tests +/// live in the `tests/` dir (a separate crate). +pub mod test_support { + use super::*; + + /// An [`OcrEngine`] that returns a fixed set of lines regardless of input, + /// so the dispatcher's storage→persist path can be tested without models. + pub struct StubOcrEngine { + pub lines: Vec, + } + + impl StubOcrEngine { + pub fn new(lines: &[&str]) -> Arc { + Arc::new(Self { lines: lines.iter().map(|s| s.to_string()).collect() }) + } + } + + impl OcrEngine for StubOcrEngine { + fn recognize(&self, _image: &[u8]) -> anyhow::Result> { + Ok(self.lines.clone()) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn lines_to_analysis_maps_lines_in_order_and_leaves_rest_empty() { + let v = lines_to_analysis(vec!["Hello".to_string(), "world!".to_string()]); + assert_eq!(v.ocr_results.len(), 2); + assert_eq!(v.ocr_results[0].text, "Hello"); + assert_eq!(v.ocr_results[1].text, "world!"); + // OCR-only: kind blank (→ Narration at persist), no tags/scene/safety. + assert!(v.ocr_results.iter().all(|r| r.kind.is_empty())); + assert!(v.tagging_results.is_empty()); + assert_eq!(v.scene_description, ""); + assert!(!v.safety_flag.is_nsfw); + assert!(v.safety_flag.content_type.is_empty()); + } + + #[test] + fn lines_to_analysis_handles_no_text() { + let v = lines_to_analysis(Vec::new()); + assert!(v.ocr_results.is_empty()); + } +} diff --git a/backend/src/app.rs b/backend/src/app.rs index 54ba854..2d69bd0 100644 --- a/backend/src/app.rs +++ b/backend/src/app.rs @@ -429,46 +429,74 @@ async fn spawn_analysis_daemon( ), Err(e) => tracing::warn!(?e, "analysis: reclaim_orphaned at startup failed"), } - let http = reqwest::Client::builder() - .timeout(cfg.request_timeout) - // Refuse to honour ambient HTTP_PROXY / HTTPS_PROXY container env. - // The vision call carries an env-managed bearer token + page image - // bytes; a stray upstream proxy would exfiltrate both. Mirrors the - // crawler client's `.no_proxy()` (see `spawn_crawler_daemon`). - .no_proxy() - .build() - .context("build analysis http client")?; - let vision = crate::analysis::vision::VisionClient::new(http, cfg); - let dispatcher = Arc::new(crate::analysis::daemon::RealAnalyzeDispatcher { - db: db.clone(), - storage, - vision, - model: cfg.model.clone(), - max_image_bytes: cfg.max_image_bytes, - }); - // When a readiness URL is configured, gate leasing on it so an - // autoscaler that idle-stops the vision container never lets a job burn - // its retries. A dedicated short-timeout client keeps the probe snappy - // and independent of the (long) per-request analysis timeout. - let readiness: Option> = - match &cfg.vision_health_url { - Some(url) if !url.is_empty() => { - let probe = reqwest::Client::builder() - .timeout(std::time::Duration::from_secs(5)) - // Same reasoning as the main analysis client: do not - // honour ambient HTTP_PROXY env. The readiness probe is - // unauthenticated but a hostile upstream still gets a - // useful side-channel on backend uptime + vision health. - .no_proxy() - .build() - .context("build vision readiness http client")?; - Some(Arc::new(crate::analysis::daemon::HttpVisionReadiness { - http: probe, - health_url: url.clone(), - })) - } - _ => None, - }; + // Pick the engine. The OCR backend runs in-process (no network, no + // readiness gate); the vision backend talks to a local LLM server. + let (dispatcher, readiness): ( + Arc, + Option>, + ) = match cfg.backend { + crate::config::AnalysisBackend::Ocr => { + // Load the `.rten` models once; a bad path is a loud boot error. + let engine = crate::analysis::ocr::OcrsEngine::from_model_paths( + &cfg.ocr_detection_model, + &cfg.ocr_recognition_model, + ) + .context("build ocrs engine")?; + let dispatcher = Arc::new(crate::analysis::ocr::OcrAnalyzeDispatcher { + db: db.clone(), + storage, + engine: Arc::new(engine), + max_image_bytes: cfg.max_image_bytes, + }); + // In-process engine is always ready — no gate. + (dispatcher, None) + } + crate::config::AnalysisBackend::Vision => { + let http = reqwest::Client::builder() + .timeout(cfg.request_timeout) + // Refuse to honour ambient HTTP_PROXY / HTTPS_PROXY container + // env. The vision call carries an env-managed bearer token + + // page image bytes; a stray upstream proxy would exfiltrate + // both. Mirrors the crawler client's `.no_proxy()` (see + // `spawn_crawler_daemon`). + .no_proxy() + .build() + .context("build analysis http client")?; + let vision = crate::analysis::vision::VisionClient::new(http, cfg); + let dispatcher = Arc::new(crate::analysis::daemon::RealAnalyzeDispatcher { + db: db.clone(), + storage, + vision, + model: cfg.model.clone(), + max_image_bytes: cfg.max_image_bytes, + }); + // When a readiness URL is configured, gate leasing on it so an + // autoscaler that idle-stops the vision container never lets a job + // burn its retries. A dedicated short-timeout client keeps the + // probe snappy and independent of the (long) per-request timeout. + let readiness: Option> = + match &cfg.vision_health_url { + Some(url) if !url.is_empty() => { + let probe = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(5)) + // Same reasoning as the main analysis client: do + // not honour ambient HTTP_PROXY env. The readiness + // probe is unauthenticated but a hostile upstream + // still gets a useful side-channel on backend + // uptime + vision health. + .no_proxy() + .build() + .context("build vision readiness http client")?; + Some(Arc::new(crate::analysis::daemon::HttpVisionReadiness { + http: probe, + health_url: url.clone(), + })) + } + _ => None, + }; + (dispatcher, readiness) + } + }; let handle = crate::analysis::daemon::spawn( db, CancellationToken::new(), @@ -480,7 +508,12 @@ async fn spawn_analysis_daemon( readiness, }, ); - tracing::info!(workers = cfg.workers, model = %cfg.model, "analysis worker daemon started"); + tracing::info!( + workers = cfg.workers, + backend = ?cfg.backend, + model = %cfg.model, + "analysis worker daemon started" + ); Ok(handle) } @@ -1447,6 +1480,10 @@ mod tests { let storage: Arc = Arc::new(LocalStorage::new(storage_dir.path())); let mut cfg = crate::config::AnalysisConfig::default(); + // Use the vision backend so the daemon doesn't try to load the ocrs + // `.rten` models (absent in unit CI). This test only exercises the + // pre-worker lease reclaim, which is engine-agnostic; no dispatch runs. + cfg.backend = crate::config::AnalysisBackend::Vision; cfg.workers = 1; cfg.job_timeout = Duration::from_secs(1); let events = Arc::new(crate::analysis::events::AnalysisEvents::new()); diff --git a/backend/src/config.rs b/backend/src/config.rs index b9cc936..4f76170 100644 --- a/backend/src/config.rs +++ b/backend/src/config.rs @@ -113,6 +113,32 @@ impl ResponseFormat { } } +/// Which engine the analysis worker dispatches each page through. A +/// deploy-time choice (the engine is either installed or not), so it lives in +/// env only and is not part of the admin-editable `AnalysisSettings`. +/// +/// * `Ocr` — in-process [`crate::analysis::ocr`] (the `ocrs` engine): fast, +/// CPU-only, English text. Writes OCR text only (no tags/scene/safety). +/// * `Vision` — the local OpenAI-compatible LLM in [`crate::analysis::vision`]: +/// full OCR + tags + scene + safety, but heavy. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum AnalysisBackend { + Ocr, + Vision, +} + +impl AnalysisBackend { + /// Lenient env parse. Defaults to `Ocr` (the Pi-friendly path) for any + /// unset or unrecognized value; `vision` opts back into the LLM engine. + fn from_str(s: &str) -> AnalysisBackend { + match s.trim().to_lowercase().as_str() { + "vision" | "llm" => AnalysisBackend::Vision, + // Default (incl. "ocr", "ocrs", and anything unrecognized). + _ => AnalysisBackend::Ocr, + } + } +} + /// AI content-analysis worker configuration: the enable gate, the local /// OpenAI-compatible vision endpoint, and the worker / request knobs. #[derive(Clone, Debug)] @@ -120,6 +146,16 @@ pub struct AnalysisConfig { /// Master switch (`ANALYSIS_ENABLED`). When `false`, no analysis jobs /// are enqueued and no worker runs. Defaults to `false`. pub enabled: bool, + /// Which engine the worker dispatches through (`ANALYSIS_BACKEND`): + /// `ocr` (default, the in-process `ocrs` engine) or `vision` (the local + /// LLM). Deploy-time, env-only — see [`AnalysisBackend`]. + pub backend: AnalysisBackend, + /// Path to the ocrs text-*detection* `.rten` model + /// (`OCRS_DETECTION_MODEL`). Only read when `backend == Ocr`. + pub ocr_detection_model: String, + /// Path to the ocrs text-*recognition* `.rten` model + /// (`OCRS_RECOGNITION_MODEL`). Only read when `backend == Ocr`. + pub ocr_recognition_model: String, /// Number of concurrent analysis workers (`ANALYSIS_WORKERS`). pub workers: usize, /// OpenAI-compatible chat/completions URL (`ANALYSIS_VISION_URL`). @@ -195,6 +231,9 @@ impl Default for AnalysisConfig { fn default() -> Self { Self { enabled: false, + backend: AnalysisBackend::Ocr, + ocr_detection_model: "/models/text-detection.rten".to_string(), + ocr_recognition_model: "/models/text-recognition.rten".to_string(), workers: 1, endpoint: "http://localhost:8000/v1/chat/completions".to_string(), vision_health_url: None, @@ -227,6 +266,17 @@ impl AnalysisConfig { let d = AnalysisConfig::default(); Self { enabled: env_bool("ANALYSIS_ENABLED", d.enabled), + backend: std::env::var("ANALYSIS_BACKEND") + .map(|s| AnalysisBackend::from_str(&s)) + .unwrap_or(d.backend), + ocr_detection_model: std::env::var("OCRS_DETECTION_MODEL") + .ok() + .filter(|s| !s.is_empty()) + .unwrap_or(d.ocr_detection_model), + ocr_recognition_model: std::env::var("OCRS_RECOGNITION_MODEL") + .ok() + .filter(|s| !s.is_empty()) + .unwrap_or(d.ocr_recognition_model), workers: env_usize("ANALYSIS_WORKERS", d.workers).max(1), endpoint: std::env::var("ANALYSIS_VISION_URL").unwrap_or(d.endpoint), vision_health_url: std::env::var("ANALYSIS_VISION_HEALTH_URL") @@ -734,11 +784,18 @@ mod tests { "ANALYSIS_MAX_SLICES", "ANALYSIS_RESPONSE_FORMAT", "ANALYSIS_FREQUENCY_PENALTY", + "ANALYSIS_BACKEND", + "OCRS_DETECTION_MODEL", + "OCRS_RECOGNITION_MODEL", ] { std::env::remove_var(k); } let cfg = AnalysisConfig::from_env(); assert!(!cfg.enabled); + // OCR is the default engine (the Pi-friendly path). + assert_eq!(cfg.backend, AnalysisBackend::Ocr); + assert_eq!(cfg.ocr_detection_model, "/models/text-detection.rten"); + assert_eq!(cfg.ocr_recognition_model, "/models/text-recognition.rten"); assert_eq!(cfg.workers, 1); assert_eq!(cfg.max_pixels, 1_000_000); assert_eq!(cfg.min_slice_height, 640); @@ -765,6 +822,36 @@ mod tests { std::env::remove_var("ANALYSIS_RESPONSE_FORMAT"); } + #[test] + fn analysis_backend_parses_and_defaults_to_ocr() { + let _g = ENV_GUARD.lock().unwrap_or_else(|p| p.into_inner()); + for (raw, want) in [ + ("ocr", AnalysisBackend::Ocr), + ("ocrs", AnalysisBackend::Ocr), + ("vision", AnalysisBackend::Vision), + ("llm", AnalysisBackend::Vision), + ("anything-else", AnalysisBackend::Ocr), + ] { + std::env::set_var("ANALYSIS_BACKEND", raw); + assert_eq!(AnalysisConfig::from_env().backend, want, "raw={raw}"); + } + // Unset → OCR. + std::env::remove_var("ANALYSIS_BACKEND"); + assert_eq!(AnalysisConfig::from_env().backend, AnalysisBackend::Ocr); + } + + #[test] + fn ocr_model_paths_parse_from_env() { + let _g = ENV_GUARD.lock().unwrap_or_else(|p| p.into_inner()); + std::env::set_var("OCRS_DETECTION_MODEL", "/opt/det.rten"); + std::env::set_var("OCRS_RECOGNITION_MODEL", "/opt/rec.rten"); + let cfg = AnalysisConfig::from_env(); + std::env::remove_var("OCRS_DETECTION_MODEL"); + std::env::remove_var("OCRS_RECOGNITION_MODEL"); + assert_eq!(cfg.ocr_detection_model, "/opt/det.rten"); + assert_eq!(cfg.ocr_recognition_model, "/opt/rec.rten"); + } + #[test] fn analysis_config_parses_from_env() { let _g = ENV_GUARD.lock().unwrap_or_else(|p| p.into_inner()); diff --git a/backend/src/settings.rs b/backend/src/settings.rs index 925b4ed..c0414cc 100644 --- a/backend/src/settings.rs +++ b/backend/src/settings.rs @@ -459,6 +459,11 @@ impl AnalysisSettings { api_key: base.api_key.clone(), // Env-only readiness probe URL preserved from the base. vision_health_url: base.vision_health_url.clone(), + // Deploy-time engine selection + ocrs model paths: env-only, not + // admin-tunable, so carry them through from the base unchanged. + backend: base.backend, + ocr_detection_model: base.ocr_detection_model.clone(), + ocr_recognition_model: base.ocr_recognition_model.clone(), }) } } diff --git a/backend/tests/analysis_ocr.rs b/backend/tests/analysis_ocr.rs new file mode 100644 index 0000000..e3c40bc --- /dev/null +++ b/backend/tests/analysis_ocr.rs @@ -0,0 +1,109 @@ +//! Integration tests for the OCR analysis backend +//! (`analysis::ocr::OcrAnalyzeDispatcher`). A stub OCR engine stands in for +//! `ocrs` (whose `.rten` models aren't shipped to CI), so these pin the +//! storage→OCR→persist wiring: the dispatcher reads the page image, runs the +//! engine, and persists the lines via the shared `persist_analysis` path — +//! landing `page_ocr_text` rows and a populated `search_doc` exactly like the +//! vision backend. Each `#[sqlx::test]` gets a fresh migrated DB. + +mod common; + +use std::sync::Arc; + +use mangalord::analysis::daemon::AnalyzeDispatcher; +use mangalord::analysis::ocr::test_support::StubOcrEngine; +use mangalord::analysis::ocr::OcrAnalyzeDispatcher; +use mangalord::domain::page_analysis::AnalysisStatus; +use mangalord::repo; +use mangalord::storage::{LocalStorage, Storage}; +use sqlx::PgPool; +use tempfile::TempDir; +use uuid::Uuid; + +/// Seed a manga → chapter → page chain whose page points at `storage_key`, +/// and return the page id. +async fn seed_page(pool: &PgPool, storage_key: &str) -> Uuid { + let manga_id: Uuid = + sqlx::query_scalar("INSERT INTO mangas (title) VALUES ('M') RETURNING id") + .fetch_one(pool) + .await + .unwrap(); + let chapter_id: Uuid = sqlx::query_scalar( + "INSERT INTO chapters (manga_id, number) VALUES ($1, 1) RETURNING id", + ) + .bind(manga_id) + .fetch_one(pool) + .await + .unwrap(); + sqlx::query_scalar( + "INSERT INTO pages (chapter_id, page_number, storage_key, content_type) \ + VALUES ($1, 1, $2, 'image/png') RETURNING id", + ) + .bind(chapter_id) + .bind(storage_key) + .fetch_one(pool) + .await + .unwrap() +} + +fn ocr_dispatcher( + pool: &PgPool, + storage: Arc, + lines: &[&str], +) -> OcrAnalyzeDispatcher { + OcrAnalyzeDispatcher { + db: pool.clone(), + storage, + engine: StubOcrEngine::new(lines), + max_image_bytes: 8 * 1024 * 1024, + } +} + +#[sqlx::test(migrations = "./migrations")] +async fn dispatch_persists_ocr_lines_and_search_doc(pool: PgPool) { + let dir = TempDir::new().unwrap(); + let storage: Arc = Arc::new(LocalStorage::new(dir.path())); + let key = "mangas/x/p1.png"; + storage.put(key, &common::fake_png_bytes()).await.unwrap(); + let page_id = seed_page(&pool, key).await; + + let dispatcher = ocr_dispatcher(&pool, Arc::clone(&storage), &["Hello there", "general"]); + dispatcher.dispatch(page_id).await.unwrap(); + + // Two OCR rows, in order, with the recognized text. + let rows: Vec<(String, i32)> = sqlx::query_as( + "SELECT text, ord FROM page_ocr_text WHERE page_id = $1 ORDER BY ord", + ) + .bind(page_id) + .fetch_all(&pool) + .await + .unwrap(); + assert_eq!(rows.len(), 2); + assert_eq!(rows[0].0, "Hello there"); + assert_eq!(rows[1].0, "general"); + + // The analysis row is `done`, stamped with the ocrs model label, and has a + // non-empty tsvector so text search works. + let row = repo::page_analysis::load(&pool, page_id).await.unwrap().unwrap(); + assert_eq!(row.status, AnalysisStatus::Done); + assert_eq!(row.model.as_deref(), Some("ocrs")); + let has_doc: bool = sqlx::query_scalar( + "SELECT search_doc IS NOT NULL AND search_doc != ''::tsvector \ + FROM page_analysis WHERE page_id = $1", + ) + .bind(page_id) + .fetch_one(&pool) + .await + .unwrap(); + assert!(has_doc, "search_doc must be populated from OCR text"); +} + +#[sqlx::test(migrations = "./migrations")] +async fn dispatch_missing_page_is_noop(pool: PgPool) { + let dir = TempDir::new().unwrap(); + let storage: Arc = Arc::new(LocalStorage::new(dir.path())); + // A page id that was never inserted — the dispatcher must treat it as a + // deleted page and succeed without writing anything. + let dispatcher = ocr_dispatcher(&pool, storage, &["whatever"]); + dispatcher.dispatch(Uuid::new_v4()).await.unwrap(); +} diff --git a/docker-compose.yml b/docker-compose.yml index 64ebdc9..16a05d5 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -143,8 +143,13 @@ services: # when both are present. Leave unset for the bundled # HashedControlPassword path above. CRAWLER_TOR_CONTROL_COOKIE_PATH: ${CRAWLER_TOR_CONTROL_COOKIE_PATH:-} - # Analysis worker (vision) — boot seeds + env-only API key. + # Analysis worker — boot seeds + env-only API key. ANALYSIS_ENABLED: ${ANALYSIS_ENABLED:-false} + # Engine selection (deploy-time): `ocr` (in-process ocrs, default) or + # `vision` (local LLM). Model paths default to the image-baked /models. + ANALYSIS_BACKEND: ${ANALYSIS_BACKEND:-ocr} + OCRS_DETECTION_MODEL: ${OCRS_DETECTION_MODEL:-/models/text-detection.rten} + OCRS_RECOGNITION_MODEL: ${OCRS_RECOGNITION_MODEL:-/models/text-recognition.rten} ANALYSIS_VISION_URL: ${ANALYSIS_VISION_URL:-} ANALYSIS_VISION_MODEL: ${ANALYSIS_VISION_MODEL:-} ANALYSIS_WORKERS: ${ANALYSIS_WORKERS:-1} diff --git a/frontend/package.json b/frontend/package.json index 50639c3..5478056 100644 --- a/frontend/package.json +++ b/frontend/package.json @@ -1,6 +1,6 @@ { "name": "mangalord-frontend", - "version": "0.89.0", + "version": "0.90.0", "private": true, "type": "module", "scripts": {