feat(analysis): add ocrs OCR backend and OCR-driven page text search
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:
@@ -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
201
backend/src/analysis/ocr.rs
Normal 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());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user