use std::collections::HashMap;
use image::ImageBuffer;
use ort::session::builder::SessionBuilder;
use crate::{
angle_net::AngleNet,
base_net::BaseNet,
crnn_net::CrnnNet,
db_net::DbNet,
layout::{LayoutAnalyzer, LayoutBox},
ocr_error::OcrError,
ocr_result::{OcrResult, Point, TextBlock, WordBox},
ocr_utils::OcrUtils,
scale_param::ScaleParam,
table_classifier::{DocOrientation, DocOrientationClassifier},
};
#[derive(Debug, Clone)]
pub struct OcrOptions {
pub return_word_box: bool,
pub lang: Option<String>,
pub use_doc_orientation: bool,
pub use_doc_unwarping: bool,
pub use_seal: bool,
pub use_formula: bool,
pub use_chart: bool,
}
impl Default for OcrOptions {
fn default() -> Self {
Self {
return_word_box: false,
lang: None,
use_doc_orientation: true, use_doc_unwarping: false,
use_seal: false,
use_formula: false,
use_chart: false,
}
}
}
#[derive(Debug)]
pub struct OcrLite {
db_net: DbNet,
angle_net: AngleNet,
crnn_net: CrnnNet,
doc_orientation_clf: Option<DocOrientationClassifier>,
}
impl Default for OcrLite {
fn default() -> Self {
Self::new()
}
}
impl OcrLite {
pub fn new() -> Self {
Self {
db_net: DbNet::new(),
angle_net: AngleNet::new(),
crnn_net: CrnnNet::new(),
doc_orientation_clf: None,
}
}
pub fn set_doc_orientation_model(&mut self, clf: DocOrientationClassifier) {
self.doc_orientation_clf = Some(clf);
}
pub fn init_models(
&mut self,
det_path: &str,
cls_path: &str,
rec_path: &str,
num_thread: usize,
) -> Result<(), OcrError> {
self.db_net.init_model(det_path, num_thread, None)?;
self.angle_net.init_model(cls_path, num_thread, None)?;
self.crnn_net.init_model(rec_path, num_thread, None)?;
Ok(())
}
pub fn init_models_with_dict(
&mut self,
det_path: &str,
cls_path: &str,
rec_path: &str,
dict_path: &str,
num_thread: usize,
) -> Result<(), OcrError> {
self.db_net.init_model(det_path, num_thread, None)?;
self.angle_net.init_model(cls_path, num_thread, None)?;
self.crnn_net
.init_model_dict_file(rec_path, num_thread, None, dict_path)?;
Ok(())
}
pub fn init_models_no_angle(
&mut self,
det_path: &str,
rec_path: &str,
dict_path: &str,
num_thread: usize,
) -> Result<(), OcrError> {
self.db_net.init_model(det_path, num_thread, None)?;
self.crnn_net
.init_model_dict_file(rec_path, num_thread, None, dict_path)?;
Ok(())
}
pub fn init_models_with_dict_and_builder(
&mut self,
det_path: &str,
cls_path: &str,
rec_path: &str,
dict_path: &str,
num_thread: usize,
builder_fn: Option<fn(ort::session::builder::SessionBuilder) -> Result<ort::session::builder::SessionBuilder, ort::Error>>,
) -> Result<(), OcrError> {
self.db_net.init_model(det_path, num_thread, builder_fn)?;
self.angle_net.init_model(cls_path, num_thread, builder_fn)?;
self.crnn_net
.init_model_dict_file(rec_path, num_thread, builder_fn, dict_path)?;
Ok(())
}
pub fn init_models_custom(
&mut self,
det_path: &str,
cls_path: &str,
rec_path: &str,
builder_fn: fn(SessionBuilder) -> Result<SessionBuilder, ort::Error>,
) -> Result<(), OcrError> {
self.db_net.init_model(det_path, 0, Some(builder_fn))?;
self.angle_net.init_model(cls_path, 0, Some(builder_fn))?;
self.crnn_net.init_model(rec_path, 0, Some(builder_fn))?;
Ok(())
}
pub fn init_models_custom_with_dict(
&mut self,
det_path: &str,
cls_path: &str,
rec_path: &str,
dict_path: &str,
builder_fn: fn(SessionBuilder) -> Result<SessionBuilder, ort::Error>,
) -> Result<(), OcrError> {
self.db_net.init_model(det_path, 0, Some(builder_fn))?;
self.angle_net.init_model(cls_path, 0, Some(builder_fn))?;
self.crnn_net
.init_model_dict_file(rec_path, 0, Some(builder_fn), dict_path)?;
Ok(())
}
pub fn init_models_from_memory(
&mut self,
det_bytes: &[u8],
cls_bytes: &[u8],
rec_bytes: &[u8],
num_thread: usize,
) -> Result<(), OcrError> {
self.db_net
.init_model_from_memory(det_bytes, num_thread, None)?;
self.angle_net
.init_model_from_memory(cls_bytes, num_thread, None)?;
self.crnn_net
.init_model_from_memory(rec_bytes, num_thread, None)?;
Ok(())
}
pub fn init_models_from_memory_custom(
&mut self,
det_bytes: &[u8],
cls_bytes: &[u8],
rec_bytes: &[u8],
builder_fn: fn(SessionBuilder) -> Result<SessionBuilder, ort::Error>,
) -> Result<(), OcrError> {
self.db_net
.init_model_from_memory(det_bytes, 0, Some(builder_fn))?;
self.angle_net
.init_model_from_memory(cls_bytes, 0, Some(builder_fn))?;
self.crnn_net
.init_model_from_memory(rec_bytes, 0, Some(builder_fn))?;
Ok(())
}
fn detect_base(
&mut self,
img_src: &image::RgbImage,
padding: u32,
max_side_len: u32,
box_score_thresh: f32,
box_thresh: f32,
un_clip_ratio: f32,
do_angle: bool,
most_angle: bool,
angle_rollback: bool,
angle_rollback_threshold: f32,
) -> Result<OcrResult, OcrError> {
let origin_max_side = img_src.width().max(img_src.height());
let mut resize;
if max_side_len == 0 || max_side_len > origin_max_side {
resize = origin_max_side;
} else {
resize = max_side_len;
}
resize += 2 * padding;
let padding_src = OcrUtils::make_padding(img_src, padding)?;
let scale = ScaleParam::get_scale_param(&padding_src, resize);
self.detect_once(
&padding_src,
&scale,
padding,
box_score_thresh,
box_thresh,
un_clip_ratio,
do_angle,
most_angle,
angle_rollback,
angle_rollback_threshold,
OcrOptions::default(),
)
}
pub fn detect(
&mut self,
img_src: &image::RgbImage,
padding: u32,
max_side_len: u32,
box_score_thresh: f32,
box_thresh: f32,
un_clip_ratio: f32,
do_angle: bool,
most_angle: bool,
) -> Result<OcrResult, OcrError> {
self.detect_base(
img_src,
padding,
max_side_len,
box_score_thresh,
box_thresh,
un_clip_ratio,
do_angle,
most_angle,
false,
0.0,
)
}
pub fn detect_angle_rollback(
&mut self,
img_src: &image::RgbImage,
padding: u32,
max_side_len: u32,
box_score_thresh: f32,
box_thresh: f32,
un_clip_ratio: f32,
do_angle: bool,
most_angle: bool,
angle_rollback_threshold: f32,
) -> Result<OcrResult, OcrError> {
self.detect_base(
img_src,
padding,
max_side_len,
box_score_thresh,
box_thresh,
un_clip_ratio,
do_angle,
most_angle,
true,
angle_rollback_threshold,
)
}
pub fn detect_from_path(
&mut self,
img_path: &str,
padding: u32,
max_side_len: u32,
box_score_thresh: f32,
box_thresh: f32,
un_clip_ratio: f32,
do_angle: bool,
most_angle: bool,
) -> Result<OcrResult, OcrError> {
let img_src = image::open(img_path)?.to_rgb8();
self.detect(
&img_src,
padding,
max_side_len,
box_score_thresh,
box_thresh,
un_clip_ratio,
do_angle,
most_angle,
)
}
fn detect_once(
&mut self,
img_src: &image::RgbImage,
scale: &ScaleParam,
padding: u32,
box_score_thresh: f32,
box_thresh: f32,
un_clip_ratio: f32,
do_angle: bool,
most_angle: bool,
angle_rollback: bool,
angle_rollback_threshold: f32,
options: OcrOptions,
) -> Result<OcrResult, OcrError> {
let text_boxes = self.db_net.get_text_boxes(
img_src,
scale,
box_score_thresh,
box_thresh,
un_clip_ratio,
)?;
let part_images = OcrUtils::get_part_images(img_src, &text_boxes);
let angles = self
.angle_net
.get_angles(&part_images, do_angle, most_angle)?;
let mut rotated_images: Vec<image::RgbImage> = Vec::with_capacity(part_images.len());
let mut angle_rollback_records =
HashMap::<usize, ImageBuffer<image::Rgb<u8>, Vec<u8>>>::new();
for (index, (angle, mut part_image)) in
angles.iter().zip(part_images.into_iter()).enumerate()
{
if angle.index == 1 {
if angle_rollback {
angle_rollback_records.insert(index, part_image.clone());
}
OcrUtils::mat_rotate_clock_wise_180(&mut part_image);
}
rotated_images.push(part_image);
}
let lines_meta = self.crnn_net.get_text_lines_with_word_ranges(
&rotated_images,
&angle_rollback_records,
angle_rollback_threshold,
)?;
let mut text_blocks = Vec::with_capacity(lines_meta.len());
for (i, (text_line, word_ranges, crop_size, target_w, t_steps)) in lines_meta.into_iter().enumerate() {
let box_points: Vec<Point> = text_boxes[i].points.iter().map(|p| Point {
x: ((p.x as f32) - padding as f32) as u32,
y: ((p.y as f32) - padding as f32) as u32,
}).collect();
let words: Vec<WordBox> = if options.return_word_box {
build_word_boxes(
&word_ranges,
&text_boxes[i].points, crop_size,
target_w,
t_steps,
angles[i].index == 1,
padding,
)
} else {
Vec::new()
};
text_blocks.push(TextBlock {
box_points,
box_score: text_boxes[i].score,
angle_index: angles[i].index,
angle_score: angles[i].score,
text: text_line.text,
text_score: text_line.text_score,
words,
});
}
Ok(OcrResult { text_blocks, page_angle: 0 })
}
pub fn detect_with_layout(
&mut self,
image: &image::RgbImage,
layout: &mut LayoutAnalyzer,
padding: u32,
max_side_len: u32,
box_score_thresh: f32,
box_thresh: f32,
un_clip_ratio: f32,
do_angle: bool,
most_angle: bool,
options: OcrOptions,
) -> Result<LayoutAwareResult, OcrError> {
let layout_boxes = layout.analyze(image)?;
let ocr = self.detect_with_options(
image,
padding,
max_side_len,
box_score_thresh,
box_thresh,
un_clip_ratio,
do_angle,
most_angle,
options,
)?;
let mut associated: Vec<TextBlockWithLayout> = ocr.text_blocks
.into_iter()
.map(|tb| {
let (cx, cy) = OcrUtils::polygon_centroid(&tb.box_points);
let contained = layout_boxes.iter().position(|lb| lb.contains(cx, cy));
let (idx, dist) = match contained {
Some(i) => (Some(i), 0.0),
None => nearest_layout_box(&layout_boxes, cx, cy),
};
TextBlockWithLayout {
block: tb,
layout_index: idx,
distance: dist,
centroid_x: cx,
centroid_y: cy,
}
})
.collect();
associated.sort_by(|a, b| {
let ra = a.layout_index
.map(|i| layout_boxes[i].reading_order)
.filter(|&r| r >= 0)
.unwrap_or(i32::MAX);
let rb = b.layout_index
.map(|i| layout_boxes[i].reading_order)
.filter(|&r| r >= 0)
.unwrap_or(i32::MAX);
ra.cmp(&rb)
.then_with(|| a.centroid_y.cmp(&b.centroid_y))
.then_with(|| a.centroid_x.cmp(&b.centroid_x))
});
Ok(LayoutAwareResult {
layout_boxes,
blocks: associated,
})
}
fn rotate_to_upright(img: image::RgbImage, orient: DocOrientation) -> image::RgbImage {
match orient {
DocOrientation::Deg0 => img,
DocOrientation::Deg90 => image::imageops::rotate270(&img),
DocOrientation::Deg180 => image::imageops::rotate180(&img),
DocOrientation::Deg270 => image::imageops::rotate90(&img),
}
}
pub fn detect_with_options(
&mut self,
img_src: &image::RgbImage,
padding: u32,
max_side_len: u32,
box_score_thresh: f32,
box_thresh: f32,
un_clip_ratio: f32,
do_angle: bool,
most_angle: bool,
options: OcrOptions,
) -> Result<OcrResult, OcrError> {
if options.use_doc_orientation {
match &self.doc_orientation_clf {
None => {} Some(clf) => {
let (orient, _conf) = clf.classify(img_src)?;
if orient != DocOrientation::Deg0 {
let rotated = Self::rotate_to_upright(img_src.clone(), orient);
let mut opts2 = options.clone();
opts2.use_doc_orientation = false; let mut result = self.detect_with_options(
&rotated, padding, max_side_len,
box_score_thresh, box_thresh, un_clip_ratio,
do_angle, most_angle, opts2,
)?;
result.page_angle = orient.degrees();
return Ok(result);
}
}
}
}
if options.use_doc_unwarping {
eprintln!("[ppocr-rs] WARN: use_doc_unwarping non implementato (TextImageUnwarping)");
}
if options.use_seal {
eprintln!("[ppocr-rs] WARN: use_seal non implementato (SealTextDet + SealTextRec)");
}
if options.use_formula {
eprintln!("[ppocr-rs] WARN: use_formula non implementato (PP-FormulaNet-L)");
}
if options.use_chart {
eprintln!("[ppocr-rs] WARN: use_chart non implementato (ChartRecognition)");
}
let origin_max_side = img_src.width().max(img_src.height());
let mut resize;
if max_side_len == 0 || max_side_len > origin_max_side {
resize = origin_max_side;
} else {
resize = max_side_len;
}
resize += 2 * padding;
let padding_src = OcrUtils::make_padding(img_src, padding)?;
let scale = ScaleParam::get_scale_param(&padding_src, resize);
self.detect_once(
&padding_src,
&scale,
padding,
box_score_thresh,
box_thresh,
un_clip_ratio,
do_angle,
most_angle,
false,
0.0,
options,
)
}
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct LayoutAwareResult {
pub layout_boxes: Vec<LayoutBox>,
pub blocks: Vec<TextBlockWithLayout>,
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct TextBlockWithLayout {
pub block: TextBlock,
pub layout_index: Option<usize>,
pub distance: f32,
pub centroid_x: u32,
pub centroid_y: u32,
}
fn nearest_layout_box(boxes: &[LayoutBox], cx: u32, cy: u32) -> (Option<usize>, f32) {
if boxes.is_empty() { return (None, f32::INFINITY); }
let mut best_idx = 0usize;
let mut best_dist = f32::INFINITY;
for (i, lb) in boxes.iter().enumerate() {
let d = lb.distance_to(cx, cy);
if d < best_dist {
best_dist = d;
best_idx = i;
}
}
(Some(best_idx), best_dist)
}
fn build_word_boxes(
word_ranges: &[crate::crnn_net::WordRange],
polygon_padded: &[Point], crop_size: (u32, u32),
target_w: usize,
t_steps: usize,
was_180_rotated: bool,
padding: u32,
) -> Vec<WordBox> {
if word_ranges.is_empty() || polygon_padded.len() != 4 || t_steps == 0 || target_w == 0 {
return Vec::new();
}
let (crop_w, crop_h) = crop_size;
if crop_h == 0 || crop_w == 0 { return Vec::new(); }
if crop_h >= crop_w {
return Vec::new();
}
let dst_h = crate::crnn_net::CRNN_DST_HEIGHT as f32;
let resized_w = ((crop_w as f32) * dst_h / (crop_h as f32)).ceil() as u32;
let resized_w = resized_w.min(target_w as u32).max(1);
let x_per_ts = (target_w as f32) / (t_steps as f32).max(1.0);
let poly: [Point; 4] = [
polygon_padded[0], polygon_padded[1], polygon_padded[2], polygon_padded[3],
];
let mut out = Vec::with_capacity(word_ranges.len());
for w in word_ranges {
let x_crnn_start = (w.start_ts as f32) * x_per_ts;
let x_crnn_end = ((w.end_ts + 1) as f32) * x_per_ts;
let x_crnn_start = x_crnn_start.min(resized_w as f32);
let x_crnn_end = x_crnn_end .min(resized_w as f32);
let ratio = (crop_w as f32) / (resized_w as f32);
let mut x_crop_start = x_crnn_start * ratio;
let mut x_crop_end = x_crnn_end * ratio;
if x_crop_end <= x_crop_start { continue; }
if was_180_rotated {
let new_start = (crop_w as f32) - x_crop_end;
let new_end = (crop_w as f32) - x_crop_start;
x_crop_start = new_start.max(0.0);
x_crop_end = new_end.max(0.0);
}
let quad = [
(x_crop_start, 0.0),
(x_crop_end, 0.0),
(x_crop_end, crop_h as f32),
(x_crop_start, crop_h as f32),
];
if let Some(image_pts_padded) = OcrUtils::inverse_warp_quad(&poly, crop_size, &quad) {
let image_pts: Vec<Point> = image_pts_padded.iter().map(|p| Point {
x: p.x.saturating_sub(padding),
y: p.y.saturating_sub(padding),
}).collect();
out.push(WordBox {
text: w.text.clone(),
box_points: image_pts,
score: w.score,
});
}
}
out
}