oar_ocr/pipeline/oarocr/

mod.rs

//! The main OCR pipeline implementation.
//!
//! This module provides the complete OCR pipeline that combines multiple
//! components to perform document orientation classification, text detection,
//! text recognition, and text line classification.
//!
//! # Orchestration Architecture
//!
//! The pipeline has been refactored to use a unified orchestration system that
//! eliminates code duplication and provides better maintainability. The key
//! improvements include:
//!
//! ## Unified Processing Logic
//!
//! Previously, methods like `process_single_image`, `process_single_image_from_memory`,
//! `process_images_individually`, and `process_images_from_memory` contained
//! duplicated orchestration logic for:
//!
//! - Parallel vs sequential processing decisions
//! - Index management and result sorting
//! - Pipeline stage execution patterns
//! - Error handling and logging
//!
//! Now, all these methods use the [`orchestration`] module's abstractions:
//!
//! - [`ImageProcessingOrchestrator`] for batch processing coordination
//! - [`PipelineExecutor`] for unified stage execution
//! - [`ProcessingStrategy`] for configurable parallelism
//! - [`ImageInputSource`] for input abstraction
//!
//! ## Benefits
//!
//! - **Reduced Duplication**: ~200 lines of duplicated code eliminated
//! - **Improved Maintainability**: Changes to orchestration logic in one place
//! - **Better Testability**: Each orchestration component can be tested independently
//! - **Type Safety**: Compile-time prevention of invalid configurations
//! - **Consistent Behavior**: All processing paths use the same underlying logic
//!
//! ## Backward Compatibility
//!
//! The public API remains unchanged. All existing code using [`OAROCR::predict`]
//! and related methods will continue to work without modification.
//!
//! [`orchestration`]: orchestration
//! [`ImageProcessingOrchestrator`]: ImageProcessingOrchestrator
//! [`PipelineExecutor`]: PipelineExecutor
//! [`ProcessingStrategy`]: ProcessingStrategy
//! [`ImageInputSource`]: ImageInputSource

mod builder;
mod components;
mod config;
mod extensible_integration;
mod image_processing;
mod orchestration;
mod result;

pub use builder::OAROCRBuilder;
pub use config::{OAROCRConfig, OnnxThreadingConfig, ParallelPolicy};
pub use extensible_integration::{ExtensibleOAROCR, ExtensibleOAROCRBuilder};
pub use image_processing::ImageProcessor;
pub use orchestration::{
    ImageInputSource, ImageProcessingOrchestrator, PipelineExecutor, PipelineStage,
    PipelineStageConfig, ProcessingStrategy,
};
pub use result::{ErrorMetrics, OAROCRResult, TextRegion};

use crate::core::{OCRError, parse_text_line_orientation, traits::StandardPredictor};
use crate::pipeline::PipelineStats;
use crate::pipeline::stages::{
    CroppingConfig, CroppingStageProcessor, RecognitionConfig, RecognitionStageProcessor,
};
use crate::predictor::{
    DocOrientationClassifier, DoctrRectifierPredictor, TextDetPredictor, TextLineClasPredictor,
    TextRecPredictor,
};
use components::ComponentBuilder;
use image::RgbImage;

use crate::pipeline::StatsManager;
use std::path::Path;
use std::sync::{Arc, Once};
use tracing::{debug, info, warn};

/// Global synchronization for rayon thread pool configuration.
/// This ensures the global thread pool is only configured once.
static THREAD_POOL_INIT: Once = Once::new();

// Parameter struct to reduce function argument count
#[derive(Debug, Clone)]
pub struct SingleImageProcessingParams<'a> {
    index: usize,
    input_img_arc: Arc<RgbImage>,
    current_img: RgbImage,
    text_boxes: Vec<crate::processors::BoundingBox>,
    orientation_angle: Option<f32>,
    rectified_img: Option<Arc<RgbImage>>,
    image_path: &'a Path,
}

