use image::imageops::{rotate180, rotate270, rotate90};
use image::RgbImage;
use crate::layout::Region;
use crate::ocr::OcrModel;
use crate::ocr_det::DetBox;
use crate::ocr_prep::{prep_region_lines, prep_region_lines_det, PrepLine};
use crate::Recognizer;
use docling_core::debug_log;
pub(crate) fn enabled() -> bool {
let raw = docling_core::env::nonempty("DOCLING_RS_OCR_ORIENTATION").unwrap_or_default();
let v = raw.trim().to_ascii_lowercase();
match v.as_str() {
"" | "auto" | "on" | "1" | "true" => true,
"off" | "0" | "false" | "none" => false,
_ => {
eprintln!(
"docling-pdf: DOCLING_RS_OCR_ORIENTATION={raw:?} is not auto|off; using auto"
);
true
}
}
}
const PROBES: usize = 6;
const UPRIGHT_CONF: f32 = 0.90;
const UPRIGHT_CHARS: usize = 20;
const OVERTURN: f32 = 1.2;
const CONF_MARGIN: f32 = 0.10;
const MIN_CHARS: usize = 8;
const MIN_CONF: f32 = 0.55;
struct Score {
weighted: f32,
chars: usize,
}
impl Score {
fn mean_conf(&self) -> f32 {
if self.chars == 0 {
0.0
} else {
self.weighted / self.chars as f32
}
}
}
fn probe(img: &RgbImage, ocr: &mut OcrModel, boxes: &[DetBox]) -> Result<Score, String> {
let page = Region {
label: "text",
score: 1.0,
l: 0.0,
t: 0.0,
r: img.width() as f32,
b: img.height() as f32,
};
let (_, mut lines) = crate::timing::timed("orient.prep", || {
if boxes.is_empty() {
prep_region_lines(img, std::slice::from_ref(&page), 1.0)
} else {
prep_region_lines_det(img, std::slice::from_ref(&page), 1.0, boxes)
}
});
let mut order: Vec<usize> = (0..lines.len()).collect();
order.sort_by(|&a, &b| lines[b].w.cmp(&lines[a].w).then(a.cmp(&b)));
order.truncate(PROBES);
order.sort_unstable();
let probes: Vec<PrepLine> = order.iter().rev().map(|&i| lines.swap_remove(i)).collect();
let (weighted, chars) = crate::timing::timed("orient.score", || ocr.score_lines(&probes))?;
Ok(Score { weighted, chars })
}
pub(crate) fn detect(img: &RgbImage, ocr: &mut Recognizer, scale: f32, boxes: &[DetBox]) -> u16 {
match ocr {
Recognizer::PpOcr(model) => detect_by_probe(img, model, boxes),
Recognizer::Tesseract(t) => t.detect_orientation(img, scale).unwrap_or(0),
}
}
pub(crate) fn rotate_boxes(boxes: &[DetBox], w: f32, h: f32, deg: u16) -> Vec<DetBox> {
boxes
.iter()
.map(|d| {
let (l, t, r, b) = match deg {
90 => (d.t, w - d.r, d.b, w - d.l),
180 => (w - d.r, h - d.b, w - d.l, h - d.t),
270 => (h - d.b, d.l, h - d.t, d.r),
_ => (d.l, d.t, d.r, d.b),
};
DetBox {
l,
t,
r,
b,
score: d.score,
}
})
.collect()
}
fn detect_by_probe(img: &RgbImage, ocr: &mut OcrModel, boxes: &[DetBox]) -> u16 {
let fail = |e: String| {
debug_log!("docling-pdf: orientation probe failed ({e}); assuming upright");
0
};
let s0 = match probe(img, ocr, boxes) {
Ok(s) => s,
Err(e) => return fail(e),
};
if s0.chars >= UPRIGHT_CHARS && s0.mean_conf() >= UPRIGHT_CONF {
debug_log!(
"docling-pdf: orientation 0° reads {} chars at {:.2} — upright, no probes",
s0.chars,
s0.mean_conf()
);
return 0;
}
let hypotheses: [(u16, RgbImage); 3] = [
(90, rotate270(img)),
(180, rotate180(img)),
(270, rotate90(img)),
];
debug_log!(
"docling-pdf: orientation 0°: {} chars at {:.2} (weighted {:.1})",
s0.chars,
s0.mean_conf(),
s0.weighted
);
let (mut best_deg, mut best) = (
0u16,
Score {
weighted: 0.0,
chars: 0,
},
);
let (w, h) = (img.width() as f32, img.height() as f32);
for (deg, rotated) in &hypotheses {
let s = match probe(rotated, ocr, &rotate_boxes(boxes, w, h, *deg)) {
Ok(s) => s,
Err(e) => return fail(e),
};
debug_log!(
"docling-pdf: orientation {deg}°: {} chars at {:.2} (weighted {:.1})",
s.chars,
s.mean_conf(),
s.weighted
);
if s.weighted > best.weighted {
(best_deg, best) = (*deg, s);
}
}
if best_deg != 0
&& best.chars >= MIN_CHARS
&& best.mean_conf() >= MIN_CONF
&& best.weighted > OVERTURN * s0.weighted
&& best.mean_conf() >= s0.mean_conf() + CONF_MARGIN
{
return best_deg;
}
0
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn boxes_rotate_with_the_image() {
let bx = DetBox {
l: 10.0,
t: 5.0,
r: 50.0,
b: 15.0,
score: 0.9,
};
let at = |deg: u16| {
let r = rotate_boxes(std::slice::from_ref(&bx), 100.0, 60.0, deg);
(r[0].l, r[0].t, r[0].r, r[0].b)
};
assert_eq!(at(0), (10.0, 5.0, 50.0, 15.0));
assert_eq!(at(90), (5.0, 50.0, 15.0, 90.0));
assert_eq!(at(180), (50.0, 45.0, 90.0, 55.0));
assert_eq!(at(270), (45.0, 10.0, 55.0, 50.0));
let mut img = RgbImage::new(100, 60);
img.put_pixel(10, 5, image::Rgb([255, 0, 0]));
let r = rotate270(&img);
assert_eq!(
r.get_pixel(5, 89)[0],
255,
"rotate270 maps (10,5) → (5, 89)"
);
let r = rotate90(&img);
assert_eq!(
r.get_pixel(54, 10)[0],
255,
"rotate90 maps (10,5) → (54, 10)"
);
}
}