df-ocr-switcher 0.1.0

//! Engine Tesseract 5.5 via tesseract5-rs (feature: tesseract-engine).
//!
//! Usa `Ocr5Engine` per text recognition, `LayoutAnalyzer` di ppocr-rs
//! per layout analysis, `DocOrientationClassifier` di ppocr-rs per rotazione.
//! Le word-bbox vengono da `TesseractHierarchy::iter_words()`.

use std::path::PathBuf;

use image::{imageops, RgbImage};
use ppocr_rs::{
    DocOrientation, DocOrientationClassifier, LayoutAnalyzer, LayoutBox,
    ModelHub, PpStructureModel,
};
use tesseract5_rs::{BoundingBox, Ocr5Engine, OcrOptions as TessOpts};

use crate::engine::{OcrBlock, OcrLayoutBox, OcrPageResult, OcrWord};
use crate::error::{Error, Result};

use super::OcrEngine;

// ─── Config ──────────────────────────────────────────────────────────────────

pub struct TesseractEngineConfig {
    /// Lingue Tesseract (es. `"ita+eng"`).
    pub lang:      String,
    /// Path tessdata dir. Se `None` usa `default_tessdata_dir()`.
    pub tessdata:  Option<PathBuf>,
    /// Path modello orientamento. Se `None` usa ModelHub.
    pub ori_model: Option<PathBuf>,
    /// Page Segmentation Mode (default `None` = auto).
    pub psm:       Option<u8>,
}

impl Default for TesseractEngineConfig {
    fn default() -> Self {
        Self {
            lang:      "ita+eng".into(),
            tessdata:  None,
            ori_model: None,
            psm:       None,
        }
    }
}

// ─── Engine ──────────────────────────────────────────────────────────────────

pub struct TesseractEngine {
    engine: Ocr5Engine,
    ori:    Option<DocOrientationClassifier>,
}

impl TesseractEngine {
    pub fn new(cfg: TesseractEngineConfig) -> Result<Self> {
        let hub = ModelHub::with_default_cache()?;

        let engine = Ocr5Engine::new(TessOpts {
            lang:           cfg.lang,
            psm:            cfg.psm,
            tessdata_dir:   cfg.tessdata,
            with_hierarchy: true,
        })?;

        let ori = {
            let ori_path = if let Some(p) = cfg.ori_model {
                p
            } else {
                match hub.ensure_single(PpStructureModel::DocOrientation) {
                    Ok(p) => p.onnx,
                    Err(e) => {
                        eprintln!("[df-ocr-switcher/tesseract] orientamento disabilitato: {e}");
                        return Ok(Self { engine, ori: None });
                    }
                }
            };
            match DocOrientationClassifier::from_path(&ori_path) {
                Ok(clf) => Some(clf),
                Err(e)  => {
                    eprintln!("[df-ocr-switcher/tesseract] orientamento disabilitato: {e}");
                    None
                }
            }
        };

        Ok(Self { engine, ori })
    }
}

impl OcrEngine for TesseractEngine {
    fn process_page(&mut self, img: &RgbImage, layout: &mut LayoutAnalyzer) -> Result<OcrPageResult> {
        // ── Step 1: orientamento (ppocr-rs DocOrientationClassifier) ──────
        let (page_angle, upright) = if let Some(clf) = &self.ori {
            let (orient, _) = clf.classify(img)?;
            let deg = orient.degrees();
            let rot = rotate_upright(img, orient);
            (deg, rot)
        } else {
            (0u32, img.clone())
        };
        let (w, h) = upright.dimensions();

        // ── Step 2: layout analysis (ppocr-rs PP-DocLayoutV3) ─────────────
        let layout_boxes_raw = layout.analyze(&upright)?;
        let layout_boxes: Vec<OcrLayoutBox> = layout_boxes_raw.iter().map(lb_to_ocr).collect();

        // ── Step 3: Tesseract OCR → TesseractHierarchy ────────────────────
        let raw = upright.as_raw();
        let out = self.engine.recognize(
            raw,
            w as i32, h as i32,
            3, (w * 3) as i32,
        )?;

        // ── Step 4: assembly blocks + words ───────────────────────────────
        let hierarchy = out.hierarchy.unwrap_or_default();
        let mut blocks = Vec::new();
        let mut words  = Vec::new();

        for line in hierarchy.iter_lines() {
            let (x1, y1, x2, y2) = bbox_to_aabb(&line.bbox);
            let cx = (x1 + x2) / 2;
            let cy = (y1 + y2) / 2;
            let layout_idx = assign_layout_idx(cx, cy, &layout_boxes_raw);

            blocks.push(OcrBlock {
                text:       line.text.clone(),
                x1, y1, x2, y2,
                confidence: line.confidence,
                layout_idx,
            });

            for w_item in &line.words {
                let (wx1, wy1, wx2, wy2) = bbox_to_aabb(&w_item.bbox);
                words.push(OcrWord {
                    text:       w_item.text.clone(),
                    x1: wx1, y1: wy1, x2: wx2, y2: wy2,
                    confidence: w_item.confidence,
                    layout_idx,
                });
            }
        }

        Ok(OcrPageResult {
            page_angle, page_width: w, page_height: h,
            layout_boxes, blocks, words,
            tables: vec![], formulas: vec![],
        })
    }
}

// ─── Helper ──────────────────────────────────────────────────────────────────

fn rotate_upright(img: &RgbImage, orient: DocOrientation) -> RgbImage {
    match orient {
        DocOrientation::Deg0   => img.clone(),
        DocOrientation::Deg90  => imageops::rotate270(img),
        DocOrientation::Deg180 => imageops::rotate180(img),
        DocOrientation::Deg270 => imageops::rotate90(img),
    }
}

fn bbox_to_aabb(b: &BoundingBox) -> (u32, u32, u32, u32) {
    (
        b.left.max(0) as u32,
        b.top.max(0) as u32,
        b.right.max(0) as u32,
        b.bottom.max(0) as u32,
    )
}

fn lb_to_ocr(lb: &LayoutBox) -> OcrLayoutBox {
    OcrLayoutBox {
        class_name:    format!("{:?}", lb.class),
        semantic:      lb.class.semantic(),
        x1: lb.xmin(), y1: lb.ymin(),
        x2: lb.xmax(), y2: lb.ymax(),
        reading_order: lb.reading_order,
    }
}

/// Assegna l'indice layout al centroide del testo.
/// Cerca prima containment esatto, poi fallback al layout box più vicino.
fn assign_layout_idx(cx: u32, cy: u32, layout_boxes: &[LayoutBox]) -> i32 {
    // containment
    if let Some(i) = layout_boxes.iter().position(|lb| lb.contains(cx, cy)) {
        return i as i32;
    }
    // nearest
    layout_boxes.iter()
        .enumerate()
        .min_by(|(_, a), (_, b)| {
            a.distance_to(cx, cy).partial_cmp(&b.distance_to(cx, cy))
                .unwrap_or(std::cmp::Ordering::Equal)
        })
        .map(|(i, _)| i as i32)
        .unwrap_or(-1)
}