/// Configures the global rayon thread pool if not already configured.
///
/// This function uses `std::sync::Once` to ensure the global thread pool
/// is only configured once, preventing race conditions and silent failures
/// that can occur when `build_global()` is called multiple times.
///
/// # Arguments
///
/// * `max_threads` - Maximum number of threads to use for the thread pool
///
/// # Returns
///
/// A Result indicating success or an OCRError if configuration fails
pub fn configure_thread_pool_once(max_threads: usize) -> crate::core::OcrResult<()> {
    let mut result = Ok(());

    THREAD_POOL_INIT.call_once(|| {
        debug!(
            "Configuring global rayon thread pool with {} threads",
            max_threads
        );
        if let Err(e) = rayon::ThreadPoolBuilder::new()
            .num_threads(max_threads)
            .build_global()
        {
            result = Err(OCRError::config_error(format!(
                "Failed to configure global thread pool: {e}"
            )));
        }
    });

    result
}

/// The main OCR pipeline that combines multiple components to perform
/// document processing and text recognition.
///
/// This struct manages the complete OCR pipeline, including document
/// orientation classification, text detection, text recognition, and
/// text line classification. It initializes and coordinates all the
/// required components based on the provided configuration.
pub struct OAROCR {
    /// Configuration for the OCR pipeline.
    config: OAROCRConfig,
    /// Document orientation classifier (optional).
    doc_orientation_classifier: Option<DocOrientationClassifier>,
    /// Document rectifier for unwarping (optional).
    doc_rectifier: Option<DoctrRectifierPredictor>,
    /// Text detector for finding text regions (required).
    text_detector: TextDetPredictor,
    /// Text line classifier for orientation (optional).
    text_line_classifier: Option<TextLineClasPredictor>,
    /// Text recognizer for reading text content (required).
    text_recognizer: TextRecPredictor,
    /// Statistics manager for the pipeline execution (thread-safe).
    stats: StatsManager,
}

impl OAROCR {
    /// Creates a new OAROCR instance with the provided configuration.
    ///
    /// This method initializes all the required components based on the
    /// configuration and builds the complete OCR pipeline.
    ///
    /// # Arguments
    ///
    /// * `config` - The configuration for the OCR pipeline
    ///
    /// # Returns
    ///
    /// A Result containing the OAROCR instance or an OCRError
    pub fn new(config: OAROCRConfig) -> crate::core::OcrResult<Self> {
        info!("Initializing OAROCR pipeline with config: {:?}", config);

        // Configure global rayon thread pool if max_threads is specified
        if let Some(max_threads) = config.max_threads() {
            configure_thread_pool_once(max_threads)?;
        }

        // Initialize optional components first
        let doc_orientation_classifier = if config.use_doc_orientation_classify {
            info!("Initializing document orientation classifier");
            Some(ComponentBuilder::build_doc_orientation_classifier(&config)?)
        } else {
            None
        };

        let doc_rectifier = if config.use_doc_unwarping {
            info!("Initializing document rectifier");
            Some(ComponentBuilder::build_doc_rectifier(&config)?)
        } else {
            None
        };

        let text_line_classifier = if config.use_textline_orientation {
            info!("Initializing text line classifier");
            Some(ComponentBuilder::build_text_line_classifier(&config)?)
        } else {
            None
        };

        // Initialize required components
        info!("Initializing text detector");
        let text_detector = ComponentBuilder::build_text_detector(&config)?;

        info!("Initializing text recognizer");
        let text_recognizer = ComponentBuilder::build_text_recognizer(&config)?;

        let pipeline = Self {
            config,
            doc_orientation_classifier,
            doc_rectifier,
            text_detector,
            text_line_classifier,
            text_recognizer,
            stats: StatsManager::new(),
        };

        info!("OAROCR pipeline initialized successfully");
        Ok(pipeline)
    }

    /// Processes one or more images through the OCR pipeline.
    ///
    /// This method runs the complete OCR pipeline on either a single image or
    /// a batch of images, including document orientation classification, text detection,
    /// text recognition, and text line classification (if enabled).
    ///
    /// Multiple images are processed in parallel using rayon for optimal performance.
    ///
    /// # Arguments
    ///
    /// * `image_paths` - A slice of paths to the image files
    ///
    /// # Returns
    ///
    /// A Result containing a vector of OAROCRResult or an OCRError
    /// Processes one or more images already loaded in memory.
    ///
    /// Prefer this API when you have RgbImage instances to avoid file I/O.
    pub fn predict(&self, images: &[RgbImage]) -> crate::core::OcrResult<Vec<OAROCRResult>> {
        let start_time = std::time::Instant::now();

        info!(
            "Starting OCR pipeline for {} in-memory image(s)",
            images.len()
        );

        let result = self.process_images_from_memory(images);

        // Update statistics based on the result
        let processing_time = start_time.elapsed();
        let total_time_ms = processing_time.as_millis() as f64;

        match &result {
            Ok(results) => {
                self.update_stats(images.len(), results.len(), 0, total_time_ms);
            }
            Err(_) => {
                self.update_stats(images.len(), 0, images.len(), total_time_ms);
            }
        }

        result
    }

