feat(analysis): add ocrs OCR backend and OCR-driven page text search
All checks were successful
deploy / test-backend (push) Successful in 28m40s
deploy / test-frontend (push) Successful in 10m36s
deploy / build-and-push (push) Successful in 11m8s
deploy / deploy (push) Successful in 12s

Adds an in-process ocrs OCR backend as the active analysis engine
(vision left dormant), enables OCR text search on the page-tag
aggregation endpoints, and reshapes the admin analysis/settings UI to
present the OCR-only surface. Bumps version 0.89.0 -> 0.90.0.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
MechaCat02
2026-06-30 19:52:43 +02:00
parent 4fb98e4a1e
commit 83c2899373
24 changed files with 1335 additions and 443 deletions

View File

@@ -9,5 +9,6 @@
pub mod daemon;
pub mod events;
pub mod ocr;
pub mod prompt;
pub mod vision;

201
backend/src/analysis/ocr.rs Normal file
View File

@@ -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<Vec<String>>;
}
/// 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<String>) -> 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<Self> {
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<Vec<String>> {
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<dyn Storage>,
pub engine: Arc<dyn OcrEngine>,
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<String>,
}
impl StubOcrEngine {
pub fn new(lines: &[&str]) -> Arc<Self> {
Arc::new(Self { lines: lines.iter().map(|s| s.to_string()).collect() })
}
}
impl OcrEngine for StubOcrEngine {
fn recognize(&self, _image: &[u8]) -> anyhow::Result<Vec<String>> {
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());
}
}