use crate::error::{RapidOcrError, Result};
use serde::{Deserialize, Serialize};
use std::path::PathBuf;
use std::sync::Arc;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum PixelFormat {
Bgra8,
Rgba8,
Rgb8,
Gray8,
GrayAlpha8,
}
impl PixelFormat {
pub const fn bytes_per_pixel(self) -> usize {
match self {
Self::Bgra8 | Self::Rgba8 => 4,
Self::Rgb8 => 3,
Self::Gray8 => 1,
Self::GrayAlpha8 => 2,
}
}
}
#[derive(Debug, Clone)]
pub struct OwnedPixelBuffer {
pub width: u32,
pub height: u32,
pub stride: usize,
pub format: PixelFormat,
pub bottom_up: bool,
pub data: Arc<[u8]>,
}
impl OwnedPixelBuffer {
pub fn validate(&self) -> Result<()> {
if self.width == 0 || self.height == 0 {
return Err(RapidOcrError::InvalidInput(
"pixel buffer width/height must be greater than zero".into(),
));
}
let row = (self.width as usize)
.checked_mul(self.format.bytes_per_pixel())
.ok_or_else(|| RapidOcrError::InvalidInput("pixel row size overflow".into()))?;
if self.stride < row {
return Err(RapidOcrError::InvalidInput(format!(
"pixel stride {} is smaller than row size {}",
self.stride, row
)));
}
let required = self
.stride
.checked_mul(self.height as usize)
.ok_or_else(|| RapidOcrError::InvalidInput("pixel buffer size overflow".into()))?;
if self.data.len() < required {
return Err(RapidOcrError::InvalidInput(format!(
"pixel buffer has {} bytes, requires at least {}",
self.data.len(),
required
)));
}
Ok(())
}
pub fn to_bgr(&self, roi: Option<RectU32>) -> Result<(Vec<u8>, ImageSize, (u32, u32))> {
self.validate()?;
let selected = roi.unwrap_or(RectU32 {
x: 0,
y: 0,
width: self.width,
height: self.height,
});
selected.validate_against(self.width, self.height)?;
let channels = self.format.bytes_per_pixel();
let row_bytes = selected.width as usize * 3;
let mut out = vec![0; row_bytes * selected.height as usize];
for y in 0..selected.height as usize {
let sy = if self.bottom_up {
self.height as usize - 1 - (selected.y as usize + y)
} else {
selected.y as usize + y
};
let start = sy * self.stride + selected.x as usize * channels;
let src = &self.data[start..start + selected.width as usize * channels];
let dst = &mut out[y * row_bytes..(y + 1) * row_bytes];
for x in 0..selected.width as usize {
let s = &src[x * channels..(x + 1) * channels];
let d = &mut dst[x * 3..x * 3 + 3];
match self.format {
PixelFormat::Bgra8 => {
let a = s[3] as u16;
d[0] = ((s[0] as u16 * a + 255 * (255 - a)) / 255) as u8;
d[1] = ((s[1] as u16 * a + 255 * (255 - a)) / 255) as u8;
d[2] = ((s[2] as u16 * a + 255 * (255 - a)) / 255) as u8;
}
PixelFormat::Rgba8 => {
let a = s[3] as u16;
d[0] = ((s[2] as u16 * a + 255 * (255 - a)) / 255) as u8;
d[1] = ((s[1] as u16 * a + 255 * (255 - a)) / 255) as u8;
d[2] = ((s[0] as u16 * a + 255 * (255 - a)) / 255) as u8;
}
PixelFormat::Rgb8 => d.copy_from_slice(&[s[2], s[1], s[0]]),
PixelFormat::Gray8 => d.fill(s[0]),
PixelFormat::GrayAlpha8 => {
let a = s[1] as u16;
let value = ((s[0] as u16 * a + 255 * (255 - a)) / 255) as u8;
d.fill(value);
}
}
}
}
Ok((
out,
ImageSize {
width: selected.width,
height: selected.height,
},
(selected.x, selected.y),
))
}
}
#[derive(Debug, Clone)]
pub enum ImageInput {
Encoded(Arc<[u8]>),
Pixels(OwnedPixelBuffer),
File(PathBuf),
Url(String),
Image(crate::config::RecImage),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct RectU32 {
pub x: u32,
pub y: u32,
pub width: u32,
pub height: u32,
}
impl RectU32 {
pub fn validate_against(self, w: u32, h: u32) -> Result<()> {
let r = self
.x
.checked_add(self.width)
.ok_or_else(|| RapidOcrError::InvalidInput("ROI x/width overflow".into()))?;
let b = self
.y
.checked_add(self.height)
.ok_or_else(|| RapidOcrError::InvalidInput("ROI y/height overflow".into()))?;
if self.width == 0 || self.height == 0 || r > w || b > h {
return Err(RapidOcrError::InvalidInput(format!(
"ROI {self:?} is outside image {w}x{h}"
)));
}
Ok(())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct TilePolicy {
pub max_width: u32,
pub max_height: u32,
pub overlap: u32,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum EnhancementPolicy {
#[default]
None,
ScreenAdaptive,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PreprocessPolicy {
pub max_decode_pixels: u64,
pub max_encoded_bytes: u64,
pub max_side: Option<u32>,
pub min_text_scale: Option<f32>,
pub tile: Option<TilePolicy>,
pub enhance: EnhancementPolicy,
}
impl Default for PreprocessPolicy {
fn default() -> Self {
Self {
max_decode_pixels: 24_000_000,
max_encoded_bytes: 128 * 1024 * 1024,
max_side: None,
min_text_scale: None,
tile: None,
enhance: EnhancementPolicy::None,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ClassifierPolicy {
Off,
IfAvailable,
Required,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct ClassifierPlan {
pub policy: ClassifierPolicy,
pub apply_rotation: bool,
}
impl Default for ClassifierPlan {
fn default() -> Self {
Self {
policy: ClassifierPolicy::IfAvailable,
apply_rotation: true,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct StagePlan {
pub detect: bool,
pub classify: ClassifierPlan,
pub recognize: bool,
}
impl Default for StagePlan {
fn default() -> Self {
Self {
detect: true,
classify: ClassifierPlan::default(),
recognize: true,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
pub enum WordOutputMode {
#[default]
Off,
Words,
Chars,
}
impl WordOutputMode {
pub fn as_str(self) -> &'static str {
match self {
Self::Off => "off",
Self::Words => "words",
Self::Chars => "chars",
}
}
pub fn from_name(value: &str) -> Option<Self> {
match value {
"off" => Some(Self::Off),
"words" => Some(Self::Words),
"chars" => Some(Self::Chars),
_ => None,
}
}
pub const ALL: [Self; 3] = [Self::Off, Self::Words, Self::Chars];
pub fn kind_for(self) -> WordKind {
match self {
Self::Chars => WordKind::Char,
Self::Off | Self::Words => WordKind::Token,
}
}
}
impl StagePlan {
pub fn validate(self, words: WordOutputMode) -> Result<()> {
if !self.detect
&& !self.recognize
&& matches!(self.classify.policy, ClassifierPolicy::Off)
&& matches!(words, WordOutputMode::Off)
{
return Err(RapidOcrError::InvalidInput(
"at least one OCR stage must be enabled".into(),
));
}
if !self.recognize && !matches!(words, WordOutputMode::Off) {
return Err(RapidOcrError::InvalidInput(
"word output requires recognition stage".into(),
));
}
Ok(())
}
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, Default)]
pub struct DetectionPolicy {
pub box_thresh: Option<f32>,
pub unclip_ratio: Option<f32>,
pub text_score: Option<f32>,
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, Default)]
pub struct RecognitionPolicy {
pub words: WordOutputMode,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum CoordinateSpace {
#[default]
Image,
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct OutputPolicy {
pub coordinate_space: CoordinateSpace,
}
impl Default for OutputPolicy {
fn default() -> Self {
Self {
coordinate_space: CoordinateSpace::Image,
}
}
}
#[derive(Debug, Clone)]
pub struct OcrRequest {
pub input: ImageInput,
pub roi: Option<RectU32>,
pub scale_hint: Option<f32>,
pub stages: StagePlan,
pub preprocess: PreprocessPolicy,
pub detection: DetectionPolicy,
pub recognition: RecognitionPolicy,
pub output: OutputPolicy,
pub formula: FormulaPolicy,
pub table: TablePolicy,
}
impl OcrRequest {
pub fn validate(&self) -> Result<()> {
self.stages.validate(self.recognition.words)?;
self.formula.validate()?;
self.table.validate()?;
if self.formula.enabled {
self.reject_off_page_premise(
"formula",
"formula regions are expressed in original image coordinates",
)?;
}
if self.table.enabled {
self.reject_off_page_premise(
"table",
"the table locator reads the whole page's text layout",
)?;
}
if self.preprocess.max_decode_pixels == 0 {
return Err(RapidOcrError::InvalidInput(
"max_decode_pixels must be greater than zero".into(),
));
}
if self.preprocess.max_encoded_bytes == 0 {
return Err(RapidOcrError::InvalidInput(
"max_encoded_bytes must be greater than zero".into(),
));
}
if self.preprocess.max_side.is_some_and(|v| v == 0) {
return Err(RapidOcrError::InvalidInput(
"max_side must be greater than zero when set".into(),
));
}
if let Some(v) = self.preprocess.min_text_scale
&& (!v.is_finite() || v <= 0.0)
{
return Err(RapidOcrError::InvalidInput(
"min_text_scale must be finite and greater than zero".into(),
));
}
if let Some(v) = self.scale_hint
&& (!v.is_finite() || v <= 0.0)
{
return Err(RapidOcrError::InvalidInput(
"scale_hint must be finite and greater than zero".into(),
));
}
for (name, value) in [
("box_thresh", self.detection.box_thresh),
("text_score", self.detection.text_score),
] {
if let Some(value) = value
&& (!value.is_finite() || !(0.0..=1.0).contains(&value))
{
return Err(RapidOcrError::InvalidInput(format!(
"{name} must be finite and in [0,1]"
)));
}
}
if let Some(value) = self.detection.unclip_ratio
&& (!value.is_finite() || value <= 0.0)
{
return Err(RapidOcrError::InvalidInput(
"unclip_ratio must be finite and greater than zero".into(),
));
}
if let Some(t) = self.preprocess.tile
&& (t.max_width == 0
|| t.max_height == 0
|| t.overlap >= t.max_width
|| t.overlap >= t.max_height)
{
return Err(RapidOcrError::InvalidInput(
"tile dimensions must be positive and overlap must be smaller than dimensions"
.into(),
));
}
Ok(())
}
fn reject_off_page_premise(&self, stage: &str, roi_reason: &str) -> Result<()> {
if self.roi.is_some() {
return Err(RapidOcrError::InvalidInput(format!(
"{stage} routing does not support `roi`: {roi_reason}"
)));
}
if self.preprocess.tile.is_some() {
return Err(RapidOcrError::InvalidInput(format!(
"{stage} routing does not support tiled preprocessing"
)));
}
Ok(())
}
}
#[derive(Debug, Clone, Default)]
pub struct PageRoute {
formula: Option<FormulaPolicy>,
table: Option<TablePolicy>,
}
impl PageRoute {
pub fn runs_formula(&self) -> bool {
self.formula.is_some()
}
pub fn runs_table(&self) -> bool {
self.table.is_some()
}
pub fn formula(&self) -> Option<&FormulaPolicy> {
self.formula.as_ref()
}
pub fn table(&self) -> Option<&TablePolicy> {
self.table.as_ref()
}
}
#[derive(Debug, Clone)]
pub struct ValidatedRequest {
request: OcrRequest,
route: PageRoute,
}
impl ValidatedRequest {
pub fn plan(request: OcrRequest) -> Result<Self> {
request.validate()?;
let route = PageRoute {
formula: request.formula.enabled.then(|| request.formula.clone()),
table: request.table.enabled.then(|| request.table.clone()),
};
Ok(Self { request, route })
}
pub fn request(&self) -> &OcrRequest {
&self.request
}
pub fn route(&self) -> &PageRoute {
&self.route
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct ImageSize {
pub width: u32,
pub height: u32,
}
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub struct Polygon {
pub points: [[f32; 2]; 4],
}
impl Polygon {
pub fn validate(self, size: ImageSize) -> Result<()> {
for [x, y] in self.points {
if !x.is_finite()
|| !y.is_finite()
|| x < 0.0
|| y < 0.0
|| x > size.width as f32
|| y > size.height as f32
{
return Err(RapidOcrError::InvalidInput(format!(
"polygon point ({x}, {y}) is outside image {:?}",
size
)));
}
}
Ok(())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum RegionSource {
Detected { detector_index: usize },
Input,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum TextOrientation {
Deg0,
Deg180,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DetectionOutcome {
pub score: f32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ClassificationOutcome {
pub orientation: TextOrientation,
pub score: f32,
pub applied_rotation: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum WordKind {
Char,
Token,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OcrWord {
pub text: String,
pub score: f32,
pub polygon: Polygon,
pub kind: WordKind,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RecognitionOutcome {
pub text: String,
pub score: f32,
pub words: Option<Vec<OcrWord>>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum RegionKind {
#[default]
Text,
Formula,
Table,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum TableStructureConfidence {
#[default]
High,
Low,
}
impl TableStructureConfidence {
pub fn as_str(self) -> &'static str {
match self {
Self::High => "high",
Self::Low => "low",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case", tag = "kind")]
pub enum TableDegradation {
GridHoles { holes: usize },
Truncated,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum TableCellMatch {
BlankBox,
Unmatched,
Single,
Multiple,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TableOutcome {
pub html: String,
pub table_html: String,
pub markdown: String,
pub rows: usize,
pub cols: usize,
pub confidence: TableStructureConfidence,
pub degradations: Vec<TableDegradation>,
pub holes: usize,
pub cells: Vec<TableCellOutcome>,
pub source_line_indices: Vec<usize>,
pub model_id: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TableCellOutcome {
pub row: usize,
pub col: usize,
pub rowspan: usize,
pub colspan: usize,
pub text: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub score: Option<f32>,
pub match_kind: TableCellMatch,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub polygon: Option<Polygon>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FormulaOutcome {
pub latex: String,
pub eos_index: Option<usize>,
pub truncated: bool,
pub model_id: String,
pub token_ids: Option<Vec<i64>>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FormulaPolicy {
pub enabled: bool,
pub model_path: Option<PathBuf>,
pub expected_model_sha256: Option<String>,
pub detector_path: Option<PathBuf>,
pub expected_detector_sha256: Option<String>,
pub confidence_threshold: f32,
pub iou_threshold: f32,
pub max_regions: usize,
pub min_area_ratio: f32,
pub input_regions: Vec<Polygon>,
pub include_token_ids: bool,
}
impl Default for FormulaPolicy {
fn default() -> Self {
Self {
enabled: false,
model_path: None,
expected_model_sha256: None,
detector_path: None,
expected_detector_sha256: None,
confidence_threshold: crate::formula::detect::DEFAULT_FORMULA_DETECT_CONFIDENCE,
iou_threshold: crate::formula::detect::DEFAULT_FORMULA_DETECT_IOU,
max_regions: 64,
min_area_ratio: 0.000_05,
input_regions: Vec::new(),
include_token_ids: false,
}
}
}
impl FormulaPolicy {
pub fn validate(&self) -> Result<()> {
if !self.enabled {
return Ok(());
}
if self.model_path.is_none() {
return Err(crate::model_verify::missing_model_path_error("formula"));
}
for (name, value) in [
("confidence_threshold", self.confidence_threshold),
("iou_threshold", self.iou_threshold),
("min_area_ratio", self.min_area_ratio),
] {
if !value.is_finite() || !(0.0..=1.0).contains(&value) {
return Err(RapidOcrError::InvalidInput(format!(
"formula {name} must be finite and in [0,1]"
)));
}
}
if self.max_regions == 0 {
return Err(RapidOcrError::InvalidInput(
"formula max_regions must be greater than zero".into(),
));
}
Ok(())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum TableTextPolicy {
#[default]
Escape,
Raw,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TablePolicy {
pub enabled: bool,
pub model_path: Option<PathBuf>,
pub expected_model_sha256: Option<String>,
pub input_side_len: usize,
pub fix_input_shape: bool,
pub intra_threads: Option<usize>,
pub min_region_score: f32,
pub max_regions: usize,
pub text_policy: TableTextPolicy,
}
impl Default for TablePolicy {
fn default() -> Self {
Self {
enabled: false,
model_path: None,
expected_model_sha256: None,
input_side_len: crate::table::structure::TABLE_STRUCTURE_SIDE_LEN,
fix_input_shape: true,
intra_threads: None,
min_region_score: 0.0,
max_regions: 8,
text_policy: TableTextPolicy::Escape,
}
}
}
impl TablePolicy {
pub fn validate(&self) -> Result<()> {
if !self.enabled {
return Ok(());
}
if self.model_path.is_none() {
return Err(crate::model_verify::missing_model_path_error("table"));
}
if self.input_side_len == 0 {
return Err(RapidOcrError::InvalidInput(
"table input_side_len must be greater than zero".into(),
));
}
if !self.min_region_score.is_finite() || !(0.0..=1.0).contains(&self.min_region_score) {
return Err(RapidOcrError::InvalidInput(
"table min_region_score must be finite and in [0,1]".into(),
));
}
if self.max_regions == 0 {
return Err(RapidOcrError::InvalidInput(
"table max_regions must be greater than zero".into(),
));
}
if self.intra_threads == Some(0) {
return Err(RapidOcrError::InvalidInput(
"table intra_threads must be greater than zero when set".into(),
));
}
Ok(())
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OcrRegion {
pub source: RegionSource,
#[serde(default)]
pub kind: RegionKind,
pub polygon: Option<Polygon>,
pub detection: Option<DetectionOutcome>,
pub classification: Option<ClassificationOutcome>,
pub recognition: Option<RecognitionOutcome>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub formula: Option<FormulaOutcome>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub table: Option<TableOutcome>,
}
impl OcrRegion {
pub fn text(
source: RegionSource,
polygon: Option<Polygon>,
detection: Option<DetectionOutcome>,
classification: Option<ClassificationOutcome>,
recognition: Option<RecognitionOutcome>,
) -> Self {
Self {
source,
kind: RegionKind::Text,
polygon,
detection,
classification,
recognition,
formula: None,
table: None,
}
}
pub fn table(
source: RegionSource,
polygon: Option<Polygon>,
detection: Option<DetectionOutcome>,
outcome: TableOutcome,
) -> Self {
Self {
source,
kind: RegionKind::Table,
polygon,
detection,
classification: None,
recognition: None,
formula: None,
table: Some(outcome),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum StageState {
#[default]
Disabled,
SkippedNoInput,
SkippedUnavailable,
Completed {
items: usize,
},
}
#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
pub struct StageTiming {
pub preprocess_ms: f32,
pub infer_ms: f32,
pub postprocess_ms: f32,
}
impl StageTiming {
pub fn total_ms(self) -> f32 {
self.preprocess_ms + self.infer_ms + self.postprocess_ms
}
}
#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize)]
pub struct StageReport {
pub state: StageState,
pub timing: Option<StageTiming>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct InputTimings {
pub decode_ms: Option<f32>,
pub resize_ms: Option<f32>,
pub crop_ms: Option<f32>,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct StageReports {
pub input: InputTimings,
pub detector: StageReport,
pub classifier: StageReport,
pub recognizer: StageReport,
#[serde(default)]
pub formula: StageReport,
#[serde(default)]
pub table: StageReport,
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct OcrTimings {
pub decode_ms: f32,
pub resize_ms: f32,
pub crop_ms: f32,
pub preprocess_ms: f32,
pub detector_preprocess_ms: f32,
pub detector_infer_ms: f32,
pub detector_postprocess_ms: f32,
pub detect_ms: f32,
pub classifier_preprocess_ms: f32,
pub classifier_infer_ms: f32,
pub classifier_postprocess_ms: f32,
pub classify_ms: f32,
pub recognizer_preprocess_ms: f32,
pub recognizer_infer_ms: f32,
pub recognizer_postprocess_ms: f32,
pub recognize_ms: f32,
#[serde(default)]
pub formula_ms: f32,
#[serde(default)]
pub table_detect_ms: f32,
#[serde(default)]
pub table_structure_infer_ms: f32,
#[serde(default)]
pub table_backfill_ms: f32,
pub postprocess_ms: f32,
pub total_ms: f32,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EngineInfo {
pub model_id: String,
pub provider: ProviderInfo,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ImageInfo {
pub original_size: ImageSize,
pub processed_size: ImageSize,
pub coordinate_space: CoordinateSpace,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct OcrOutput {
pub schema_version: u32,
pub image: ImageInfo,
pub stages: StageReports,
pub regions: Vec<OcrRegion>,
pub timings: OcrTimings,
pub engine: EngineInfo,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TextOrder {
Reading,
Detection,
}
impl OcrOutput {
pub fn len(&self) -> usize {
self.regions.len()
}
pub fn text_len(&self) -> usize {
self.regions
.iter()
.filter(|region| region.recognition.is_some())
.count()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn reading_order(&self) -> Vec<usize> {
self.reading_order_groups().into_iter().flatten().collect()
}
pub(crate) fn reading_order_groups(&self) -> Vec<Vec<usize>> {
let mut boxes = Vec::with_capacity(self.regions.len());
let mut without_polygon = Vec::new();
for (id, region) in self.regions.iter().enumerate() {
let Some(polygon) = region.polygon else {
without_polygon.push(id);
continue;
};
let Some(reading_box) = ReadingBox::from_polygon(id, polygon) else {
without_polygon.push(id);
continue;
};
boxes.push(reading_box);
}
let mut groups = order_reading_groups(&boxes);
groups.extend(without_polygon.into_iter().map(|id| vec![id]));
groups
}
pub fn plain_text(&self, order: TextOrder) -> String {
let ids = match order {
TextOrder::Reading => self.reading_order(),
TextOrder::Detection => (0..self.regions.len()).collect(),
};
ids.into_iter()
.filter_map(|i| self.regions[i].recognition.as_ref())
.map(|v| v.text.as_str())
.collect::<Vec<_>>()
.join("\n")
}
pub fn formula_latex(&self, order: TextOrder) -> Vec<(usize, &str)> {
let ids = match order {
TextOrder::Reading => self.reading_order(),
TextOrder::Detection => (0..self.regions.len()).collect(),
};
ids.into_iter()
.filter_map(|id| {
let region = self.regions.get(id)?;
let formula = region.formula.as_ref()?;
Some((id, formula.latex.as_str()))
})
.collect()
}
pub fn formula_count(&self) -> usize {
self.regions
.iter()
.filter(|region| region.kind == RegionKind::Formula)
.count()
}
pub fn table_count(&self) -> usize {
self.regions
.iter()
.filter(|region| region.kind == RegionKind::Table)
.count()
}
pub fn tables(&self, order: TextOrder) -> Vec<(usize, &TableOutcome)> {
let ids = match order {
TextOrder::Reading => self.reading_order(),
TextOrder::Detection => (0..self.regions.len()).collect(),
};
ids.into_iter()
.filter_map(|id| {
let region = self.regions.get(id)?;
Some((id, region.table.as_ref()?))
})
.collect()
}
pub fn validate(&self) -> Result<()> {
if self.schema_version != 1 {
return Err(RapidOcrError::InvalidInput(format!(
"unsupported OCR schema version {}",
self.schema_version
)));
}
for region in &self.regions {
match region.kind {
RegionKind::Formula => {
if region.formula.is_none() {
return Err(RapidOcrError::InvalidInput(
"formula region must carry a formula outcome".into(),
));
}
if region.recognition.is_some() {
return Err(RapidOcrError::InvalidInput(
"formula region must not carry a CTC recognition outcome".into(),
));
}
}
RegionKind::Text => {
if region.formula.is_some() {
return Err(RapidOcrError::InvalidInput(
"text region must not carry a formula outcome".into(),
));
}
if region.table.is_some() {
return Err(RapidOcrError::InvalidInput(
"text region must not carry a table outcome".into(),
));
}
}
RegionKind::Table => {
if region.table.is_none() {
return Err(RapidOcrError::InvalidInput(
"table region must carry a table outcome".into(),
));
}
if region.formula.is_some() {
return Err(RapidOcrError::InvalidInput(
"table region must not carry a formula outcome".into(),
));
}
if region.recognition.is_some() {
return Err(RapidOcrError::InvalidInput(
"table region must not carry a CTC recognition outcome".into(),
));
}
}
}
if let Some(polygon) = region.polygon {
polygon.validate(self.image.original_size)?;
}
if let Some(detection) = ®ion.detection
&& (!detection.score.is_finite() || !(0.0..=1.0).contains(&detection.score))
{
return Err(RapidOcrError::InvalidInput(
"detection score must be finite and in [0,1]".into(),
));
}
if let Some(classification) = ®ion.classification
&& (!classification.score.is_finite()
|| !(0.0..=1.0).contains(&classification.score))
{
return Err(RapidOcrError::InvalidInput(
"classification score must be finite and in [0,1]".into(),
));
}
if let Some(recognition) = ®ion.recognition {
if !recognition.score.is_finite() || !(0.0..=1.0).contains(&recognition.score) {
return Err(RapidOcrError::InvalidInput(
"recognition score must be finite and in [0,1]".into(),
));
}
if let Some(words) = &recognition.words {
for word in words {
if !word.score.is_finite() || !(0.0..=1.0).contains(&word.score) {
return Err(RapidOcrError::InvalidInput(
"word score must be finite and in [0,1]".into(),
));
}
word.polygon.validate(self.image.original_size)?;
}
}
}
}
Ok(())
}
}
#[derive(Debug, Clone, Copy)]
struct ReadingBox {
id: usize,
x0: f32,
y0: f32,
x1: f32,
y1: f32,
}
impl ReadingBox {
fn from_polygon(id: usize, polygon: Polygon) -> Option<Self> {
let mut x0 = f32::INFINITY;
let mut y0 = f32::INFINITY;
let mut x1 = f32::NEG_INFINITY;
let mut y1 = f32::NEG_INFINITY;
for [x, y] in polygon.points {
if !x.is_finite() || !y.is_finite() {
return None;
}
x0 = x0.min(x);
y0 = y0.min(y);
x1 = x1.max(x);
y1 = y1.max(y);
}
Some(Self { id, x0, y0, x1, y1 })
}
fn height(self) -> f32 {
(self.y1 - self.y0).max(0.0)
}
}
#[derive(Debug, Clone, Copy)]
struct Gap {
start: f32,
end: f32,
}
impl Gap {
fn extent(self) -> f32 {
(self.end - self.start).max(0.0)
}
fn center(self) -> f32 {
(self.start + self.end) * 0.5
}
}
#[cfg(test)]
fn order_reading_boxes(boxes: &[ReadingBox]) -> Vec<usize> {
order_reading_groups(boxes).into_iter().flatten().collect()
}
fn order_reading_groups(boxes: &[ReadingBox]) -> Vec<Vec<usize>> {
let mut groups = Vec::with_capacity(boxes.len());
append_reading_order(boxes, &mut groups);
groups
}
fn append_reading_order(boxes: &[ReadingBox], groups: &mut Vec<Vec<usize>>) {
if boxes.is_empty() {
return;
}
if boxes.len() == 1 {
groups.push(vec![boxes[0].id]);
return;
}
let median_height = median_box_height(boxes);
let vertical_gap = largest_vertical_gap(boxes);
let horizontal_gap = largest_horizontal_gap(boxes);
if let Some(gap) = horizontal_gap
&& gap.extent() >= median_height * SECTION_BREAK_MEDIAN_HEIGHT_FACTOR
&& append_horizontal_split(boxes, gap, groups)
{
return;
}
if let Some(gap) = vertical_gap {
let threshold = (median_height * 0.5).max(2.0);
if gap.extent() >= threshold {
let split = gap.center();
let left: Vec<ReadingBox> = boxes.iter().copied().filter(|b| b.x1 <= split).collect();
let right: Vec<ReadingBox> = boxes.iter().copied().filter(|b| b.x0 >= split).collect();
if !left.is_empty() && !right.is_empty() && left.len() + right.len() == boxes.len() {
append_reading_order(&left, groups);
append_reading_order(&right, groups);
return;
}
}
}
if let Some(gap) = horizontal_gap
&& gap.extent() >= 1.0
&& append_horizontal_split(boxes, gap, groups)
{
return;
}
append_reading_order_by_lines(boxes, groups);
}
fn append_horizontal_split(boxes: &[ReadingBox], gap: Gap, groups: &mut Vec<Vec<usize>>) -> bool {
let split = gap.center();
let top: Vec<ReadingBox> = boxes.iter().copied().filter(|b| b.y1 <= split).collect();
let bottom: Vec<ReadingBox> = boxes.iter().copied().filter(|b| b.y0 >= split).collect();
if top.is_empty() || bottom.is_empty() || top.len() + bottom.len() != boxes.len() {
return false;
}
append_reading_order(&top, groups);
append_reading_order(&bottom, groups);
true
}
fn append_reading_order_by_lines(boxes: &[ReadingBox], groups: &mut Vec<Vec<usize>>) {
let mut remaining = boxes.to_vec();
remaining.sort_by(|a, b| {
a.y0.total_cmp(&b.y0)
.then_with(|| a.x0.total_cmp(&b.x0))
.then_with(|| a.id.cmp(&b.id))
});
let mut lines: Vec<Vec<ReadingBox>> = Vec::new();
for item in remaining {
let mut best_line: Option<(usize, f32)> = None;
for (index, line) in lines.iter().enumerate() {
let line_top = line.iter().map(|b| b.y0).fold(f32::INFINITY, f32::min);
let line_bottom = line.iter().map(|b| b.y1).fold(f32::NEG_INFINITY, f32::max);
let line_height = (line_bottom - line_top).max(0.0);
let overlap = (line_bottom.min(item.y1) - line_top.max(item.y0)).max(0.0);
let min_height = line_height.min(item.height()).max(0.0001);
if overlap / min_height >= 0.5
&& best_line.is_none_or(|(_, best_overlap)| overlap > best_overlap)
{
best_line = Some((index, overlap));
}
}
if let Some((index, _)) = best_line {
lines[index].push(item);
} else {
lines.push(vec![item]);
}
}
lines.sort_by(|a, b| {
let a_top = a.iter().map(|b| b.y0).fold(f32::INFINITY, f32::min);
let b_top = b.iter().map(|b| b.y0).fold(f32::INFINITY, f32::min);
a_top.total_cmp(&b_top).then_with(|| {
let a_left = a.iter().map(|b| b.x0).fold(f32::INFINITY, f32::min);
let b_left = b.iter().map(|b| b.x0).fold(f32::INFINITY, f32::min);
a_left.total_cmp(&b_left)
})
});
for mut line in lines {
line.sort_by(|a, b| {
a.x0.total_cmp(&b.x0)
.then_with(|| a.y0.total_cmp(&b.y0))
.then_with(|| a.id.cmp(&b.id))
});
groups.push(line.into_iter().map(|b| b.id).collect());
}
}
const SECTION_BREAK_MEDIAN_HEIGHT_FACTOR: f32 = 1.5;
fn median_box_height(boxes: &[ReadingBox]) -> f32 {
let mut heights: Vec<f32> = boxes
.iter()
.map(|b| b.height())
.filter(|height| height.is_finite() && *height > 0.0)
.collect();
if heights.is_empty() {
return 1.0;
}
heights.sort_by(|a, b| a.total_cmp(b));
heights[heights.len() / 2]
}
fn largest_vertical_gap(boxes: &[ReadingBox]) -> Option<Gap> {
largest_gap(boxes.iter().map(|b| (b.x0, b.x1)))
}
fn largest_horizontal_gap(boxes: &[ReadingBox]) -> Option<Gap> {
largest_gap(boxes.iter().map(|b| (b.y0, b.y1)))
}
fn largest_gap<I>(intervals: I) -> Option<Gap>
where
I: IntoIterator<Item = (f32, f32)>,
{
let mut intervals: Vec<(f32, f32)> = intervals
.into_iter()
.filter(|(a, b)| a.is_finite() && b.is_finite())
.map(|(a, b)| if a <= b { (a, b) } else { (b, a) })
.collect();
if intervals.len() < 2 {
return None;
}
intervals.sort_by(|a, b| a.0.total_cmp(&b.0));
let mut merged: Vec<(f32, f32)> = Vec::with_capacity(intervals.len());
for (start, end) in intervals {
if let Some(last) = merged.last_mut()
&& start <= last.1
{
last.1 = last.1.max(end);
continue;
}
merged.push((start, end));
}
merged
.windows(2)
.map(|pair| Gap {
start: pair[0].1,
end: pair[1].0,
})
.filter(|gap| gap.end > gap.start)
.max_by(|a, b| {
a.extent()
.total_cmp(&b.extent())
.then_with(|| b.start.total_cmp(&a.start))
})
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ResolvedProvider {
Cpu,
DirectMl,
Cuda,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ProviderPreference {
Cpu,
DirectMl { device_id: usize },
Cuda { device_id: usize },
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct ProviderResolutionInfo {
pub requested: ProviderPreference,
pub selected_ep: ResolvedProvider,
pub fallback_to_cpu: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct ProviderInfo {
pub detector: ProviderResolutionInfo,
pub classifier: Option<ProviderResolutionInfo>,
pub recognizer: ProviderResolutionInfo,
}
pub trait OcrEngine {
fn model_id(&self) -> &str;
fn provider(&self) -> ProviderInfo;
fn recognize(&mut self, request: OcrRequest) -> Result<OcrOutput>;
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn rejects_empty_plan() {
let p = StagePlan {
detect: false,
classify: ClassifierPlan {
policy: ClassifierPolicy::Off,
apply_rotation: false,
},
recognize: false,
};
assert!(p.validate(WordOutputMode::Off).is_err());
}
#[test]
fn words_require_recognition() {
let p = StagePlan {
recognize: false,
..StagePlan::default()
};
assert!(p.validate(WordOutputMode::Words).is_err());
}
#[test]
fn the_word_mode_vocabulary_is_a_closed_lowercase_set() {
for mode in WordOutputMode::ALL {
assert_eq!(WordOutputMode::from_name(mode.as_str()), Some(mode));
}
assert_eq!(
WordOutputMode::ALL.map(WordOutputMode::as_str),
["off", "words", "chars"]
);
for bad in [
"Off", "CHARS", "char", "word", "", " off", "off ", "1", "true",
] {
assert_eq!(WordOutputMode::from_name(bad), None, "{bad}");
}
}
#[test]
fn the_word_kind_wire_names_are_snake_case() {
for (kind, name) in [(WordKind::Char, "char"), (WordKind::Token, "token")] {
let value = serde_json::to_value(kind).expect("kind should serialize");
assert_eq!(value, serde_json::json!(name));
let parsed: WordKind = serde_json::from_value(value).expect("kind should deserialize");
assert_eq!(parsed, kind, "{name}");
}
for stale in ["Char", "Token"] {
assert!(serde_json::from_str::<WordKind>(&format!("\"{stale}\"")).is_err());
}
}
#[test]
fn the_word_mode_labels_the_boxes_it_produces() {
assert_eq!(WordOutputMode::Chars.kind_for(), WordKind::Char);
assert_eq!(WordOutputMode::Words.kind_for(), WordKind::Token);
assert_eq!(WordOutputMode::Off.kind_for(), WordKind::Token);
}
#[test]
fn stage_plan_accepts_all_non_empty_combinations() {
let classifier = ClassifierPlan {
policy: ClassifierPolicy::IfAvailable,
apply_rotation: true,
};
for (detect, classify, recognize) in [
(true, false, false),
(false, true, false),
(false, false, true),
(true, true, false),
(true, false, true),
(false, true, true),
(true, true, true),
] {
let plan = StagePlan {
detect,
classify: if classify {
classifier
} else {
ClassifierPlan {
policy: ClassifierPolicy::Off,
apply_rotation: false,
}
},
recognize,
};
assert!(
plan.validate(WordOutputMode::Off).is_ok(),
"invalid combination: {detect}/{classify}/{recognize}"
);
}
}
#[test]
fn polygon_rejects_non_finite_or_out_of_bounds_points() {
let size = ImageSize {
width: 100,
height: 80,
};
assert!(
Polygon {
points: [[0.0, 0.0], [100.0, 0.0], [100.0, 80.0], [0.0, 80.0]]
}
.validate(size)
.is_ok()
);
assert!(
Polygon {
points: [[f32::NAN, 0.0]; 4]
}
.validate(size)
.is_err()
);
assert!(
Polygon {
points: [[101.0, 0.0]; 4]
}
.validate(size)
.is_err()
);
}
#[test]
fn request_validation_rejects_invalid_policy_values() {
let request = OcrRequest {
input: ImageInput::Encoded(Arc::from([] as [u8; 0])),
roi: None,
scale_hint: None,
stages: StagePlan::default(),
preprocess: PreprocessPolicy {
min_text_scale: Some(f32::NAN),
..PreprocessPolicy::default()
},
detection: DetectionPolicy {
text_score: Some(1.5),
..DetectionPolicy::default()
},
recognition: RecognitionPolicy::default(),
output: OutputPolicy::default(),
formula: FormulaPolicy::default(),
table: TablePolicy::default(),
};
assert!(request.validate().is_err());
}
fn page_request(formula_enabled: bool, table_enabled: bool) -> OcrRequest {
let mut request = OcrRequest {
input: ImageInput::Encoded(Arc::from([] as [u8; 0])),
roi: None,
scale_hint: None,
stages: StagePlan::default(),
preprocess: PreprocessPolicy::default(),
detection: DetectionPolicy::default(),
recognition: RecognitionPolicy::default(),
output: OutputPolicy::default(),
formula: FormulaPolicy::default(),
table: TablePolicy::default(),
};
if formula_enabled {
request.formula = FormulaPolicy {
enabled: true,
model_path: Some(PathBuf::from("model.onnx")),
..FormulaPolicy::default()
};
}
if table_enabled {
request.table = TablePolicy {
enabled: true,
model_path: Some(PathBuf::from("model.onnx")),
..TablePolicy::default()
};
}
request
}
#[test]
fn plan_derives_route_from_enabled_stages() {
for (formula_enabled, table_enabled) in
[(false, false), (true, false), (false, true), (true, true)]
{
let planned = ValidatedRequest::plan(page_request(formula_enabled, table_enabled))
.expect("a well-formed page request must plan");
let route = planned.route();
assert_eq!(route.runs_formula(), formula_enabled);
assert_eq!(route.runs_table(), table_enabled);
assert_eq!(route.formula().is_some(), formula_enabled);
assert_eq!(route.table().is_some(), table_enabled);
}
}
#[test]
fn plan_allows_roi_and_tile_without_page_stages() {
let mut request = page_request(false, false);
request.roi = Some(RectU32 {
x: 0,
y: 0,
width: 10,
height: 10,
});
ValidatedRequest::plan(request).expect("plain text may use roi");
let mut request = page_request(false, false);
request.preprocess.tile = Some(TilePolicy {
max_width: 100,
max_height: 100,
overlap: 10,
});
ValidatedRequest::plan(request).expect("plain text may use tiled preprocessing");
}
#[test]
fn plan_rejects_off_page_premise_for_each_enabled_stage() {
for (formula, table, stage) in [
(true, false, "formula"),
(false, true, "table"),
(true, true, "formula"),
] {
let mut with_roi = page_request(formula, table);
with_roi.roi = Some(RectU32 {
x: 0,
y: 0,
width: 10,
height: 10,
});
let message = ValidatedRequest::plan(with_roi)
.expect_err("page stages forbid roi")
.to_string();
assert!(
message.contains(stage) && message.contains("roi"),
"stage={stage} premise=roi message={message}"
);
let mut with_tile = page_request(formula, table);
with_tile.preprocess.tile = Some(TilePolicy {
max_width: 100,
max_height: 100,
overlap: 10,
});
let message = ValidatedRequest::plan(with_tile)
.expect_err("page stages forbid tiled")
.to_string();
assert!(
message.contains(stage) && message.contains("tiled"),
"stage={stage} premise=tiled message={message}"
);
}
}
#[test]
fn gray_alpha_pixels_are_composited() {
let pixels = OwnedPixelBuffer {
width: 1,
height: 1,
stride: 2,
format: PixelFormat::GrayAlpha8,
bottom_up: false,
data: Arc::from([0_u8, 0_u8]),
};
let (bgr, _, _) = pixels.to_bgr(None).expect("valid gray alpha");
assert_eq!(bgr, vec![255, 255, 255]);
}
#[test]
fn default_preprocess_policy_defers_max_side_to_engine_config() {
let policy = PreprocessPolicy::default();
assert_eq!(policy.max_side, None);
assert_eq!(policy.max_encoded_bytes, 128 * 1024 * 1024);
}
fn reading_box(id: usize, x0: f32, y0: f32, x1: f32, y1: f32) -> ReadingBox {
ReadingBox { id, x0, y0, x1, y1 }
}
#[test]
fn reading_order_uses_columns_before_rows() {
let boxes = vec![
reading_box(0, 10.0, 10.0, 60.0, 30.0),
reading_box(1, 10.0, 40.0, 60.0, 60.0),
reading_box(2, 100.0, 10.0, 150.0, 30.0),
reading_box(3, 100.0, 40.0, 150.0, 60.0),
];
assert_eq!(order_reading_boxes(&boxes), vec![0, 1, 2, 3]);
}
#[test]
fn reading_order_handles_full_width_heading_then_columns() {
let boxes = vec![
reading_box(0, 10.0, 0.0, 190.0, 20.0),
reading_box(1, 10.0, 30.0, 80.0, 50.0),
reading_box(2, 10.0, 60.0, 80.0, 80.0),
reading_box(3, 110.0, 30.0, 190.0, 50.0),
reading_box(4, 110.0, 60.0, 190.0, 80.0),
];
assert_eq!(order_reading_boxes(&boxes), vec![0, 1, 2, 3, 4]);
}
#[test]
fn reading_order_handles_three_columns() {
let boxes = vec![
reading_box(0, 0.0, 10.0, 30.0, 30.0),
reading_box(1, 0.0, 40.0, 30.0, 60.0),
reading_box(2, 50.0, 10.0, 80.0, 30.0),
reading_box(3, 50.0, 40.0, 80.0, 60.0),
reading_box(4, 100.0, 10.0, 130.0, 30.0),
reading_box(5, 100.0, 40.0, 130.0, 60.0),
];
assert_eq!(order_reading_boxes(&boxes), vec![0, 1, 2, 3, 4, 5]);
}
#[derive(serde::Deserialize)]
struct RealLayoutFixture {
source: String,
boxes: Vec<RealLayoutBox>,
expected: Vec<usize>,
}
#[derive(serde::Deserialize)]
struct RealLayoutBox {
x0: f32,
y0: f32,
x1: f32,
y1: f32,
}
#[test]
fn reading_order_matches_real_10_columns_layout() {
let fixture: RealLayoutFixture = serde_json::from_str(include_str!(
"../tests/fixtures/reading_order_10_columns.json"
))
.expect("real layout fixture should parse");
let boxes: Vec<ReadingBox> = fixture
.boxes
.iter()
.enumerate()
.map(|(id, b)| reading_box(id, b.x0, b.y0, b.x1, b.y1))
.collect();
assert_eq!(
order_reading_boxes(&boxes),
fixture.expected,
"reading order mismatch for {}",
fixture.source
);
}
#[test]
fn ocr_timings_stage_breakdown_sums_to_each_stage_total() {
let timings = OcrTimings {
decode_ms: 4.0,
resize_ms: 5.0,
crop_ms: 1.0,
preprocess_ms: 12.0,
detector_preprocess_ms: 3.0,
detector_infer_ms: 400.0,
detector_postprocess_ms: 17.0,
detect_ms: 420.0,
classifier_preprocess_ms: 1.0,
classifier_infer_ms: 2.0,
classifier_postprocess_ms: 1.0,
classify_ms: 4.0,
recognizer_preprocess_ms: 4.0,
recognizer_infer_ms: 200.0,
recognizer_postprocess_ms: 30.0,
recognize_ms: 234.0,
formula_ms: 9.0,
table_detect_ms: 0.0,
table_structure_infer_ms: 0.0,
table_backfill_ms: 0.0,
postprocess_ms: 6.0,
total_ms: 12.0 + 420.0 + 4.0 + 234.0 + 6.0 + 9.0,
};
for (name, parts, total) in [
(
"detector",
[
timings.detector_preprocess_ms,
timings.detector_infer_ms,
timings.detector_postprocess_ms,
],
timings.detect_ms,
),
(
"classifier",
[
timings.classifier_preprocess_ms,
timings.classifier_infer_ms,
timings.classifier_postprocess_ms,
],
timings.classify_ms,
),
(
"recognizer",
[
timings.recognizer_preprocess_ms,
timings.recognizer_infer_ms,
timings.recognizer_postprocess_ms,
],
timings.recognize_ms,
),
] {
let sum: f32 = parts.iter().sum();
assert!(
(sum - total).abs() < 1e-3,
"{name} stage breakdown ({sum}) must equal its stage total ({total})"
);
}
let outer = f64::from(timings.preprocess_ms);
assert!(
outer >= 0.0,
"the outer preprocess window cannot be negative: {outer}"
);
}
#[test]
fn a_disabled_table_policy_is_valid_even_without_a_model() {
let policy = TablePolicy::default();
assert!(!policy.enabled);
assert!(policy.model_path.is_none());
policy.validate().expect("a disabled policy must validate");
}
#[test]
fn the_table_policy_rejections_are_distinguishable() {
let model = PathBuf::from("slanet-plus.onnx");
let base = TablePolicy {
enabled: true,
model_path: Some(model),
..TablePolicy::default()
};
base.validate().expect("the baseline must be valid");
let cases: [(&str, TablePolicy, &str); 5] = [
(
"no model path",
TablePolicy {
model_path: None,
..base.clone()
},
"requires `model_path`",
),
(
"zero side length",
TablePolicy {
input_side_len: 0,
..base.clone()
},
"input_side_len must be greater than zero",
),
(
"non-finite score",
TablePolicy {
min_region_score: f32::NAN,
..base.clone()
},
"min_region_score must be finite and in [0,1]",
),
(
"score above one",
TablePolicy {
min_region_score: 1.5,
..base.clone()
},
"min_region_score must be finite and in [0,1]",
),
(
"zero max regions",
TablePolicy {
max_regions: 0,
..base.clone()
},
"max_regions must be greater than zero",
),
];
for (name, policy, expected) in cases {
let error = policy
.validate()
.expect_err(&format!("{name} must be refused"));
assert!(
error.to_string().contains(expected),
"{name}: expected {expected:?}, got {error}"
);
}
let error = TablePolicy {
intra_threads: Some(0),
..base
}
.validate()
.expect_err("intra_threads = 0 must be refused");
assert!(
error
.to_string()
.contains("intra_threads must be greater than zero"),
"unexpected error: {error}"
);
}
#[test]
fn the_table_wire_names_are_snake_case() {
let json = serde_json::to_value(RegionKind::Table).expect("must serialize");
assert_eq!(json, serde_json::json!("table"));
let json = serde_json::to_value(TableStructureConfidence::Low).expect("must serialize");
assert_eq!(json, serde_json::json!("low"));
let json = serde_json::to_value(TableCellMatch::BlankBox).expect("must serialize");
assert_eq!(json, serde_json::json!("blank_box"));
let json =
serde_json::to_value(TableDegradation::GridHoles { holes: 3 }).expect("must serialize");
assert_eq!(json, serde_json::json!({"kind": "grid_holes", "holes": 3}));
let json = serde_json::to_value(TableDegradation::Truncated).expect("must serialize");
assert_eq!(json, serde_json::json!({"kind": "truncated"}));
let json = serde_json::to_value(TableTextPolicy::Escape).expect("must serialize");
assert_eq!(json, serde_json::json!("escape"));
}
#[test]
fn the_confidence_name_agrees_with_serde() {
for confidence in [
TableStructureConfidence::High,
TableStructureConfidence::Low,
] {
let serde_name = serde_json::to_value(confidence).expect("confidence must serialize");
assert_eq!(
serde_name,
serde_json::json!(confidence.as_str()),
"as_str() and serde disagree for {confidence:?}"
);
}
assert_ne!(
TableStructureConfidence::High.as_str(),
TableStructureConfidence::Low.as_str()
);
}
#[test]
fn the_table_policy_round_trips_through_json() {
let policy = TablePolicy {
enabled: true,
model_path: Some(PathBuf::from("models/table/slanet-plus.onnx")),
expected_model_sha256: Some("a".repeat(64)),
input_side_len: 512,
fix_input_shape: false,
intra_threads: Some(4),
min_region_score: 0.25,
max_regions: 3,
text_policy: TableTextPolicy::Raw,
};
let json = serde_json::to_string(&policy).expect("must serialize");
let back: TablePolicy = serde_json::from_str(&json).expect("must deserialize");
assert_eq!(back.input_side_len, 512);
assert_eq!(back.intra_threads, Some(4));
assert_eq!(back.max_regions, 3);
assert_eq!(back.text_policy, TableTextPolicy::Raw);
assert!(!back.fix_input_shape);
}
#[test]
fn an_absent_table_outcome_is_omitted_from_json() {
let region = OcrRegion::text(
RegionSource::Input,
None,
None,
None,
Some(RecognitionOutcome {
text: "ok".into(),
score: 0.5,
words: None,
}),
);
let json = serde_json::to_value(®ion).expect("must serialize");
assert!(
json.get("table").is_none(),
"an absent table outcome must not be serialized: {json}"
);
assert_eq!(json["kind"], serde_json::json!("text"));
}
}