    /// Internal: process in-memory images individually (with parallelism thresholds)
    fn process_images_from_memory(
        &self,
        images: &[RgbImage],
    ) -> crate::core::OcrResult<Vec<OAROCRResult>> {
        // Use the new orchestration abstraction
        let orchestrator = ImageProcessingOrchestrator::new(self);

        // Convert images to indexed inputs
        let inputs: Vec<(usize, &RgbImage)> = images.iter().enumerate().collect();

        // Use auto strategy based on image threshold
        let image_threshold = self.config.image_threshold();
        let strategy = ProcessingStrategy::Auto(image_threshold);
        let stage_config = PipelineStageConfig::default(); // Full pipeline

        orchestrator.process_batch(inputs, strategy, stage_config)
    }

    /// Processes a single image from the detection stage onwards.
    ///
    /// This method continues processing after detection has already been performed,
    /// used in dynamic batching scenarios where detection is batched separately.
    fn process_single_image_from_detection(
        &self,
        params: SingleImageProcessingParams,
    ) -> crate::core::OcrResult<OAROCRResult> {
        // Destructure parameters
        let SingleImageProcessingParams {
            index,
            input_img_arc,
            current_img,
            text_boxes,
            orientation_angle,
            rectified_img,
            image_path,
        } = params;

        // Stage 4: Text box cropping (can be parallelized)
        let cropping_config = CroppingConfig::default();

        let cropping_stage_result = CroppingStageProcessor::process_single(
            &current_img,
            &text_boxes,
            Some(&cropping_config),
        )?;

        let cropped_images = cropping_stage_result.data.cropped_images;
        let failed_crops = cropping_stage_result.data.failed_crops;

        // Continue with text line orientation and recognition as in the original method
        // (This is the same logic as in process_single_image from stage 5 onwards)

        // Stage 5: Text line orientation classification
        let mut text_line_orientations: Vec<Option<f32>> = Vec::new();
        let mut failed_orientations = 0;
        if self.config.use_textline_orientation && !text_boxes.is_empty() {
            if let Some(ref classifier) = self.text_line_classifier {
                let valid_images: Vec<RgbImage> = cropped_images
                    .iter()
                    .filter_map(|o| o.as_ref().cloned())
                    .collect();
                let valid_images_count = valid_images.len();
                if !valid_images.is_empty() {
                    match classifier.predict(valid_images, None) {
                        Ok(result) => {
                            let mut result_idx = 0usize;
                            for cropped_img_opt in &cropped_images {
                                if cropped_img_opt.is_some() {
                                    if let (Some(labels), Some(score_list)) = (
                                        result.label_names.get(result_idx),
                                        result.scores.get(result_idx),
                                    ) {
                                        if let (Some(label), Some(&score)) =
                                            (labels.first(), score_list.first())
                                        {
                                            let confidence_threshold = self
                                                .config
                                                .text_line_orientation_stage
                                                .as_ref()
                                                .and_then(|config| config.confidence_threshold);

                                            let orientation_result = parse_text_line_orientation(
                                                label.as_ref(),
                                                score,
                                                confidence_threshold,
                                            );

                                            if orientation_result.is_confident {
                                                text_line_orientations
                                                    .push(Some(orientation_result.angle));
                                            } else {
                                                text_line_orientations.push(None);
                                            }
                                        } else {
                                            text_line_orientations.push(None);
                                        }
                                    } else {
                                        text_line_orientations.push(None);
                                    }
                                    result_idx += 1;
                                } else {
                                    text_line_orientations.push(None);
                                }
                            }
                        }
                        Err(e) => {
                            failed_orientations = valid_images_count;
                            warn!(
                                "Text line orientation classification failed for {} images: {}",
                                valid_images_count, e
                            );
                            text_line_orientations.resize(text_boxes.len(), None);
                        }
                    }
                } else {
                    text_line_orientations.resize(text_boxes.len(), None);
                }
            } else {
                text_line_orientations.resize(text_boxes.len(), None);
            }
        } else {
            text_line_orientations.resize(text_boxes.len(), None);
        }

        // Stage 6: Text recognition (using existing logic)
        let recognition_config = RecognitionConfig::from_legacy_config(
            self.config.use_textline_orientation,
            self.config.aspect_ratio_bucketing.clone(),
        );

        let recognition_stage_result = RecognitionStageProcessor::process_single(
            cropped_images,
            Some(&text_line_orientations),
            Some(&self.text_recognizer),
            Some(&recognition_config),
        )?;

        let rec_texts = recognition_stage_result.data.rec_texts;
        let rec_scores = recognition_stage_result.data.rec_scores;
        let failed_recognitions = recognition_stage_result.data.failed_recognitions;

        // Stage 7: Final filtering and result assembly
        let score_thresh = self.config.recognition.score_thresh.unwrap_or(0.0);
        let mut final_texts: Vec<Option<Arc<str>>> = Vec::new();
        let mut final_scores: Vec<Option<f32>> = Vec::new();
        let mut final_orientations: Vec<Option<f32>> = Vec::new();
        for ((text, score), orientation) in rec_texts
            .into_iter()
            .zip(rec_scores)
            .zip(text_line_orientations.iter().cloned())
        {
            if score >= score_thresh {
                final_texts.push(Some(text));
                final_scores.push(Some(score));
                final_orientations.push(orientation);
            } else {
                final_texts.push(None);
                final_scores.push(None);
                final_orientations.push(orientation);
            }
        }

        // Create error metrics
        let error_metrics = ErrorMetrics {
            failed_crops,
            failed_recognitions,
            failed_orientations,
            total_text_boxes: text_boxes.len(),
        };

        // Create text regions from parallel vectors
        let text_regions = OAROCRResult::create_text_regions_from_vectors(
            &text_boxes,
            &final_texts,
            &final_scores,
            &final_orientations,
        );

        Ok(OAROCRResult {
            input_path: Arc::from(image_path.to_string_lossy().as_ref()),
            index,
            input_img: input_img_arc,
            text_regions,
            orientation_angle,
            rectified_img,
            error_metrics,
        })
    }

