feat(analysis): slice tall pages to the pixel budget for reliable OCR
Very tall webtoon pages were squashed by the fixed longest-edge downscale, so OCR was garbage. The vision client now slices to the model's pixel budget at native resolution and assembles one result: - plan_slices: native-resolution bands sized to max_pixels (portrait OR landscape), only when height/width > tall_aspect_threshold; min_slice_ height guards pathologically wide pages (reduce width); max_slices caps the count (coarser fallback). Normal pages still take one combined call. - Two-pass, OCR-first: Pass A OCRs each band (near-native), merge_ocr stitches them with seam-scoped fuzzy dedup (keep the longer transcript, preserve order/kind, don't collapse non-adjacent repeats); Pass B feeds the whole downscaled image + the merged OCR text to get tags/scene/safety grounded in the real dialogue. Only OCR is merged. - Config: replace max_image_dim with max_pixels / min_slice_height / slice_overlap / tall_aspect_threshold / max_slices (+ env_f64); bump job_timeout default 180->600 (N sequential slice calls per page); raise MAX_OCR_PIECES 60->200. New prompts/schemas: OCR_PROMPT/ocr_json_schema, GROUNDING_PROMPT/grounding_json_schema. Tests: plan_slices (single/portrait/landscape/pathological-wide/cap + coverage/overlap), merge_ocr (seam dedup, keep-longer, non-adjacent repeats, order), render_whole/render_slice, the new request builders, and updated config defaults/env. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
2
backend/Cargo.lock
generated
2
backend/Cargo.lock
generated
@@ -1517,7 +1517,7 @@ checksum = "c41e0c4fef86961ac6d6f8a82609f55f31b05e4fce149ac5710e439df7619ba4"
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "mangalord"
|
name = "mangalord"
|
||||||
version = "0.77.1"
|
version = "0.78.0"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"anyhow",
|
"anyhow",
|
||||||
"argon2",
|
"argon2",
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "mangalord"
|
name = "mangalord"
|
||||||
version = "0.77.1"
|
version = "0.78.0"
|
||||||
edition = "2021"
|
edition = "2021"
|
||||||
default-run = "mangalord"
|
default-run = "mangalord"
|
||||||
|
|
||||||
|
|||||||
@@ -22,8 +22,10 @@ pub const MIN_TAGS: usize = 5;
|
|||||||
pub const MAX_TAGS: usize = 10;
|
pub const MAX_TAGS: usize = 10;
|
||||||
|
|
||||||
/// Output bounds enforced at sanitize time so one verbose response can't
|
/// Output bounds enforced at sanitize time so one verbose response can't
|
||||||
/// bloat the row or the search document.
|
/// bloat the row or the search document. The OCR cap is generous because a
|
||||||
pub const MAX_OCR_PIECES: usize = 60;
|
/// long sliced page legitimately accumulates many text pieces (the tsvector
|
||||||
|
/// handles it).
|
||||||
|
pub const MAX_OCR_PIECES: usize = 200;
|
||||||
pub const MAX_OCR_TEXT_CHARS: usize = 500;
|
pub const MAX_OCR_TEXT_CHARS: usize = 500;
|
||||||
pub const MAX_SCENE_CHARS: usize = 1000;
|
pub const MAX_SCENE_CHARS: usize = 1000;
|
||||||
|
|
||||||
@@ -50,6 +52,81 @@ pub const SYSTEM_PROMPT: &str = concat!(
|
|||||||
"content. Output must be valid minified JSON and nothing else."
|
"content. Output must be valid minified JSON and nothing else."
|
||||||
);
|
);
|
||||||
|
|
||||||
|
// --- Two-pass prompts (long-page slicing) -----------------------------------
|
||||||
|
|
||||||
|
/// Pass A: OCR-only over a single vertical slice of a tall page. Kept
|
||||||
|
/// narrow so each slice call is fast and never truncates.
|
||||||
|
pub const OCR_PROMPT: &str = concat!(
|
||||||
|
"You are an OCR engine for manga/comic pages. The image is one vertical ",
|
||||||
|
"slice of a larger page. Transcribe EVERY visible text element verbatim ",
|
||||||
|
"into ocr_results, each with its kind ",
|
||||||
|
"(speech|thought|narration|sfx|title|caption), in top-to-bottom order. ",
|
||||||
|
"If there is no text, return {\"ocr_results\":[]}. Output ONLY one ",
|
||||||
|
"minified JSON object, no prose, no markdown."
|
||||||
|
);
|
||||||
|
|
||||||
|
/// Pass B: tags + scene + safety for the whole page, grounded in the merged
|
||||||
|
/// OCR text (supplied as a separate user message part by the client).
|
||||||
|
pub const GROUNDING_PROMPT: &str = concat!(
|
||||||
|
"You are a manga/comic page analyzer. You are given the page image ",
|
||||||
|
"(possibly downscaled) AND the OCR text already extracted from it. Using ",
|
||||||
|
"both, return ONLY one minified JSON object with: tagging_results (5-10 ",
|
||||||
|
"short lowercase content tags — characters, actions, setting, mood, ",
|
||||||
|
"genre; include explicit/sexual tags when present); scene_description ",
|
||||||
|
"(one or two sentences on setting, characters, and action, referencing ",
|
||||||
|
"the dialogue where relevant); safety_flag (is_nsfw boolean + ",
|
||||||
|
"content_type from sexual|nudity|gore|violence|disturbing, [] if none). ",
|
||||||
|
"No prose, no markdown."
|
||||||
|
);
|
||||||
|
|
||||||
|
/// Pass-A schema: `{ ocr_results }` only.
|
||||||
|
pub fn ocr_json_schema() -> serde_json::Value {
|
||||||
|
serde_json::json!({
|
||||||
|
"type": "object",
|
||||||
|
"additionalProperties": false,
|
||||||
|
"properties": {
|
||||||
|
"ocr_results": {
|
||||||
|
"type": "array",
|
||||||
|
"items": {
|
||||||
|
"type": "object",
|
||||||
|
"additionalProperties": false,
|
||||||
|
"properties": {
|
||||||
|
"text": { "type": "string" },
|
||||||
|
"kind": { "type": "string", "enum": OCR_KINDS }
|
||||||
|
},
|
||||||
|
"required": ["text", "kind"]
|
||||||
|
}
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"required": ["ocr_results"]
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Pass-B schema: `{ tagging_results, scene_description, safety_flag }`.
|
||||||
|
pub fn grounding_json_schema() -> serde_json::Value {
|
||||||
|
serde_json::json!({
|
||||||
|
"type": "object",
|
||||||
|
"additionalProperties": false,
|
||||||
|
"properties": {
|
||||||
|
"tagging_results": { "type": "array", "items": { "type": "string" } },
|
||||||
|
"scene_description": { "type": "string" },
|
||||||
|
"safety_flag": {
|
||||||
|
"type": "object",
|
||||||
|
"additionalProperties": false,
|
||||||
|
"properties": {
|
||||||
|
"is_nsfw": { "type": "boolean" },
|
||||||
|
"content_type": {
|
||||||
|
"type": "array",
|
||||||
|
"items": { "type": "string", "enum": CONTENT_WARNINGS }
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"required": ["is_nsfw", "content_type"]
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"required": ["tagging_results", "scene_description", "safety_flag"]
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
/// JSON Schema for the analysis output, used in `response_format:
|
/// JSON Schema for the analysis output, used in `response_format:
|
||||||
/// json_schema` mode (OpenAI structured outputs / LM Studio). Mirrors
|
/// json_schema` mode (OpenAI structured outputs / LM Studio). Mirrors
|
||||||
/// [`crate::domain::page_analysis::VisionAnalysis`]; `additionalProperties:
|
/// [`crate::domain::page_analysis::VisionAnalysis`]; `additionalProperties:
|
||||||
|
|||||||
@@ -1,22 +1,30 @@
|
|||||||
//! The OpenAI-compatible vision HTTP client.
|
//! The OpenAI-compatible vision HTTP client + the long-page slicing wrapper.
|
||||||
//!
|
//!
|
||||||
//! Flow: downscale the page image to fit the model's context budget →
|
//! [`VisionClient::analyze`] decides per page:
|
||||||
//! base64 data-URL → chat/completions request with an `image_url` content
|
//! * **Normal-aspect page** → one combined call (OCR + tags + scene + safety).
|
||||||
//! part → parse `choices[0].message.content` as JSON → sanitize/bound the
|
//! * **Tall page** → slice it into budget-sized bands at native resolution,
|
||||||
//! result. Parsing and sanitization are pure functions so they're tested
|
//! OCR each (Pass A), merge with seam-dedup, then one grounding call that
|
||||||
//! without a live model; only [`VisionClient::analyze`] does I/O.
|
//! gets the whole (downscaled) image **plus** the merged OCR text and
|
||||||
|
//! returns tags + scene + safety (Pass B). Only OCR is merged.
|
||||||
|
//!
|
||||||
|
//! The geometry ([`plan_slices`]), merge ([`merge_ocr`]) and request builders
|
||||||
|
//! are pure so they're unit-tested without a live model; only the `post_chat`
|
||||||
|
//! path does I/O.
|
||||||
|
|
||||||
|
use std::collections::HashSet;
|
||||||
use std::io::Cursor;
|
use std::io::Cursor;
|
||||||
|
|
||||||
use anyhow::{anyhow, Context};
|
use anyhow::{anyhow, Context};
|
||||||
use base64::Engine;
|
use base64::Engine;
|
||||||
|
use image::imageops::FilterType;
|
||||||
|
use image::DynamicImage;
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
|
|
||||||
use crate::analysis::prompt::{
|
use crate::analysis::prompt::{
|
||||||
self, MAX_OCR_PIECES, MAX_OCR_TEXT_CHARS, MAX_SCENE_CHARS, MAX_TAGS,
|
self, MAX_OCR_PIECES, MAX_OCR_TEXT_CHARS, MAX_SCENE_CHARS, MAX_TAGS,
|
||||||
};
|
};
|
||||||
use crate::config::{AnalysisConfig, ResponseFormat};
|
use crate::config::{AnalysisConfig, ResponseFormat};
|
||||||
use crate::domain::page_analysis::{ContentWarning, VisionAnalysis};
|
use crate::domain::page_analysis::{ContentWarning, OcrResult, VisionAnalysis};
|
||||||
|
|
||||||
/// Vision client built from [`AnalysisConfig`]. Cheap to clone (holds a
|
/// Vision client built from [`AnalysisConfig`]. Cheap to clone (holds a
|
||||||
/// `reqwest::Client`, which is internally `Arc`-backed).
|
/// `reqwest::Client`, which is internally `Arc`-backed).
|
||||||
@@ -27,8 +35,18 @@ pub struct VisionClient {
|
|||||||
model: String,
|
model: String,
|
||||||
api_key: Option<String>,
|
api_key: Option<String>,
|
||||||
max_tokens: u32,
|
max_tokens: u32,
|
||||||
max_image_dim: u32,
|
|
||||||
response_format: ResponseFormat,
|
response_format: ResponseFormat,
|
||||||
|
slice: SliceParams,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Pure geometry inputs, mirrored from [`AnalysisConfig`].
|
||||||
|
#[derive(Clone, Copy, Debug)]
|
||||||
|
struct SliceParams {
|
||||||
|
max_pixels: u32,
|
||||||
|
min_slice_height: u32,
|
||||||
|
overlap: f64,
|
||||||
|
tall_threshold: f64,
|
||||||
|
max_slices: usize,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl VisionClient {
|
impl VisionClient {
|
||||||
@@ -39,28 +57,102 @@ impl VisionClient {
|
|||||||
model: cfg.model.clone(),
|
model: cfg.model.clone(),
|
||||||
api_key: cfg.api_key.clone(),
|
api_key: cfg.api_key.clone(),
|
||||||
max_tokens: cfg.max_tokens,
|
max_tokens: cfg.max_tokens,
|
||||||
max_image_dim: cfg.max_image_dim,
|
|
||||||
response_format: cfg.response_format,
|
response_format: cfg.response_format,
|
||||||
|
slice: SliceParams {
|
||||||
|
max_pixels: cfg.max_pixels,
|
||||||
|
min_slice_height: cfg.min_slice_height,
|
||||||
|
overlap: cfg.slice_overlap,
|
||||||
|
tall_threshold: cfg.tall_aspect_threshold,
|
||||||
|
max_slices: cfg.max_slices,
|
||||||
|
},
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Analyze one page image. `mime` is the stored content type (used for
|
/// Analyze one page image. `mime` is the stored content type (used only
|
||||||
/// the data-URL when no downscale happens).
|
/// for the fallback when the image can't be decoded locally).
|
||||||
pub async fn analyze(&self, image: &[u8], mime: &str) -> anyhow::Result<VisionAnalysis> {
|
pub async fn analyze(&self, image: &[u8], mime: &str) -> anyhow::Result<VisionAnalysis> {
|
||||||
let (bytes, mime) = match downscale(image, self.max_image_dim) {
|
// Undecodable locally → send the raw bytes in a single combined call
|
||||||
Some((b, m)) => (b, m.to_string()),
|
// and let the server cope (preserves the prior behavior).
|
||||||
None => (image.to_vec(), mime.to_string()),
|
let Some(img) = image::load_from_memory(image).ok() else {
|
||||||
|
let url = format!("data:{mime};base64,{}", b64(image));
|
||||||
|
let body =
|
||||||
|
build_request_body(&self.model, self.max_tokens, &url, self.response_format);
|
||||||
|
return parse_chat_completion(&self.post_chat(body).await?);
|
||||||
};
|
};
|
||||||
let b64 = base64::engine::general_purpose::STANDARD.encode(&bytes);
|
|
||||||
let data_url = format!("data:{mime};base64,{b64}");
|
|
||||||
|
|
||||||
let body = build_request_body(
|
match plan_slices(img.width(), img.height(), &self.slice) {
|
||||||
&self.model,
|
Plan::Single { .. } => {
|
||||||
self.max_tokens,
|
let jpeg = render_whole(&img, self.slice.max_pixels)
|
||||||
&data_url,
|
.ok_or_else(|| anyhow!("failed to encode page image"))?;
|
||||||
self.response_format,
|
let body = build_request_body(
|
||||||
);
|
&self.model,
|
||||||
|
self.max_tokens,
|
||||||
|
&data_url(&jpeg),
|
||||||
|
self.response_format,
|
||||||
|
);
|
||||||
|
parse_chat_completion(&self.post_chat(body).await?)
|
||||||
|
}
|
||||||
|
Plan::Sliced { width, height, bands } => {
|
||||||
|
tracing::debug!(
|
||||||
|
bands = bands.len(),
|
||||||
|
width,
|
||||||
|
height,
|
||||||
|
"analysis: slicing tall page"
|
||||||
|
);
|
||||||
|
// Work image at the (possibly width-reduced) slice space.
|
||||||
|
let work = if width == img.width() && height == img.height() {
|
||||||
|
img.clone()
|
||||||
|
} else {
|
||||||
|
img.resize_exact(width, height, FilterType::Triangle)
|
||||||
|
};
|
||||||
|
|
||||||
|
// Pass A: OCR each band.
|
||||||
|
let mut slices: Vec<Vec<OcrResult>> = Vec::with_capacity(bands.len());
|
||||||
|
for (y0, y1) in &bands {
|
||||||
|
let jpeg = render_slice(&work, *y0, *y1, self.slice.max_pixels)
|
||||||
|
.ok_or_else(|| anyhow!("failed to encode page slice"))?;
|
||||||
|
let body = build_ocr_body(
|
||||||
|
&self.model,
|
||||||
|
self.max_tokens,
|
||||||
|
self.response_format,
|
||||||
|
&data_url(&jpeg),
|
||||||
|
);
|
||||||
|
let parsed = parse_chat_completion(&self.post_chat(body).await?)?;
|
||||||
|
slices.push(parsed.ocr_results);
|
||||||
|
}
|
||||||
|
let merged = merge_ocr(slices);
|
||||||
|
|
||||||
|
// Pass B: ground tags/scene/safety on the whole image + OCR.
|
||||||
|
let whole = render_whole(&img, self.slice.max_pixels)
|
||||||
|
.ok_or_else(|| anyhow!("failed to encode page image"))?;
|
||||||
|
let ocr_text = merged
|
||||||
|
.iter()
|
||||||
|
.map(|o| format!("[{}] {}", o.kind, o.text))
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
.join("\n");
|
||||||
|
let body = build_grounding_body(
|
||||||
|
&self.model,
|
||||||
|
self.max_tokens,
|
||||||
|
self.response_format,
|
||||||
|
&data_url(&whole),
|
||||||
|
&ocr_text,
|
||||||
|
);
|
||||||
|
let g = parse_chat_completion(&self.post_chat(body).await?)?;
|
||||||
|
|
||||||
|
Ok(sanitize(VisionAnalysis {
|
||||||
|
ocr_results: merged,
|
||||||
|
tagging_results: g.tagging_results,
|
||||||
|
scene_description: g.scene_description,
|
||||||
|
safety_flag: g.safety_flag,
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// POST a chat/completions body and return the parsed JSON value,
|
||||||
|
/// surfacing the server's error body on a non-2xx (different
|
||||||
|
/// OpenAI-compatible servers reject different fields).
|
||||||
|
async fn post_chat(&self, body: serde_json::Value) -> anyhow::Result<serde_json::Value> {
|
||||||
let mut req = self.http.post(&self.endpoint).json(&body);
|
let mut req = self.http.post(&self.endpoint).json(&body);
|
||||||
if let Some(key) = &self.api_key {
|
if let Some(key) = &self.api_key {
|
||||||
req = req.bearer_auth(key);
|
req = req.bearer_auth(key);
|
||||||
@@ -68,39 +160,310 @@ impl VisionClient {
|
|||||||
let resp = req.send().await.context("vision request failed")?;
|
let resp = req.send().await.context("vision request failed")?;
|
||||||
let status = resp.status();
|
let status = resp.status();
|
||||||
if !status.is_success() {
|
if !status.is_success() {
|
||||||
// Surface the server's error body — different OpenAI-compatible
|
|
||||||
// servers reject different fields (e.g. LM Studio rejects
|
|
||||||
// response_format=json_object), and the body says which.
|
|
||||||
let body = resp.text().await.unwrap_or_default();
|
let body = resp.text().await.unwrap_or_default();
|
||||||
anyhow::bail!(
|
anyhow::bail!(
|
||||||
"vision endpoint returned {status}: {}",
|
"vision endpoint returned {status}: {}",
|
||||||
body.trim().chars().take(800).collect::<String>()
|
body.trim().chars().take(800).collect::<String>()
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
let value: serde_json::Value =
|
resp.json().await.context("vision response was not JSON")
|
||||||
resp.json().await.context("vision response was not JSON")?;
|
|
||||||
parse_chat_completion(&value)
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Build the chat/completions request body. Pure so the per-format shape is
|
// ---------------------------------------------------------------------------
|
||||||
/// unit-tested without a live server. The image is sent as an `image_url`
|
// Geometry
|
||||||
/// content part with a base64 data URL (OpenAI vision convention).
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
#[derive(Debug, PartialEq, Eq)]
|
||||||
|
enum Plan {
|
||||||
|
/// Whole-page single call (dims are the working dims, informational).
|
||||||
|
Single { width: u32, height: u32 },
|
||||||
|
/// Slice the page into `bands` (y0,y1) over a `width × height` working
|
||||||
|
/// image; `width` is reduced from native only for pathologically wide
|
||||||
|
/// pages.
|
||||||
|
Sliced {
|
||||||
|
width: u32,
|
||||||
|
height: u32,
|
||||||
|
bands: Vec<(u32, u32)>,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Plan how to feed a `w×h` page to the model: one call, or N budget-sized
|
||||||
|
/// vertical slices at native resolution (portrait *or* landscape). See the
|
||||||
|
/// module/plan docs for the rules.
|
||||||
|
fn plan_slices(w: u32, h: u32, p: &SliceParams) -> Plan {
|
||||||
|
if w == 0 || h == 0 {
|
||||||
|
return Plan::Single { width: w, height: h };
|
||||||
|
}
|
||||||
|
let budget = p.max_pixels.max(1) as f64;
|
||||||
|
let overlap = p.overlap.clamp(0.0, 0.9);
|
||||||
|
|
||||||
|
// Aspect guard: a page wider than this would yield hair-thin strips, so
|
||||||
|
// (and only then) reduce the working width.
|
||||||
|
let max_slice_width = (budget / p.min_slice_height.max(1) as f64).floor().max(1.0);
|
||||||
|
let (work_w, work_h) = if (w as f64) > max_slice_width {
|
||||||
|
let s = max_slice_width / w as f64;
|
||||||
|
(
|
||||||
|
max_slice_width as u32,
|
||||||
|
((h as f64) * s).round().max(1.0) as u32,
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
(w, h)
|
||||||
|
};
|
||||||
|
|
||||||
|
// Budget-sized slice height for this width.
|
||||||
|
let slice_h_budget = (budget / work_w as f64).floor().clamp(1.0, work_h as f64);
|
||||||
|
if (work_h as f64) <= slice_h_budget * p.tall_threshold.max(1.0) {
|
||||||
|
return Plan::Single {
|
||||||
|
width: work_w,
|
||||||
|
height: work_h,
|
||||||
|
};
|
||||||
|
}
|
||||||
|
|
||||||
|
// Would the budget-sized slices exceed the cap? If so, grow the slice
|
||||||
|
// height so exactly `max_slices` bands cover the page (coarser; the
|
||||||
|
// oversized crop is downscaled to budget by `render_slice`).
|
||||||
|
let step_for = |sh: f64| (sh * (1.0 - overlap)).max(1.0);
|
||||||
|
let bands_for = |sh: f64| (((work_h as f64 - sh) / step_for(sh)).ceil() as i64 + 1).max(1);
|
||||||
|
let slice_h = if bands_for(slice_h_budget) > p.max_slices as i64 {
|
||||||
|
(work_h as f64 / (1.0 + (p.max_slices as f64 - 1.0) * (1.0 - overlap)))
|
||||||
|
.ceil()
|
||||||
|
.clamp(1.0, work_h as f64)
|
||||||
|
} else {
|
||||||
|
slice_h_budget
|
||||||
|
} as u32;
|
||||||
|
|
||||||
|
let step = ((slice_h as f64) * (1.0 - overlap)).round().max(1.0) as u32;
|
||||||
|
let mut bands: Vec<(u32, u32)> = Vec::new();
|
||||||
|
let mut y0 = 0u32;
|
||||||
|
loop {
|
||||||
|
let y1 = (y0 + slice_h).min(work_h);
|
||||||
|
bands.push((y0, y1));
|
||||||
|
if y1 >= work_h || bands.len() >= p.max_slices {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
y0 += step;
|
||||||
|
}
|
||||||
|
// Guarantee coverage to the bottom edge.
|
||||||
|
if let Some(last) = bands.last_mut() {
|
||||||
|
last.1 = work_h;
|
||||||
|
}
|
||||||
|
Plan::Sliced {
|
||||||
|
width: work_w,
|
||||||
|
height: work_h,
|
||||||
|
bands,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Image rendering
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// Resize `img` down to `max_pixels` (preserving aspect) only if it exceeds
|
||||||
|
/// the budget, then JPEG-encode. An image already within budget is encoded
|
||||||
|
/// at native resolution.
|
||||||
|
fn render_whole(img: &DynamicImage, max_pixels: u32) -> Option<Vec<u8>> {
|
||||||
|
let pixels = img.width() as u64 * img.height() as u64;
|
||||||
|
let out = if pixels > max_pixels as u64 && pixels > 0 {
|
||||||
|
let s = (max_pixels as f64 / pixels as f64).sqrt();
|
||||||
|
let nw = ((img.width() as f64 * s).round() as u32).max(1);
|
||||||
|
let nh = ((img.height() as f64 * s).round() as u32).max(1);
|
||||||
|
img.resize_exact(nw, nh, FilterType::Triangle)
|
||||||
|
} else {
|
||||||
|
img.clone()
|
||||||
|
};
|
||||||
|
encode_jpeg(&out)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Crop a vertical band `[y0, y1)` from the working image and encode it,
|
||||||
|
/// downscaling to budget only if the band is oversized (the coarse fallback).
|
||||||
|
fn render_slice(work: &DynamicImage, y0: u32, y1: u32, max_pixels: u32) -> Option<Vec<u8>> {
|
||||||
|
let h = y1.saturating_sub(y0).max(1).min(work.height().saturating_sub(y0).max(1));
|
||||||
|
let crop = work.crop_imm(0, y0, work.width(), h);
|
||||||
|
render_whole(&crop, max_pixels)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn encode_jpeg(img: &DynamicImage) -> Option<Vec<u8>> {
|
||||||
|
let mut buf = Vec::new();
|
||||||
|
img.to_rgb8()
|
||||||
|
.write_to(&mut Cursor::new(&mut buf), image::ImageFormat::Jpeg)
|
||||||
|
.ok()?;
|
||||||
|
Some(buf)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn b64(bytes: &[u8]) -> String {
|
||||||
|
base64::engine::general_purpose::STANDARD.encode(bytes)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn data_url(jpeg: &[u8]) -> String {
|
||||||
|
format!("data:image/jpeg;base64,{}", b64(jpeg))
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// OCR merge (seam dedup)
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// How many pieces at each seam edge to consider for dedup.
|
||||||
|
const SEAM: usize = 8;
|
||||||
|
|
||||||
|
/// Merge per-slice OCR (top-to-bottom) into one ordered list, de-duplicating
|
||||||
|
/// text that appears in both slices' overlap region. Dedup is *seam-scoped*
|
||||||
|
/// (only the first `SEAM` pieces of a slice vs. the last `SEAM` already
|
||||||
|
/// emitted), so a genuinely repeated line in non-adjacent slices is kept.
|
||||||
|
/// When two pieces match, the longer (more complete) transcription wins.
|
||||||
|
fn merge_ocr(slices: Vec<Vec<OcrResult>>) -> Vec<OcrResult> {
|
||||||
|
let mut out: Vec<OcrResult> = Vec::new();
|
||||||
|
// Index in `out` where the *previous* slice's pieces began — dedup
|
||||||
|
// compares only against that slice's tail, never earlier ones.
|
||||||
|
let mut prev_slice_start = 0usize;
|
||||||
|
for (si, slice) in slices.into_iter().enumerate() {
|
||||||
|
let cur_slice_start = out.len();
|
||||||
|
for (pi, piece) in slice.into_iter().enumerate() {
|
||||||
|
if si > 0 && pi < SEAM {
|
||||||
|
let win_start = cur_slice_start.saturating_sub(SEAM).max(prev_slice_start);
|
||||||
|
let np = normalize_text(&piece.text);
|
||||||
|
if !np.is_empty() {
|
||||||
|
if let Some(j) = (win_start..cur_slice_start)
|
||||||
|
.find(|&j| similar(&normalize_text(&out[j].text), &np))
|
||||||
|
{
|
||||||
|
if piece.text.chars().count() > out[j].text.chars().count() {
|
||||||
|
out[j] = piece;
|
||||||
|
}
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
out.push(piece);
|
||||||
|
}
|
||||||
|
prev_slice_start = cur_slice_start;
|
||||||
|
}
|
||||||
|
out
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Lowercased, alphanumeric-only, whitespace-collapsed form for comparison.
|
||||||
|
fn normalize_text(s: &str) -> String {
|
||||||
|
let mut t = String::with_capacity(s.len());
|
||||||
|
for c in s.chars() {
|
||||||
|
if c.is_alphanumeric() {
|
||||||
|
t.extend(c.to_lowercase());
|
||||||
|
} else {
|
||||||
|
t.push(' ');
|
||||||
|
}
|
||||||
|
}
|
||||||
|
t.split_whitespace().collect::<Vec<_>>().join(" ")
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Heuristic similarity over already-normalized strings: equal, containment
|
||||||
|
/// (one inside the other, ≥3 chars), or word-set Jaccard ≥ 0.6.
|
||||||
|
fn similar(a: &str, b: &str) -> bool {
|
||||||
|
if a == b {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
if a.is_empty() || b.is_empty() {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
let (short, long) = if a.len() <= b.len() { (a, b) } else { (b, a) };
|
||||||
|
if short.len() >= 3 && long.contains(short) {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
let wa: HashSet<&str> = a.split_whitespace().collect();
|
||||||
|
let wb: HashSet<&str> = b.split_whitespace().collect();
|
||||||
|
if wa.is_empty() || wb.is_empty() {
|
||||||
|
return false;
|
||||||
|
}
|
||||||
|
let inter = wa.intersection(&wb).count() as f64;
|
||||||
|
let union = wa.union(&wb).count() as f64;
|
||||||
|
inter / union >= 0.6
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Request builders
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/// Combined single-call body (OCR + tags + scene + safety). Public + this
|
||||||
|
/// exact signature so the existing tests keep working.
|
||||||
pub fn build_request_body(
|
pub fn build_request_body(
|
||||||
model: &str,
|
model: &str,
|
||||||
max_tokens: u32,
|
max_tokens: u32,
|
||||||
data_url: &str,
|
data_url: &str,
|
||||||
response_format: ResponseFormat,
|
response_format: ResponseFormat,
|
||||||
) -> serde_json::Value {
|
) -> serde_json::Value {
|
||||||
|
build_body(
|
||||||
|
model,
|
||||||
|
max_tokens,
|
||||||
|
response_format,
|
||||||
|
prompt::SYSTEM_PROMPT,
|
||||||
|
"page_analysis",
|
||||||
|
prompt::output_json_schema(),
|
||||||
|
data_url,
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Pass-A OCR-only body.
|
||||||
|
fn build_ocr_body(
|
||||||
|
model: &str,
|
||||||
|
max_tokens: u32,
|
||||||
|
response_format: ResponseFormat,
|
||||||
|
data_url: &str,
|
||||||
|
) -> serde_json::Value {
|
||||||
|
build_body(
|
||||||
|
model,
|
||||||
|
max_tokens,
|
||||||
|
response_format,
|
||||||
|
prompt::OCR_PROMPT,
|
||||||
|
"page_ocr",
|
||||||
|
prompt::ocr_json_schema(),
|
||||||
|
data_url,
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Pass-B grounding body — whole image + the merged OCR text as context.
|
||||||
|
fn build_grounding_body(
|
||||||
|
model: &str,
|
||||||
|
max_tokens: u32,
|
||||||
|
response_format: ResponseFormat,
|
||||||
|
data_url: &str,
|
||||||
|
ocr_text: &str,
|
||||||
|
) -> serde_json::Value {
|
||||||
|
let context = format!("OCR text extracted from this page:\n{ocr_text}");
|
||||||
|
build_body(
|
||||||
|
model,
|
||||||
|
max_tokens,
|
||||||
|
response_format,
|
||||||
|
prompt::GROUNDING_PROMPT,
|
||||||
|
"page_grounding",
|
||||||
|
prompt::grounding_json_schema(),
|
||||||
|
data_url,
|
||||||
|
Some(&context),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Build a chat/completions body with a system prompt, an image part, an
|
||||||
|
/// optional trailing text part, and the requested `response_format`.
|
||||||
|
fn build_body(
|
||||||
|
model: &str,
|
||||||
|
max_tokens: u32,
|
||||||
|
response_format: ResponseFormat,
|
||||||
|
system_prompt: &str,
|
||||||
|
schema_name: &str,
|
||||||
|
schema: serde_json::Value,
|
||||||
|
data_url: &str,
|
||||||
|
user_text: Option<&str>,
|
||||||
|
) -> serde_json::Value {
|
||||||
|
let mut user_content =
|
||||||
|
vec![json!({ "type": "image_url", "image_url": { "url": data_url } })];
|
||||||
|
if let Some(text) = user_text {
|
||||||
|
user_content.push(json!({ "type": "text", "text": text }));
|
||||||
|
}
|
||||||
let mut body = json!({
|
let mut body = json!({
|
||||||
"model": model,
|
"model": model,
|
||||||
"temperature": 0,
|
"temperature": 0,
|
||||||
"max_tokens": max_tokens,
|
"max_tokens": max_tokens,
|
||||||
"messages": [
|
"messages": [
|
||||||
{ "role": "system", "content": prompt::SYSTEM_PROMPT },
|
{ "role": "system", "content": system_prompt },
|
||||||
{ "role": "user", "content": [
|
{ "role": "user", "content": user_content }
|
||||||
{ "type": "image_url", "image_url": { "url": data_url } }
|
|
||||||
]}
|
|
||||||
]
|
]
|
||||||
});
|
});
|
||||||
let rf = match response_format {
|
let rf = match response_format {
|
||||||
@@ -108,11 +471,7 @@ pub fn build_request_body(
|
|||||||
ResponseFormat::JsonObject => Some(json!({ "type": "json_object" })),
|
ResponseFormat::JsonObject => Some(json!({ "type": "json_object" })),
|
||||||
ResponseFormat::JsonSchema => Some(json!({
|
ResponseFormat::JsonSchema => Some(json!({
|
||||||
"type": "json_schema",
|
"type": "json_schema",
|
||||||
"json_schema": {
|
"json_schema": { "name": schema_name, "strict": true, "schema": schema }
|
||||||
"name": "page_analysis",
|
|
||||||
"strict": true,
|
|
||||||
"schema": prompt::output_json_schema()
|
|
||||||
}
|
|
||||||
})),
|
})),
|
||||||
};
|
};
|
||||||
if let Some(rf) = rf {
|
if let Some(rf) = rf {
|
||||||
@@ -121,25 +480,13 @@ pub fn build_request_body(
|
|||||||
body
|
body
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Downscale `bytes` so its longest edge is `max_dim`, re-encoding as JPEG.
|
// ---------------------------------------------------------------------------
|
||||||
/// Returns `None` (use the original bytes) when the image already fits or
|
// Response parsing / sanitizing
|
||||||
/// can't be decoded — a non-decodable page is passed through untouched and
|
// ---------------------------------------------------------------------------
|
||||||
/// the model server can reject it if it wants.
|
|
||||||
fn downscale(bytes: &[u8], max_dim: u32) -> Option<(Vec<u8>, &'static str)> {
|
|
||||||
let img = image::load_from_memory(bytes).ok()?;
|
|
||||||
if img.width().max(img.height()) <= max_dim {
|
|
||||||
return None;
|
|
||||||
}
|
|
||||||
let scaled = img.resize(max_dim, max_dim, image::imageops::FilterType::Triangle);
|
|
||||||
let mut buf = Vec::new();
|
|
||||||
scaled
|
|
||||||
.to_rgb8()
|
|
||||||
.write_to(&mut Cursor::new(&mut buf), image::ImageFormat::Jpeg)
|
|
||||||
.ok()?;
|
|
||||||
Some((buf, "image/jpeg"))
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Extract the model's JSON object from a chat/completions response.
|
/// Extract and parse the model's JSON object from a chat/completions
|
||||||
|
/// response into a (sanitized) [`VisionAnalysis`]. Lenient field defaults
|
||||||
|
/// let the same parser serve the combined, OCR-only and grounding shapes.
|
||||||
pub fn parse_chat_completion(value: &serde_json::Value) -> anyhow::Result<VisionAnalysis> {
|
pub fn parse_chat_completion(value: &serde_json::Value) -> anyhow::Result<VisionAnalysis> {
|
||||||
let content = value
|
let content = value
|
||||||
.get("choices")
|
.get("choices")
|
||||||
@@ -155,9 +502,7 @@ pub fn parse_chat_completion(value: &serde_json::Value) -> anyhow::Result<Vision
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Return the substring from the first `{` to the last `}` (inclusive),
|
/// Return the substring from the first `{` to the last `}` (inclusive),
|
||||||
/// tolerating markdown fences or stray prose around the JSON object. Falls
|
/// tolerating markdown fences or stray prose around the JSON object.
|
||||||
/// back to the trimmed input if no braces are found (so the parse fails
|
|
||||||
/// with a clear error rather than silently matching nothing).
|
|
||||||
fn extract_json_object(s: &str) -> &str {
|
fn extract_json_object(s: &str) -> &str {
|
||||||
match (s.find('{'), s.rfind('}')) {
|
match (s.find('{'), s.rfind('}')) {
|
||||||
(Some(start), Some(end)) if end >= start => &s[start..=end],
|
(Some(start), Some(end)) if end >= start => &s[start..=end],
|
||||||
@@ -166,10 +511,8 @@ fn extract_json_object(s: &str) -> &str {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Bound the model output: drop empty OCR text, truncate over-long fields,
|
/// Bound the model output: drop empty OCR text, truncate over-long fields,
|
||||||
/// clamp tags to [`MAX_TAGS`] via the page-tag normalizer, and keep only
|
/// clamp tags via the page-tag normalizer, and keep only recognized
|
||||||
/// recognized content-warning values. Persist applies the same rules as a
|
/// content-warning values.
|
||||||
/// safety net, but doing it here keeps the row tidy and the worker's logs
|
|
||||||
/// honest about what it stored.
|
|
||||||
pub fn sanitize(mut a: VisionAnalysis) -> VisionAnalysis {
|
pub fn sanitize(mut a: VisionAnalysis) -> VisionAnalysis {
|
||||||
a.ocr_results.retain(|r| !r.text.trim().is_empty());
|
a.ocr_results.retain(|r| !r.text.trim().is_empty());
|
||||||
a.ocr_results.truncate(MAX_OCR_PIECES);
|
a.ocr_results.truncate(MAX_OCR_PIECES);
|
||||||
@@ -177,9 +520,6 @@ pub fn sanitize(mut a: VisionAnalysis) -> VisionAnalysis {
|
|||||||
r.text = truncate_chars(r.text.trim(), MAX_OCR_TEXT_CHARS);
|
r.text = truncate_chars(r.text.trim(), MAX_OCR_TEXT_CHARS);
|
||||||
}
|
}
|
||||||
|
|
||||||
// Normalize tags through the shared page-tag rules (lowercase, collapse
|
|
||||||
// whitespace, reject wildcards/control/invisible chars); drop the
|
|
||||||
// unmappable, dedup, cap.
|
|
||||||
let mut seen = std::collections::HashSet::new();
|
let mut seen = std::collections::HashSet::new();
|
||||||
let mut tags = Vec::new();
|
let mut tags = Vec::new();
|
||||||
for raw in std::mem::take(&mut a.tagging_results) {
|
for raw in std::mem::take(&mut a.tagging_results) {
|
||||||
@@ -196,7 +536,6 @@ pub fn sanitize(mut a: VisionAnalysis) -> VisionAnalysis {
|
|||||||
|
|
||||||
a.scene_description = truncate_chars(a.scene_description.trim(), MAX_SCENE_CHARS);
|
a.scene_description = truncate_chars(a.scene_description.trim(), MAX_SCENE_CHARS);
|
||||||
|
|
||||||
// Keep only recognized warnings (canonical lowercase), deduped.
|
|
||||||
let mut seen_w = std::collections::HashSet::new();
|
let mut seen_w = std::collections::HashSet::new();
|
||||||
a.safety_flag.content_type = std::mem::take(&mut a.safety_flag.content_type)
|
a.safety_flag.content_type = std::mem::take(&mut a.safety_flag.content_type)
|
||||||
.into_iter()
|
.into_iter()
|
||||||
@@ -207,8 +546,6 @@ pub fn sanitize(mut a: VisionAnalysis) -> VisionAnalysis {
|
|||||||
a
|
a
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Truncate to at most `n` chars (not bytes), avoiding a panic on a
|
|
||||||
/// multi-byte boundary.
|
|
||||||
fn truncate_chars(s: &str, n: usize) -> String {
|
fn truncate_chars(s: &str, n: usize) -> String {
|
||||||
s.chars().take(n).collect()
|
s.chars().take(n).collect()
|
||||||
}
|
}
|
||||||
@@ -221,6 +558,30 @@ mod tests {
|
|||||||
json!({ "choices": [ { "message": { "content": content } } ] })
|
json!({ "choices": [ { "message": { "content": content } } ] })
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn params() -> SliceParams {
|
||||||
|
SliceParams {
|
||||||
|
max_pixels: 1_000_000,
|
||||||
|
min_slice_height: 640,
|
||||||
|
overlap: 0.12,
|
||||||
|
tall_threshold: 1.6,
|
||||||
|
max_slices: 16,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn ocr(text: &str, kind: &str) -> OcrResult {
|
||||||
|
OcrResult {
|
||||||
|
text: text.into(),
|
||||||
|
kind: kind.into(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn jpeg_of(w: u32, h: u32) -> Vec<u8> {
|
||||||
|
let img = DynamicImage::ImageRgb8(image::RgbImage::from_pixel(w, h, image::Rgb([9, 9, 9])));
|
||||||
|
encode_jpeg(&img).unwrap()
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- parse / sanitize (unchanged behavior) ---
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn parses_a_clean_json_response() {
|
fn parses_a_clean_json_response() {
|
||||||
let body = chat_body(
|
let body = chat_body(
|
||||||
@@ -236,48 +597,23 @@ mod tests {
|
|||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn strips_markdown_fences_and_prose() {
|
fn ocr_only_response_parses_with_defaults() {
|
||||||
let body = chat_body(
|
let body = chat_body(r#"{"ocr_results":[{"text":"Hi","kind":"speech"}]}"#);
|
||||||
"Here you go:\n```json\n{\"scene_description\":\"x\"}\n```\nthanks",
|
|
||||||
);
|
|
||||||
let v = parse_chat_completion(&body).unwrap();
|
let v = parse_chat_completion(&body).unwrap();
|
||||||
assert_eq!(v.scene_description, "x");
|
assert_eq!(v.ocr_results.len(), 1);
|
||||||
|
assert!(v.tagging_results.is_empty());
|
||||||
|
assert_eq!(v.scene_description, "");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn strips_markdown_fences_and_prose() {
|
||||||
|
let body = chat_body("Here:\n```json\n{\"scene_description\":\"x\"}\n```\nthanks");
|
||||||
|
assert_eq!(parse_chat_completion(&body).unwrap().scene_description, "x");
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn rejects_non_json_content() {
|
fn rejects_non_json_content() {
|
||||||
let body = chat_body("I cannot analyze this image.");
|
assert!(parse_chat_completion(&chat_body("nope")).is_err());
|
||||||
assert!(parse_chat_completion(&body).is_err());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn rejects_missing_content() {
|
|
||||||
assert!(parse_chat_completion(&json!({"choices": []})).is_err());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn sanitize_clamps_tags_and_drops_invalid() {
|
|
||||||
let body = chat_body(
|
|
||||||
r#"{"tagging_results":["A","a","b","b","c","d","e","f","g","h","i","j",
|
|
||||||
"bad%wild"," "]}"#,
|
|
||||||
);
|
|
||||||
let v = parse_chat_completion(&body).unwrap();
|
|
||||||
// "A"/"a" dedup to one; wildcard + blank dropped; capped at 10.
|
|
||||||
assert!(v.tagging_results.len() <= MAX_TAGS);
|
|
||||||
assert!(v.tagging_results.contains(&"a".to_string()));
|
|
||||||
assert!(!v.tagging_results.iter().any(|t| t.contains('%')));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn sanitize_drops_empty_ocr_and_truncates() {
|
|
||||||
let long = "x".repeat(MAX_OCR_TEXT_CHARS + 50);
|
|
||||||
let body = chat_body(&format!(
|
|
||||||
r#"{{"ocr_results":[{{"text":" ","kind":"speech"}},
|
|
||||||
{{"text":"{long}","kind":"narration"}}]}}"#
|
|
||||||
));
|
|
||||||
let v = parse_chat_completion(&body).unwrap();
|
|
||||||
assert_eq!(v.ocr_results.len(), 1, "blank OCR dropped");
|
|
||||||
assert_eq!(v.ocr_results[0].text.chars().count(), MAX_OCR_TEXT_CHARS);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -289,46 +625,40 @@ mod tests {
|
|||||||
assert_eq!(v.safety_flag.content_type, vec!["sexual", "gore"]);
|
assert_eq!(v.safety_flag.content_type, vec!["sexual", "gore"]);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
// --- request builders ---
|
||||||
fn downscale_passes_through_small_images() {
|
|
||||||
// 8x8 PNG — well under any sane max_dim, so no re-encode.
|
|
||||||
let img = image::RgbImage::from_pixel(8, 8, image::Rgb([10, 20, 30]));
|
|
||||||
let mut buf = Vec::new();
|
|
||||||
image::DynamicImage::ImageRgb8(img)
|
|
||||||
.write_to(&mut Cursor::new(&mut buf), image::ImageFormat::Png)
|
|
||||||
.unwrap();
|
|
||||||
assert!(downscale(&buf, 1024).is_none());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn downscale_shrinks_large_images_to_max_dim() {
|
fn combined_body_uses_full_schema_and_image() {
|
||||||
let img = image::RgbImage::from_pixel(2000, 1000, image::Rgb([1, 2, 3]));
|
let body = build_request_body("m", 100, "data:image/png;base64,AA", ResponseFormat::JsonSchema);
|
||||||
let mut buf = Vec::new();
|
|
||||||
image::DynamicImage::ImageRgb8(img)
|
|
||||||
.write_to(&mut Cursor::new(&mut buf), image::ImageFormat::Png)
|
|
||||||
.unwrap();
|
|
||||||
let (out, mime) = downscale(&buf, 512).expect("should downscale");
|
|
||||||
assert_eq!(mime, "image/jpeg");
|
|
||||||
let decoded = image::load_from_memory(&out).unwrap();
|
|
||||||
assert_eq!(decoded.width().max(decoded.height()), 512);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn downscale_passes_through_undecodable_bytes() {
|
|
||||||
assert!(downscale(b"not an image", 512).is_none());
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn request_body_carries_model_image_and_prompt() {
|
|
||||||
let body = build_request_body("m", 100, "data:image/png;base64,AA", ResponseFormat::None);
|
|
||||||
assert_eq!(body["model"], "m");
|
|
||||||
assert_eq!(body["max_tokens"], 100);
|
|
||||||
// System prompt + a user image_url part.
|
|
||||||
assert_eq!(body["messages"][0]["role"], "system");
|
assert_eq!(body["messages"][0]["role"], "system");
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
body["messages"][1]["content"][0]["image_url"]["url"],
|
body["messages"][1]["content"][0]["image_url"]["url"],
|
||||||
"data:image/png;base64,AA"
|
"data:image/png;base64,AA"
|
||||||
);
|
);
|
||||||
|
assert_eq!(body["response_format"]["json_schema"]["name"], "page_analysis");
|
||||||
|
let req = &body["response_format"]["json_schema"]["schema"]["required"];
|
||||||
|
assert!(req.as_array().unwrap().iter().any(|v| v == "safety_flag"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn ocr_body_carries_only_the_ocr_schema() {
|
||||||
|
let body = build_ocr_body("m", 100, ResponseFormat::JsonSchema, "d");
|
||||||
|
assert_eq!(body["response_format"]["json_schema"]["name"], "page_ocr");
|
||||||
|
let req = &body["response_format"]["json_schema"]["schema"]["required"];
|
||||||
|
assert_eq!(req, &json!(["ocr_results"]));
|
||||||
|
// No user text part — just the image.
|
||||||
|
assert_eq!(body["messages"][1]["content"].as_array().unwrap().len(), 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn grounding_body_includes_ocr_context_and_grounding_schema() {
|
||||||
|
let body = build_grounding_body("m", 100, ResponseFormat::JsonSchema, "d", "[speech] Hi");
|
||||||
|
assert_eq!(body["response_format"]["json_schema"]["name"], "page_grounding");
|
||||||
|
let parts = body["messages"][1]["content"].as_array().unwrap();
|
||||||
|
assert_eq!(parts.len(), 2, "image + text parts");
|
||||||
|
assert_eq!(parts[0]["type"], "image_url");
|
||||||
|
assert_eq!(parts[1]["type"], "text");
|
||||||
|
assert!(parts[1]["text"].as_str().unwrap().contains("[speech] Hi"));
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -337,20 +667,142 @@ mod tests {
|
|||||||
assert!(body.get("response_format").is_none());
|
assert!(body.get("response_format").is_none());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// --- geometry: plan_slices ---
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn response_format_json_object_sets_type() {
|
fn normal_aspect_page_is_a_single_call() {
|
||||||
let body = build_request_body("m", 100, "d", ResponseFormat::JsonObject);
|
assert!(matches!(
|
||||||
assert_eq!(body["response_format"]["type"], "json_object");
|
plan_slices(800, 1000, ¶ms()),
|
||||||
|
Plan::Single { .. }
|
||||||
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn response_format_json_schema_embeds_the_strict_schema() {
|
fn tall_narrow_page_slices_into_portrait_bands() {
|
||||||
let body = build_request_body("m", 100, "d", ResponseFormat::JsonSchema);
|
let p = params();
|
||||||
assert_eq!(body["response_format"]["type"], "json_schema");
|
let Plan::Sliced { width, height, bands } = plan_slices(800, 8000, &p) else {
|
||||||
assert_eq!(body["response_format"]["json_schema"]["strict"], true);
|
panic!("expected Sliced");
|
||||||
// The schema is the real analysis schema (top-level required keys).
|
};
|
||||||
let required =
|
assert_eq!((width, height), (800, 8000));
|
||||||
&body["response_format"]["json_schema"]["schema"]["required"];
|
assert!(bands.len() >= 2 && bands.len() <= p.max_slices);
|
||||||
assert!(required.as_array().unwrap().iter().any(|v| v == "safety_flag"));
|
assert_eq!(bands.first().unwrap().0, 0, "starts at top");
|
||||||
|
assert_eq!(bands.last().unwrap().1, 8000, "covers the bottom");
|
||||||
|
// Overlapping + advancing bands; each ≈ the pixel budget.
|
||||||
|
for win in bands.windows(2) {
|
||||||
|
assert!(win[1].0 < win[0].1, "adjacent bands overlap");
|
||||||
|
assert!(win[1].0 > win[0].0, "bands advance");
|
||||||
|
}
|
||||||
|
for (y0, y1) in &bands {
|
||||||
|
let px = (y1 - y0) as u64 * width as u64;
|
||||||
|
assert!(px <= 1_000_000 + 1, "band within budget, got {px}");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn wide_page_yields_landscape_bands() {
|
||||||
|
// 1400 < maxSliceWidth(1562) so width is preserved; budget/1400≈714
|
||||||
|
// → bands are wider than tall.
|
||||||
|
let Plan::Sliced { width, bands, .. } = plan_slices(1400, 4000, ¶ms()) else {
|
||||||
|
panic!("expected Sliced");
|
||||||
|
};
|
||||||
|
assert_eq!(width, 1400);
|
||||||
|
for (y0, y1) in &bands {
|
||||||
|
assert!(y1 - y0 < width, "landscape band (height < width)");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn pathologically_wide_page_reduces_width() {
|
||||||
|
let p = params();
|
||||||
|
let max_slice_width = (p.max_pixels as f64 / p.min_slice_height as f64).floor() as u32;
|
||||||
|
let Plan::Sliced { width, .. } = plan_slices(3000, 9000, &p) else {
|
||||||
|
panic!("expected Sliced");
|
||||||
|
};
|
||||||
|
assert!(width <= max_slice_width && width < 3000, "width reduced to the guard");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn very_long_page_is_capped_at_max_slices() {
|
||||||
|
let p = params();
|
||||||
|
let Plan::Sliced { bands, height, .. } = plan_slices(800, 200_000, &p) else {
|
||||||
|
panic!("expected Sliced");
|
||||||
|
};
|
||||||
|
assert!(bands.len() <= p.max_slices);
|
||||||
|
assert_eq!(bands.last().unwrap().1, height, "still covers the bottom");
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- merge_ocr ---
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn merge_dedups_a_line_repeated_across_the_seam() {
|
||||||
|
let merged = merge_ocr(vec![
|
||||||
|
vec![ocr("Hello", "speech"), ocr("world", "speech")],
|
||||||
|
vec![ocr("world", "speech"), ocr("bye", "speech")],
|
||||||
|
]);
|
||||||
|
let texts: Vec<&str> = merged.iter().map(|o| o.text.as_str()).collect();
|
||||||
|
assert_eq!(texts, vec!["Hello", "world", "bye"]);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn merge_keeps_the_longer_transcription() {
|
||||||
|
let merged = merge_ocr(vec![
|
||||||
|
vec![ocr("I don", "speech")],
|
||||||
|
vec![ocr("I don't think so", "speech")],
|
||||||
|
]);
|
||||||
|
assert_eq!(merged.len(), 1);
|
||||||
|
assert_eq!(merged[0].text, "I don't think so");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn merge_keeps_non_adjacent_repeats() {
|
||||||
|
// "boom" in slice 0 and slice 2 must both survive (not a seam dup).
|
||||||
|
let merged = merge_ocr(vec![
|
||||||
|
vec![ocr("boom", "sfx")],
|
||||||
|
vec![ocr("crash", "sfx")],
|
||||||
|
vec![ocr("boom", "sfx")],
|
||||||
|
]);
|
||||||
|
let booms = merged.iter().filter(|o| o.text == "boom").count();
|
||||||
|
assert_eq!(booms, 2);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn merge_preserves_distinct_lines_and_order() {
|
||||||
|
let merged = merge_ocr(vec![
|
||||||
|
vec![ocr("one", "narration")],
|
||||||
|
vec![ocr("two", "speech")],
|
||||||
|
]);
|
||||||
|
assert_eq!(merged.len(), 2);
|
||||||
|
assert_eq!(merged[0].kind, "narration");
|
||||||
|
assert_eq!(merged[1].text, "two");
|
||||||
|
}
|
||||||
|
|
||||||
|
// --- rendering ---
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn render_whole_passes_small_images_through() {
|
||||||
|
let jpeg = jpeg_of(100, 80);
|
||||||
|
let img = image::load_from_memory(&jpeg).unwrap();
|
||||||
|
let out = render_whole(&img, 1_000_000).unwrap();
|
||||||
|
let dec = image::load_from_memory(&out).unwrap();
|
||||||
|
assert_eq!((dec.width(), dec.height()), (100, 80));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn render_whole_fits_oversized_to_budget() {
|
||||||
|
let jpeg = jpeg_of(2000, 1000); // 2.0 MP
|
||||||
|
let img = image::load_from_memory(&jpeg).unwrap();
|
||||||
|
let out = render_whole(&img, 500_000).unwrap();
|
||||||
|
let dec = image::load_from_memory(&out).unwrap();
|
||||||
|
let px = dec.width() as u64 * dec.height() as u64;
|
||||||
|
assert!(px <= 520_000, "fit to ~budget, got {px}");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn render_slice_crops_the_band() {
|
||||||
|
let jpeg = jpeg_of(100, 400);
|
||||||
|
let img = image::load_from_memory(&jpeg).unwrap();
|
||||||
|
let out = render_slice(&img, 50, 150, 1_000_000).unwrap();
|
||||||
|
let dec = image::load_from_memory(&out).unwrap();
|
||||||
|
assert_eq!((dec.width(), dec.height()), (100, 100));
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -118,8 +118,23 @@ pub struct AnalysisConfig {
|
|||||||
/// length`) and the parse fails. The page image + prompt are only a few
|
/// length`) and the parse fails. The page image + prompt are only a few
|
||||||
/// hundred tokens, so a large output budget still fits an 8k window.
|
/// hundred tokens, so a large output budget still fits an 8k window.
|
||||||
pub max_tokens: u32,
|
pub max_tokens: u32,
|
||||||
/// Longest image edge (px) before downscaling (`ANALYSIS_MAX_IMAGE_DIM`).
|
/// Per-slice / per-image pixel budget (`ANALYSIS_MAX_PIXELS`). The model
|
||||||
pub max_image_dim: u32,
|
/// resizes any input to roughly this many pixels anyway, so we slice/fit
|
||||||
|
/// to it at native resolution rather than pre-squashing by a fixed edge.
|
||||||
|
pub max_pixels: u32,
|
||||||
|
/// Minimum slice height (px), the aspect guard for very wide pages
|
||||||
|
/// (`ANALYSIS_MIN_SLICE_HEIGHT`). Implies `max_slice_width =
|
||||||
|
/// max_pixels / min_slice_height`.
|
||||||
|
pub min_slice_height: u32,
|
||||||
|
/// Vertical overlap between adjacent slices as a fraction of slice height
|
||||||
|
/// (`ANALYSIS_SLICE_OVERLAP`), so text straddling a cut survives.
|
||||||
|
pub slice_overlap: f64,
|
||||||
|
/// Slice only when `height/width` exceeds this (`ANALYSIS_TALL_ASPECT`);
|
||||||
|
/// normal-aspect pages take a single combined call.
|
||||||
|
pub tall_aspect_threshold: f64,
|
||||||
|
/// Hard cap on slices per page (`ANALYSIS_MAX_SLICES`); beyond it slices
|
||||||
|
/// grow coarser (and get downscaled to budget) rather than multiplying.
|
||||||
|
pub max_slices: usize,
|
||||||
/// Hard cap on a page image's stored size; larger pages are skipped
|
/// Hard cap on a page image's stored size; larger pages are skipped
|
||||||
/// (`ANALYSIS_MAX_IMAGE_BYTES`).
|
/// (`ANALYSIS_MAX_IMAGE_BYTES`).
|
||||||
pub max_image_bytes: usize,
|
pub max_image_bytes: usize,
|
||||||
@@ -137,9 +152,16 @@ impl Default for AnalysisConfig {
|
|||||||
model: String::new(),
|
model: String::new(),
|
||||||
api_key: None,
|
api_key: None,
|
||||||
request_timeout: Duration::from_secs(120),
|
request_timeout: Duration::from_secs(120),
|
||||||
job_timeout: Duration::from_secs(180),
|
// A sliced long page is N sequential calls under one job, so the
|
||||||
|
// whole-job budget must be generous (request_timeout stays
|
||||||
|
// per-call).
|
||||||
|
job_timeout: Duration::from_secs(600),
|
||||||
max_tokens: 4096,
|
max_tokens: 4096,
|
||||||
max_image_dim: 1024,
|
max_pixels: 1_000_000,
|
||||||
|
min_slice_height: 640,
|
||||||
|
slice_overlap: 0.12,
|
||||||
|
tall_aspect_threshold: 1.6,
|
||||||
|
max_slices: 16,
|
||||||
max_image_bytes: 8 * 1024 * 1024,
|
max_image_bytes: 8 * 1024 * 1024,
|
||||||
response_format: ResponseFormat::JsonSchema,
|
response_format: ResponseFormat::JsonSchema,
|
||||||
}
|
}
|
||||||
@@ -166,7 +188,13 @@ impl AnalysisConfig {
|
|||||||
d.job_timeout.as_secs(),
|
d.job_timeout.as_secs(),
|
||||||
)),
|
)),
|
||||||
max_tokens: env_u64("ANALYSIS_MAX_TOKENS", d.max_tokens as u64) as u32,
|
max_tokens: env_u64("ANALYSIS_MAX_TOKENS", d.max_tokens as u64) as u32,
|
||||||
max_image_dim: env_u64("ANALYSIS_MAX_IMAGE_DIM", d.max_image_dim as u64) as u32,
|
max_pixels: env_u64("ANALYSIS_MAX_PIXELS", d.max_pixels as u64) as u32,
|
||||||
|
min_slice_height: env_u64("ANALYSIS_MIN_SLICE_HEIGHT", d.min_slice_height as u64)
|
||||||
|
.max(1) as u32,
|
||||||
|
slice_overlap: env_f64("ANALYSIS_SLICE_OVERLAP", d.slice_overlap).clamp(0.0, 0.9),
|
||||||
|
tall_aspect_threshold: env_f64("ANALYSIS_TALL_ASPECT", d.tall_aspect_threshold)
|
||||||
|
.max(1.0),
|
||||||
|
max_slices: env_usize("ANALYSIS_MAX_SLICES", d.max_slices).max(1),
|
||||||
max_image_bytes: env_usize("ANALYSIS_MAX_IMAGE_BYTES", d.max_image_bytes),
|
max_image_bytes: env_usize("ANALYSIS_MAX_IMAGE_BYTES", d.max_image_bytes),
|
||||||
response_format: std::env::var("ANALYSIS_RESPONSE_FORMAT")
|
response_format: std::env::var("ANALYSIS_RESPONSE_FORMAT")
|
||||||
.map(|s| ResponseFormat::from_str(&s))
|
.map(|s| ResponseFormat::from_str(&s))
|
||||||
@@ -506,6 +534,13 @@ fn env_i64(name: &str, default: i64) -> i64 {
|
|||||||
.unwrap_or(default)
|
.unwrap_or(default)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn env_f64(name: &str, default: f64) -> f64 {
|
||||||
|
std::env::var(name)
|
||||||
|
.ok()
|
||||||
|
.and_then(|s| s.parse().ok())
|
||||||
|
.unwrap_or(default)
|
||||||
|
}
|
||||||
|
|
||||||
fn env_usize(name: &str, default: usize) -> usize {
|
fn env_usize(name: &str, default: usize) -> usize {
|
||||||
std::env::var(name)
|
std::env::var(name)
|
||||||
.ok()
|
.ok()
|
||||||
@@ -578,7 +613,11 @@ mod tests {
|
|||||||
"ANALYSIS_MODEL",
|
"ANALYSIS_MODEL",
|
||||||
"ANALYSIS_API_KEY",
|
"ANALYSIS_API_KEY",
|
||||||
"ANALYSIS_MAX_TOKENS",
|
"ANALYSIS_MAX_TOKENS",
|
||||||
"ANALYSIS_MAX_IMAGE_DIM",
|
"ANALYSIS_MAX_PIXELS",
|
||||||
|
"ANALYSIS_MIN_SLICE_HEIGHT",
|
||||||
|
"ANALYSIS_SLICE_OVERLAP",
|
||||||
|
"ANALYSIS_TALL_ASPECT",
|
||||||
|
"ANALYSIS_MAX_SLICES",
|
||||||
"ANALYSIS_RESPONSE_FORMAT",
|
"ANALYSIS_RESPONSE_FORMAT",
|
||||||
] {
|
] {
|
||||||
std::env::remove_var(k);
|
std::env::remove_var(k);
|
||||||
@@ -586,7 +625,10 @@ mod tests {
|
|||||||
let cfg = AnalysisConfig::from_env();
|
let cfg = AnalysisConfig::from_env();
|
||||||
assert!(!cfg.enabled);
|
assert!(!cfg.enabled);
|
||||||
assert_eq!(cfg.workers, 1);
|
assert_eq!(cfg.workers, 1);
|
||||||
assert_eq!(cfg.max_image_dim, 1024);
|
assert_eq!(cfg.max_pixels, 1_000_000);
|
||||||
|
assert_eq!(cfg.min_slice_height, 640);
|
||||||
|
assert_eq!(cfg.max_slices, 16);
|
||||||
|
assert_eq!(cfg.tall_aspect_threshold, 1.6);
|
||||||
assert!(cfg.api_key.is_none());
|
assert!(cfg.api_key.is_none());
|
||||||
assert_eq!(cfg.response_format, ResponseFormat::JsonSchema);
|
assert_eq!(cfg.response_format, ResponseFormat::JsonSchema);
|
||||||
}
|
}
|
||||||
@@ -614,14 +656,18 @@ mod tests {
|
|||||||
std::env::set_var("ANALYSIS_WORKERS", "4");
|
std::env::set_var("ANALYSIS_WORKERS", "4");
|
||||||
std::env::set_var("ANALYSIS_VISION_URL", "http://vis/v1/chat");
|
std::env::set_var("ANALYSIS_VISION_URL", "http://vis/v1/chat");
|
||||||
std::env::set_var("ANALYSIS_MODEL", "qwen2-vl");
|
std::env::set_var("ANALYSIS_MODEL", "qwen2-vl");
|
||||||
std::env::set_var("ANALYSIS_MAX_IMAGE_DIM", "768");
|
std::env::set_var("ANALYSIS_MAX_PIXELS", "768000");
|
||||||
|
std::env::set_var("ANALYSIS_MAX_SLICES", "8");
|
||||||
|
std::env::set_var("ANALYSIS_SLICE_OVERLAP", "0.2");
|
||||||
let cfg = AnalysisConfig::from_env();
|
let cfg = AnalysisConfig::from_env();
|
||||||
for k in [
|
for k in [
|
||||||
"ANALYSIS_ENABLED",
|
"ANALYSIS_ENABLED",
|
||||||
"ANALYSIS_WORKERS",
|
"ANALYSIS_WORKERS",
|
||||||
"ANALYSIS_VISION_URL",
|
"ANALYSIS_VISION_URL",
|
||||||
"ANALYSIS_MODEL",
|
"ANALYSIS_MODEL",
|
||||||
"ANALYSIS_MAX_IMAGE_DIM",
|
"ANALYSIS_MAX_PIXELS",
|
||||||
|
"ANALYSIS_MAX_SLICES",
|
||||||
|
"ANALYSIS_SLICE_OVERLAP",
|
||||||
] {
|
] {
|
||||||
std::env::remove_var(k);
|
std::env::remove_var(k);
|
||||||
}
|
}
|
||||||
@@ -629,7 +675,9 @@ mod tests {
|
|||||||
assert_eq!(cfg.workers, 4);
|
assert_eq!(cfg.workers, 4);
|
||||||
assert_eq!(cfg.endpoint, "http://vis/v1/chat");
|
assert_eq!(cfg.endpoint, "http://vis/v1/chat");
|
||||||
assert_eq!(cfg.model, "qwen2-vl");
|
assert_eq!(cfg.model, "qwen2-vl");
|
||||||
assert_eq!(cfg.max_image_dim, 768);
|
assert_eq!(cfg.max_pixels, 768_000);
|
||||||
|
assert_eq!(cfg.max_slices, 8);
|
||||||
|
assert_eq!(cfg.slice_overlap, 0.2);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
{
|
{
|
||||||
"name": "mangalord-frontend",
|
"name": "mangalord-frontend",
|
||||||
"version": "0.77.1",
|
"version": "0.78.0",
|
||||||
"private": true,
|
"private": true,
|
||||||
"type": "module",
|
"type": "module",
|
||||||
"scripts": {
|
"scripts": {
|
||||||
|
|||||||
Reference in New Issue
Block a user