use std::path::Path;
use image::{ExtendedColorType, ImageEncoder, ImageFormat, ImageReader, Limits, RgbImage};
use crate::runner::protocol::bodies::{OcrInput, OcrOk, OcrRequest, OcrSpan};
pub const IMAGE_OCR_MAX_AREA_PX: u64 = 2_359_296;
pub const IMAGE_OCR_VALIDATED_LONG_EDGE: u32 = 1600;
pub const IMAGE_OCR_DECODE_ALLOC: u64 = 64 * 1024 * 1024;
#[derive(Debug)]
pub struct ImageOcrError {
pub code: &'static str,
pub message: String,
}
impl ImageOcrError {
fn new(code: &'static str, message: impl Into<String>) -> ImageOcrError {
ImageOcrError {
code,
message: message.into(),
}
}
}
fn stage_pixels(request: &OcrRequest) -> Result<(RgbImage, Vec<String>), ImageOcrError> {
if !matches!(request.kind, OcrInput::Image) {
return Err(ImageOcrError::new(
"ocr-protocol-error",
"the image-pixel-ocr worker recognizes standalone rasters only",
));
}
let input = Path::new(&request.input);
let format = input
.extension()
.and_then(|e| e.to_str())
.unwrap_or_default();
let bytes = std::fs::read(input).map_err(|e| {
ImageOcrError::new(
"io",
format!("cannot read the staged input {}: {e}", request.input),
)
})?;
decode_to_rgb(&bytes, format)
}
pub fn recognize(request: &OcrRequest) -> Result<OcrOk, ImageOcrError> {
let (rgb, warnings) = stage_pixels(request)?;
let spans = recognize_pixels(&rgb)?;
Ok(OcrOk { spans, warnings })
}
#[cfg(feature = "test-adapters")]
pub fn recognize_fake(request: &OcrRequest) -> Result<OcrOk, ImageOcrError> {
let (rgb, warnings) = stage_pixels(request)?;
let spans = recognize_pixels_fake(&rgb)?;
Ok(OcrOk { spans, warnings })
}
fn decode_limits() -> Limits {
let mut limits = Limits::no_limits();
limits.max_alloc = Some(IMAGE_OCR_DECODE_ALLOC);
limits
}
fn decode_to_rgb(bytes: &[u8], format: &str) -> Result<(RgbImage, Vec<String>), ImageOcrError> {
let fmt = match format {
"png" => ImageFormat::Png,
"jpeg" => ImageFormat::Jpeg,
"webp" => ImageFormat::WebP,
other => {
return Err(ImageOcrError::new(
"image-ocr-decode-failed",
format!("the image-pixel-ocr worker does not decode {other}"),
));
}
};
let decode_failed = |e: image::ImageError| {
ImageOcrError::new(
"image-ocr-decode-failed",
format!("cannot decode the raster: {e}"),
)
};
let mut reader = ImageReader::with_format(std::io::Cursor::new(bytes), fmt);
reader.limits(decode_limits());
let (width, height) = reader.into_dimensions().map_err(decode_failed)?;
let area = u64::from(width) * u64::from(height);
if area > IMAGE_OCR_MAX_AREA_PX {
return Err(ImageOcrError::new(
"image-ocr-area-exceeded",
format!(
"{width}x{height} = {area} px over the {IMAGE_OCR_MAX_AREA_PX} px encoder-input area cap"
),
));
}
let mut warnings = Vec::new();
if width.max(height) > IMAGE_OCR_VALIDATED_LONG_EDGE {
warnings.push(format!(
"image_ocr_long_edge: {width}x{height} within the area cap but outside validated evidence, logged for validation"
));
}
let mut reader = ImageReader::with_format(std::io::Cursor::new(bytes), fmt);
reader.limits(decode_limits());
let image = reader.decode().map_err(decode_failed)?;
if (image.width(), image.height()) != (width, height) {
return Err(ImageOcrError::new(
"image-ocr-decode-failed",
"the decoded dimensions disagree with the header",
));
}
Ok((image.to_rgb8(), warnings))
}
fn encode_png(rgb: &RgbImage) -> Result<Vec<u8>, ImageOcrError> {
let mut buffer = Vec::new();
image::codecs::png::PngEncoder::new(&mut buffer)
.write_image(
rgb.as_raw(),
rgb.width(),
rgb.height(),
ExtendedColorType::Rgb8,
)
.map_err(|e| {
ImageOcrError::new(
"image-ocr-decode-failed",
format!("cannot re-encode the canonical raster: {e}"),
)
})?;
Ok(buffer)
}
#[cfg(feature = "test-adapters")]
fn recognize_pixels_fake(rgb: &RgbImage) -> Result<Vec<OcrSpan>, ImageOcrError> {
let _canonical = encode_png(rgb)?;
let digest = crate::hash::hash_bytes(rgb.as_raw());
Ok(vec![
OcrSpan {
text: format!(
"recognized from a {}x{} raster {}",
rgb.width(),
rgb.height(),
&digest[..16]
),
confidence: 0.94,
},
OcrSpan {
text: "faint line".to_string(),
confidence: 0.31,
},
])
}
fn recognize_pixels(rgb: &RgbImage) -> Result<Vec<OcrSpan>, ImageOcrError> {
let canonical = encode_png(rgb)?;
let components = runtime::resolve_components();
runtime::verify_inventory(&runtime::CORE_ROLES, &components)?;
let devices = runtime::enumerate_devices(&components);
runtime::check_accelerator(&devices)?;
let stdout = runtime::engine_stdout(&components, &canonical)?;
let text = runtime::parse_fenced_envelope(&stdout)?;
Ok(vec![OcrSpan {
text,
confidence: 1.0,
}])
}
mod runtime {
use std::path::PathBuf;
use super::ImageOcrError;
pub(super) const CORE_ROLES: [&str; 5] = [
"engine-cli",
"vision-weights",
"vision-projector",
"ocr-prompt",
"ocr-limit-policy",
];
pub(super) const DETERMINISM_PROOF_PENDING: &str = "live-engine cross-restart determinism proof pending; required before the engine path is enabled";
pub(super) const BEGIN_SENTINEL: &[u8] = b"\x1eTM-OCR-BEGIN\x1e";
pub(super) const END_SENTINEL: &[u8] = b"\x1eTM-OCR-END\x1e";
pub(super) struct Component {
pub(super) role: &'static str,
pub(super) path: PathBuf,
pub(super) expected_blake3: String,
}
pub(super) fn resolve_components() -> Vec<Component> {
Vec::new()
}
pub(super) fn verify_inventory(
required: &[&str],
resolved: &[Component],
) -> Result<(), ImageOcrError> {
for role in required {
let Some(component) = resolved.iter().find(|component| component.role == *role) else {
return Err(ImageOcrError::new(
"image-ocr-runtime-missing",
format!("runtime component {role} is not present"),
));
};
let actual = crate::hash::hash_file(&component.path).map_err(|_| {
ImageOcrError::new(
"image-ocr-runtime-missing",
format!("runtime component {role} cannot be read"),
)
})?;
if actual != component.expected_blake3 {
return Err(ImageOcrError::new(
"image-ocr-runtime-mismatch",
format!("runtime component {role} does not match its pinned hash"),
));
}
}
Ok(())
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum DeviceClass {
Accelerator,
#[cfg_attr(not(test), allow(dead_code))]
Cpu,
}
pub(super) struct Device {
pub(super) class: DeviceClass,
pub(super) selected: bool,
}
pub(super) fn enumerate_devices(_components: &[Component]) -> Vec<Device> {
Vec::new()
}
pub(super) fn check_accelerator(devices: &[Device]) -> Result<(), ImageOcrError> {
if devices
.iter()
.any(|device| device.class == DeviceClass::Accelerator && device.selected)
{
Ok(())
} else {
Err(ImageOcrError::new(
"image-ocr-no-accelerator",
"no accelerator-class device is present and selected, refusing a cpu fallback",
))
}
}
pub(super) fn engine_stdout(
_components: &[Component],
_canonical_png: &[u8],
) -> Result<Vec<u8>, ImageOcrError> {
Err(ImageOcrError::new(
"image-ocr-runtime-missing",
format!(
"no pinned vision runtime is wired into this build; {DETERMINISM_PROOF_PENDING}"
),
))
}
pub(super) fn parse_fenced_envelope(stdout: &[u8]) -> Result<String, ImageOcrError> {
let protocol_error = |detail: &str| {
ImageOcrError::new(
"ocr-protocol-error",
format!("engine stdout is not a clean recognition envelope: {detail}"),
)
};
let begins = count_occurrences(stdout, BEGIN_SENTINEL);
let ends = count_occurrences(stdout, END_SENTINEL);
if begins != 1 {
return Err(protocol_error(
"the BEGIN sentinel does not appear exactly once",
));
}
if ends != 1 {
return Err(protocol_error(
"the END sentinel does not appear exactly once",
));
}
let begin = find(stdout, BEGIN_SENTINEL).expect("one BEGIN was counted");
let end = find(stdout, END_SENTINEL).expect("one END was counted");
let text_start = begin + BEGIN_SENTINEL.len();
if begin != 0 {
return Err(protocol_error("bytes appear before the BEGIN sentinel"));
}
if text_start > end {
return Err(protocol_error("the sentinels are out of order or overlap"));
}
let after_end = end + END_SENTINEL.len();
let tail = &stdout[after_end..];
if !(tail.is_empty() || tail == b"\n") {
return Err(protocol_error("bytes appear after the END sentinel"));
}
let text_bytes = &stdout[text_start..end];
String::from_utf8(text_bytes.to_vec())
.map_err(|_| protocol_error("the recognized text is not valid utf-8"))
}
pub(super) fn count_occurrences(haystack: &[u8], needle: &[u8]) -> usize {
if needle.is_empty() || haystack.len() < needle.len() {
return 0;
}
let mut count = 0;
let mut index = 0;
while index + needle.len() <= haystack.len() {
if &haystack[index..index + needle.len()] == needle {
count += 1;
index += needle.len();
} else {
index += 1;
}
}
count
}
pub(super) fn find(haystack: &[u8], needle: &[u8]) -> Option<usize> {
if needle.is_empty() || haystack.len() < needle.len() {
return None;
}
(0..=haystack.len() - needle.len()).find(|&i| &haystack[i..i + needle.len()] == needle)
}
}
#[cfg(test)]
mod tests {
use super::runtime::{
BEGIN_SENTINEL, CORE_ROLES, Component, DETERMINISM_PROOF_PENDING, Device, DeviceClass,
END_SENTINEL, check_accelerator, engine_stdout, enumerate_devices, parse_fenced_envelope,
resolve_components, verify_inventory,
};
use super::*;
use std::io::Write;
fn png(width: u32, height: u32) -> Vec<u8> {
let image = RgbImage::from_pixel(width, height, image::Rgb([200, 40, 40]));
encode_png(&image).unwrap()
}
fn fence(text: &[u8]) -> Vec<u8> {
let mut out = Vec::new();
out.extend_from_slice(BEGIN_SENTINEL);
out.extend_from_slice(text);
out.extend_from_slice(END_SENTINEL);
out
}
#[test]
fn decodes_a_small_raster_to_rgb() {
let (rgb, warnings) = decode_to_rgb(&png(8, 8), "png").unwrap();
assert_eq!((rgb.width(), rgb.height()), (8, 8));
assert!(warnings.is_empty());
}
#[test]
fn each_supported_format_decodes() {
let (rgb, _) = decode_to_rgb(&png(16, 16), "png").unwrap();
let mut jpeg = Vec::new();
image::codecs::jpeg::JpegEncoder::new(&mut jpeg)
.encode_image(&rgb)
.unwrap();
assert!(decode_to_rgb(&jpeg, "jpeg").is_ok());
let mut webp = Vec::new();
image::codecs::webp::WebPEncoder::new_lossless(&mut webp)
.encode(
rgb.as_raw(),
rgb.width(),
rgb.height(),
ExtendedColorType::Rgb8,
)
.unwrap();
assert!(decode_to_rgb(&webp, "webp").is_ok());
}
#[test]
fn the_area_cap_fails_closed_with_no_resize() {
let err = decode_to_rgb(&png(2048, 2048), "png").unwrap_err();
assert_eq!(err.code, "image-ocr-area-exceeded");
assert!(decode_to_rgb(&png(1536, 1536), "png").is_ok());
let (_, warnings) = decode_to_rgb(&png(1600, 900), "png").unwrap();
assert!(warnings.is_empty());
}
#[test]
fn a_long_edge_within_the_area_cap_warns_once_and_proceeds() {
let (rgb, warnings) = decode_to_rgb(&png(3000, 786), "png").unwrap();
assert_eq!((rgb.width(), rgb.height()), (3000, 786));
assert_eq!(warnings.len(), 1);
assert!(warnings[0].starts_with("image_ocr_long_edge:"));
}
#[test]
fn alpha_is_flattened_so_transparent_and_opaque_twins_match() {
let transparent = image::RgbaImage::from_pixel(4, 4, image::Rgba([10, 20, 30, 0]));
let opaque = image::RgbaImage::from_pixel(4, 4, image::Rgba([10, 20, 30, 255]));
let mut a = Vec::new();
image::codecs::png::PngEncoder::new(&mut a)
.write_image(transparent.as_raw(), 4, 4, ExtendedColorType::Rgba8)
.unwrap();
let mut b = Vec::new();
image::codecs::png::PngEncoder::new(&mut b)
.write_image(opaque.as_raw(), 4, 4, ExtendedColorType::Rgba8)
.unwrap();
let (rgb_a, _) = decode_to_rgb(&a, "png").unwrap();
let (rgb_b, _) = decode_to_rgb(&b, "png").unwrap();
assert_eq!(rgb_a.as_raw(), rgb_b.as_raw());
}
#[test]
fn corrupt_and_truncated_bytes_fail_with_decode_failed() {
for bytes in [
b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR".as_slice(),
b"not an image at all".as_slice(),
] {
let err = decode_to_rgb(bytes, "png").unwrap_err();
assert_eq!(err.code, "image-ocr-decode-failed");
}
}
#[test]
fn a_clean_envelope_parses_with_and_without_a_trailing_newline() {
assert_eq!(
parse_fenced_envelope(&fence(b"hello world")).unwrap(),
"hello world"
);
let mut with_newline = fence(b"line");
with_newline.push(b'\n');
assert_eq!(parse_fenced_envelope(&with_newline).unwrap(), "line");
}
#[test]
fn every_envelope_violation_is_a_protocol_error() {
let mut doubled_begin = Vec::new();
doubled_begin.extend_from_slice(BEGIN_SENTINEL);
doubled_begin.extend_from_slice(&fence(b"x"));
let mut out_of_order = Vec::new();
out_of_order.extend_from_slice(END_SENTINEL);
out_of_order.extend_from_slice(b"x");
out_of_order.extend_from_slice(BEGIN_SENTINEL);
let mut leading = b"banner\n".to_vec();
leading.extend_from_slice(&fence(b"x"));
let mut trailing = fence(b"x");
trailing.extend_from_slice(b"timing 12ms");
let mut leaked_log = fence(b"x");
leaked_log.extend_from_slice(b"\nengine: loaded\n");
for bytes in [
BEGIN_SENTINEL.to_vec(), END_SENTINEL.to_vec(), doubled_begin, out_of_order, leading, trailing, leaked_log, b"!!!!!!!!".to_vec(), ] {
let err = parse_fenced_envelope(&bytes).unwrap_err();
assert_eq!(err.code, "ocr-protocol-error", "{bytes:?}");
}
}
#[test]
fn the_core_build_resolves_no_runtime_so_recognition_fails_closed() {
assert!(resolve_components().is_empty());
let err = verify_inventory(&CORE_ROLES, &resolve_components()).unwrap_err();
assert_eq!(err.code, "image-ocr-runtime-missing");
}
#[test]
fn the_inventory_flags_missing_and_mismatch_distinctly() {
let missing = verify_inventory(&CORE_ROLES, &[]).unwrap_err();
assert_eq!(missing.code, "image-ocr-runtime-missing");
let mut file = tempfile::NamedTempFile::new().unwrap();
file.write_all(b"weights").unwrap();
file.flush().unwrap();
let good = crate::hash::hash_file(file.path()).unwrap();
let present = [Component {
role: "engine-cli",
path: file.path().to_path_buf(),
expected_blake3: good.clone(),
}];
assert!(verify_inventory(&["engine-cli"], &present).is_ok());
let wrong = [Component {
role: "engine-cli",
path: file.path().to_path_buf(),
expected_blake3: "0".repeat(64),
}];
let mismatch = verify_inventory(&["engine-cli"], &wrong).unwrap_err();
assert_eq!(mismatch.code, "image-ocr-runtime-mismatch");
}
#[test]
fn the_accelerator_check_fails_closed_without_a_selected_accelerator() {
let empty = check_accelerator(&enumerate_devices(&[])).unwrap_err();
assert_eq!(empty.code, "image-ocr-no-accelerator");
let cpu = check_accelerator(&[Device {
class: DeviceClass::Cpu,
selected: true,
}])
.unwrap_err();
assert_eq!(cpu.code, "image-ocr-no-accelerator");
assert!(
check_accelerator(&[Device {
class: DeviceClass::Accelerator,
selected: false,
}])
.is_err()
);
assert!(
check_accelerator(&[Device {
class: DeviceClass::Accelerator,
selected: true,
}])
.is_ok()
);
}
#[test]
fn the_engine_exec_fails_closed_and_carries_the_pending_determinism_note() {
let err = engine_stdout(&[], b"png").unwrap_err();
assert_eq!(err.code, "image-ocr-runtime-missing");
assert!(!DETERMINISM_PROOF_PENDING.is_empty());
assert!(err.message.contains(DETERMINISM_PROOF_PENDING));
}
#[test]
fn the_production_recognize_fails_closed_in_every_feature_combination() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("input.png");
std::fs::write(&path, png(8, 8)).unwrap();
let request = OcrRequest {
input: path.to_str().unwrap().to_string(),
kind: OcrInput::Image,
};
let err = recognize(&request).unwrap_err();
assert_eq!(err.code, "image-ocr-runtime-missing");
}
#[test]
fn a_well_formed_but_content_free_envelope_is_empty_output_not_protocol_error() {
let text = parse_fenced_envelope(&fence(b" \t\r\n\x07")).unwrap();
assert!(!text.chars().any(crate::convert::is_meaningful));
let span = OcrSpan {
text,
confidence: 1.0,
};
let err = crate::convert::subprocess::render_ocr_spans(&[span], 128).unwrap_err();
assert_eq!(err.code, "empty_output");
}
#[test]
fn the_inventory_is_scoped_to_the_five_core_roles() {
assert_eq!(CORE_ROLES.len(), 5);
assert!(!CORE_ROLES.iter().any(|role| role.starts_with("raster-")));
}
}