    /// Gets the pipeline statistics.
    ///
    /// # Returns
    ///
    /// A copy of the current PipelineStats
    pub fn get_stats(&self) -> PipelineStats {
        self.stats.get_stats()
    }

    /// Updates the pipeline statistics after processing images.
    ///
    /// # Arguments
    ///
    /// * `processed_count` - Number of images processed
    /// * `successful_count` - Number of successful predictions
    /// * `failed_count` - Number of failed predictions
    /// * `inference_time_ms` - Total inference time in milliseconds
    fn update_stats(
        &self,
        processed_count: usize,
        successful_count: usize,
        failed_count: usize,
        inference_time_ms: f64,
    ) {
        self.stats.update_stats(
            processed_count,
            successful_count,
            failed_count,
            inference_time_ms,
        );
    }

    /// Resets the pipeline statistics.
    pub fn reset_stats(&self) {
        self.stats.reset_stats();
    }

    /// Gets the pipeline configuration.
    ///
    /// # Returns
    ///
    /// A reference to the OAROCRConfig
    pub fn get_config(&self) -> &OAROCRConfig {
        &self.config
    }

    // Private helper methods for reducing complexity in main processing methods
    //
    // The following helper functions were extracted to address the long function problem
    // identified in the codebase audit. They break down complex operations into focused,
    // single-responsibility functions with clear contracts.
    //
    // Benefits achieved:
    // - Reduced cognitive load: Each function has a single, clear purpose
    // - Improved testability: Each helper can be tested in isolation
    // - Better error handling: Focused error contexts and clear propagation
    // - Enhanced maintainability: Changes to specific logic are localized
    // - Easier auditing: Small functions are easier to review and verify
    //
    // The main method that benefited most from this refactoring:
    // - process_images_with_cross_recognition_batching: 301 lines → 28 lines (90% reduction)
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_oarocr_builder_text_rec_score_thresh() {
        // Test that the text_rec_score_thresh method properly sets the score threshold
        let builder = OAROCRBuilder::new(
            "dummy_det_path".to_string(),
            "dummy_rec_path".to_string(),
            "dummy_dict_path".to_string(),
        )
        .text_rec_score_threshold(0.8);

        assert_eq!(builder.get_config().recognition.score_thresh, Some(0.8));
    }

    #[test]
    fn test_orchestration_abstraction_imports() {
        // Test that the orchestration abstractions are properly exported
        use crate::pipeline::oarocr::{
            ImageInputSource, ImageProcessingOrchestrator, PipelineExecutor, PipelineStage,
            PipelineStageConfig, ProcessingStrategy,
        };
        use image::RgbImage;
        use std::path::Path;
        use std::sync::Arc;

        // Test that we can create instances of the orchestration types
        let strategy = ProcessingStrategy::Sequential;
        let config = PipelineStageConfig::default();

        // Test enum variants
        let _orientation_stage = PipelineStage::Orientation;
        let _detection_stage = PipelineStage::Detection;

        // Test ImageInputSource variants
        let path = Path::new("test.jpg");
        let _path_source = ImageInputSource::Path(path);

        let img = RgbImage::new(100, 100);
        let _memory_source = ImageInputSource::Memory(&img);

        let img_arc = Arc::new(img);
        let _loaded_source = ImageInputSource::LoadedWithPath(img_arc, path);

        // Test that the types have the expected properties
        assert!(!strategy.should_use_parallel(10));
        assert_eq!(config.start_from, PipelineStage::Orientation);

        // Test that we can reference the orchestrator and executor types
        // (We can't actually create them without a valid OAROCR instance)
        let _orchestrator_type = std::any::type_name::<ImageProcessingOrchestrator>();
        let _executor_type = std::any::type_name::<PipelineExecutor>();
    }

    #[test]
    fn test_processing_strategy_behavior() {
        // Test ProcessingStrategy behavior
        let sequential = ProcessingStrategy::Sequential;
        let parallel = ProcessingStrategy::Parallel;
        let auto_5 = ProcessingStrategy::Auto(5);

        // Sequential should never use parallel
        assert!(!sequential.should_use_parallel(1));
        assert!(!sequential.should_use_parallel(100));

        // Parallel should always use parallel
        assert!(parallel.should_use_parallel(1));
        assert!(parallel.should_use_parallel(100));

        // Auto should use threshold
        assert!(!auto_5.should_use_parallel(3));
        assert!(!auto_5.should_use_parallel(5));
        assert!(auto_5.should_use_parallel(6));
        assert!(auto_5.should_use_parallel(10));
    }

    #[test]
    fn test_pipeline_stage_config_customization() {
        let mut config = PipelineStageConfig::default();

        // Test default values
        assert_eq!(config.start_from, PipelineStage::Orientation);
        assert!(config.skip_stages.is_empty());
        assert!(config.custom_params.is_none());

        // Test customization
        config.start_from = PipelineStage::Detection;
        config.skip_stages.insert(PipelineStage::Recognition);

        assert_eq!(config.start_from, PipelineStage::Detection);
        assert!(config.skip_stages.contains(&PipelineStage::Recognition));
        assert!(!config.skip_stages.contains(&PipelineStage::Orientation));
    }

    #[test]
    fn test_oarocr_builder_doc_orientation_confidence_threshold() {
        // Test that the doc_orientation_confidence_threshold method properly sets the threshold
        let builder = OAROCRBuilder::new(
            "dummy_det_path".to_string(),
            "dummy_rec_path".to_string(),
            "dummy_dict_path".to_string(),
        )
        .doc_orientation_threshold(0.8);

        assert!(builder.get_config().orientation_stage.is_some());
        assert_eq!(
            builder
                .get_config()
                .orientation_stage
                .as_ref()
                .unwrap()
                .confidence_threshold,
            Some(0.8)
        );
    }

    #[test]
    fn test_oarocr_builder_textline_orientation_confidence_threshold() {
        // Test that the textline_orientation_confidence_threshold method properly sets the threshold
        let builder = OAROCRBuilder::new(
            "dummy_det_path".to_string(),
            "dummy_rec_path".to_string(),
            "dummy_dict_path".to_string(),
        )
        .textline_orientation_threshold(0.9);

        assert!(builder.get_config().text_line_orientation_stage.is_some());
        assert_eq!(
            builder
                .get_config()
                .text_line_orientation_stage
                .as_ref()
                .unwrap()
                .confidence_threshold,
            Some(0.9)
        );
    }

    #[test]
    fn test_oarocr_result_alignment_preservation() {
        // Test that OAROCRResult maintains 1:1 correspondence between text_boxes and recognition results
        use crate::processors::BoundingBox;
        use crate::processors::Point;
        use image::RgbImage;
        use std::sync::Arc;

        // Create mock data
        let text_boxes = vec![
            BoundingBox {
                points: vec![
                    Point { x: 0.0, y: 0.0 },
                    Point { x: 10.0, y: 0.0 },
                    Point { x: 10.0, y: 10.0 },
                    Point { x: 0.0, y: 10.0 },
                ],
            },
            BoundingBox {
                points: vec![
                    Point { x: 20.0, y: 0.0 },
                    Point { x: 30.0, y: 0.0 },
                    Point { x: 30.0, y: 10.0 },
                    Point { x: 20.0, y: 10.0 },
                ],
            },
            BoundingBox {
                points: vec![
                    Point { x: 40.0, y: 0.0 },
                    Point { x: 50.0, y: 0.0 },
                    Point { x: 50.0, y: 10.0 },
                    Point { x: 40.0, y: 10.0 },
                ],
            },
        ];

        // Create recognition results where the middle one is filtered out (None)
        let rec_texts = vec![
            Some(Arc::from("Hello")),
            None, // This was filtered out due to low confidence
            Some(Arc::from("World")),
        ];

        let rec_scores = vec![
            Some(0.9),
            None, // This was filtered out due to low confidence
            Some(0.8),
        ];

        // Create text regions from parallel vectors
        let text_regions = OAROCRResult::create_text_regions_from_vectors(
            &text_boxes,
            &rec_texts,
            &rec_scores,
            &[None, None, None],
        );

        let result = OAROCRResult {
            input_path: Arc::from("test.jpg"),
            index: 0,
            input_img: Arc::new(RgbImage::new(100, 100)),
            text_regions,
            orientation_angle: None,
            rectified_img: None,
            error_metrics: ErrorMetrics::default(),
        };

        // Verify that the text regions were created correctly
        assert_eq!(result.text_regions.len(), 3);

        // Verify that we can access text regions with their recognition results
        for (i, region) in result.text_regions.iter().enumerate() {
            // Each region should have a bounding box
            assert!(region.bounding_box.points.len() >= 4);

            match i {
                0 => {
                    // First region should have recognition result
                    assert!(region.text.is_some());
                    assert!(region.confidence.is_some());
                    assert_eq!(region.text.as_ref().unwrap().as_ref(), "Hello");
                    assert_eq!(region.confidence.unwrap(), 0.9);
                }
                1 => {
                    // Second region should have no recognition result (filtered out)
                    assert!(region.text.is_none());
                    assert!(region.confidence.is_none());
                }
                2 => {
                    // Third region should have recognition result
                    assert!(region.text.is_some());
                    assert!(region.confidence.is_some());
                    assert_eq!(region.text.as_ref().unwrap().as_ref(), "World");
                    assert_eq!(region.confidence.unwrap(), 0.8);
                }
                _ => panic!("Unexpected index"),
            }
        }
    }

    #[test]
    fn test_thread_pool_configuration_once() {
        // Test that thread pool configuration can be called multiple times without error
        // This tests the fix for the global rayon thread pool race condition

        // First call should succeed
        let result1 = configure_thread_pool_once(2);
        assert!(
            result1.is_ok(),
            "First thread pool configuration should succeed"
        );

        // Second call should also succeed (should be ignored due to Once)
        let result2 = configure_thread_pool_once(4);
        assert!(
            result2.is_ok(),
            "Second thread pool configuration should succeed (ignored)"
        );

        // Third call with different thread count should also succeed
        let result3 = configure_thread_pool_once(1);
        assert!(
            result3.is_ok(),
            "Third thread pool configuration should succeed (ignored)"
        );
    }
}