kreuzberg 4.0.8 → 4.1.0

data/vendor/kreuzberg/src/pdf/{hierarchy.rs → hierarchy/extraction.rs}

@@ -1,173 +1,21 @@
-//! PDF text hierarchy extraction using pdfium character positions.
+//! PDF text hierarchy extraction and text block analysis.
 //!
 //! This module provides functions for extracting character information from PDFs,
-//! preserving font size and position data for text hierarchy analysis.
-//!
-//! Note: Requires the "pdf" feature to be enabled.
+//! merging characters into text blocks, and assigning hierarchy levels based on
+//! font size analysis.
 
-use super::error::{PdfError, Result};
+use super::bounding_box::BoundingBox;
+use super::clustering::FontSizeCluster;
 use crate::core::config::ExtractionConfig;
+use crate::pdf::error::{PdfError, Result};
 use pdfium_render::prelude::*;
 
 // Magic number constants
 const DEFAULT_FONT_SIZE: f32 = 12.0;
-const WEIGHTED_DISTANCE_X_WEIGHT: f32 = 5.0;
-const WEIGHTED_DISTANCE_Y_WEIGHT: f32 = 1.0;
-const KMEANS_MAX_ITERATIONS: usize = 100;
-const KMEANS_CONVERGENCE_THRESHOLD: f32 = 0.01;
 const MERGE_INTERSECTION_THRESHOLD: f32 = 0.05;
 const MERGE_X_THRESHOLD_MULTIPLIER: f32 = 2.0;
 const MERGE_Y_THRESHOLD_MULTIPLIER: f32 = 1.5;
 
-/// A bounding box for text or elements.
-#[derive(Debug, Clone, Copy, PartialEq)]
-pub struct BoundingBox {
-    /// Left x-coordinate
-    pub left: f32,
-    /// Top y-coordinate
-    pub top: f32,
-    /// Right x-coordinate
-    pub right: f32,
-    /// Bottom y-coordinate
-    pub bottom: f32,
-}
-
-impl BoundingBox {
-    /// Calculate the Intersection over Union (IOU) between this bounding box and another.
-    ///
-    /// IOU = intersection_area / union_area
-    ///
-    /// # Arguments
-    ///
-    /// * `other` - The other bounding box to compare with
-    ///
-    /// # Returns
-    ///
-    /// The IOU value between 0.0 and 1.0
-    pub fn iou(&self, other: &BoundingBox) -> f32 {
-        let intersection_area = self.calculate_intersection_area(other);
-        let self_area = self.calculate_area();
-        let other_area = other.calculate_area();
-        let union_area = self_area + other_area - intersection_area;
-
-        if union_area <= 0.0 {
-            0.0
-        } else {
-            intersection_area / union_area
-        }
-    }
-
-    /// Calculate the weighted distance between the centers of two bounding boxes.
-    ///
-    /// The distance is weighted with X-axis having weight 5.0 and Y-axis having weight 1.0.
-    /// This reflects the greater importance of horizontal distance in text layout.
-    ///
-    /// # Arguments
-    ///
-    /// * `other` - The other bounding box to compare with
-    ///
-    /// # Returns
-    ///
-    /// The weighted distance value
-    pub fn weighted_distance(&self, other: &BoundingBox) -> f32 {
-        let (self_center_x, self_center_y) = self.center();
-        let (other_center_x, other_center_y) = other.center();
-
-        let dx = (self_center_x - other_center_x).abs();
-        let dy = (self_center_y - other_center_y).abs();
-
-        dx * WEIGHTED_DISTANCE_X_WEIGHT + dy * WEIGHTED_DISTANCE_Y_WEIGHT
-    }
-
-    /// Calculate the intersection ratio relative to this bounding box's area.
-    ///
-    /// intersection_ratio = intersection_area / self_area
-    ///
-    /// # Arguments
-    ///
-    /// * `other` - The other bounding box to compare with
-    ///
-    /// # Returns
-    ///
-    /// The intersection ratio between 0.0 and 1.0
-    pub fn intersection_ratio(&self, other: &BoundingBox) -> f32 {
-        let intersection_area = self.calculate_intersection_area(other);
-        let self_area = self.calculate_area();
-
-        if self_area <= 0.0 {
-            0.0
-        } else {
-            intersection_area / self_area
-        }
-    }
-
-    /// Check if this bounding box contains another bounding box.
-    pub fn contains(&self, other: &BoundingBox) -> bool {
-        other.left >= self.left && other.right <= self.right && other.top >= self.top && other.bottom <= self.bottom
-    }
-
-    /// Calculate the center coordinates of this bounding box.
-    pub fn center(&self) -> (f32, f32) {
-        ((self.left + self.right) / 2.0, (self.top + self.bottom) / 2.0)
-    }
-
-    /// Merge this bounding box with another, creating a box that contains both.
-    pub fn merge(&self, other: &BoundingBox) -> BoundingBox {
-        BoundingBox {
-            left: self.left.min(other.left),
-            top: self.top.min(other.top),
-            right: self.right.max(other.right),
-            bottom: self.bottom.max(other.bottom),
-        }
-    }
-
-    /// Calculate a relaxed IOU with an expansion factor.
-    pub fn relaxed_iou(&self, other: &BoundingBox, relaxation: f32) -> f32 {
-        let self_width = self.right - self.left;
-        let self_height = self.bottom - self.top;
-        let self_expansion = relaxation * self_width.min(self_height).max(0.0);
-
-        let other_width = other.right - other.left;
-        let other_height = other.bottom - other.top;
-        let other_expansion = relaxation * other_width.min(other_height).max(0.0);
-
-        let expanded_self = BoundingBox {
-            left: (self.left - self_expansion).max(0.0),
-            top: (self.top - self_expansion).max(0.0),
-            right: self.right + self_expansion,
-            bottom: self.bottom + self_expansion,
-        };
-
-        let expanded_other = BoundingBox {
-            left: (other.left - other_expansion).max(0.0),
-            top: (other.top - other_expansion).max(0.0),
-            right: other.right + other_expansion,
-            bottom: other.bottom + other_expansion,
-        };
-
-        expanded_self.iou(&expanded_other)
-    }
-
-    /// Calculate the area of this bounding box.
-    fn calculate_area(&self) -> f32 {
-        let width = (self.right - self.left).max(0.0);
-        let height = (self.bottom - self.top).max(0.0);
-        width * height
-    }
-
-    /// Calculate the intersection area between this bounding box and another.
-    fn calculate_intersection_area(&self, other: &BoundingBox) -> f32 {
-        let left = self.left.max(other.left);
-        let top = self.top.max(other.top);
-        let right = self.right.min(other.right);
-        let bottom = self.bottom.min(other.bottom);
-
-        let width = (right - left).max(0.0);
-        let height = (bottom - top).max(0.0);
-        width * height
-    }
-}
-
 /// Character information extracted from PDF with font metrics.
 #[derive(Debug, Clone)]
 pub struct CharData {
@@ -196,15 +44,6 @@ pub struct TextBlock {
     pub font_size: f32,
 }
 
-/// A cluster of text blocks with the same font size characteristics.
-#[derive(Debug, Clone)]
-pub struct FontSizeCluster {
-    /// The centroid (mean) font size of this cluster
-    pub centroid: f32,
-    /// The text blocks that belong to this cluster
-    pub members: Vec<TextBlock>,
-}
-
 /// Result of KMeans clustering on font sizes.
 ///
 /// Contains cluster labels for each block, where cluster index indicates
@@ -401,185 +240,6 @@ pub fn assign_hierarchy_levels_from_clusters(
     result
 }
 
-/// Cluster text blocks by font size using k-means algorithm.
-///
-/// Uses k-means clustering to group text blocks by their font size, which helps
-/// identify document hierarchy levels (H1, H2, Body, etc.). The algorithm:
-/// 1. Extracts font sizes from text blocks
-/// 2. Applies k-means clustering to group similar font sizes
-/// 3. Sorts clusters by centroid size in descending order (largest = H1)
-/// 4. Returns clusters with their member blocks
-///
-/// # Arguments
-///
-/// * `blocks` - Slice of TextBlock objects to cluster
-/// * `k` - Number of clusters to create
-///
-/// # Returns
-///
-/// Result with vector of FontSizeCluster ordered by size (descending),
-/// or an error if clustering fails
-///
-/// # Example
-///
-/// ```rust,no_run
-/// # #[cfg(feature = "pdf")]
-/// # {
-/// use kreuzberg::pdf::hierarchy::{TextBlock, BoundingBox, cluster_font_sizes};
-///
-/// let blocks = vec![
-///     TextBlock {
-///         text: "Title".to_string(),
-///         bbox: BoundingBox { left: 0.0, top: 0.0, right: 100.0, bottom: 24.0 },
-///         font_size: 24.0,
-///     },
-///     TextBlock {
-///         text: "Body".to_string(),
-///         bbox: BoundingBox { left: 0.0, top: 30.0, right: 100.0, bottom: 42.0 },
-///         font_size: 12.0,
-///     },
-/// ];
-///
-/// let clusters = cluster_font_sizes(&blocks, 2).unwrap();
-/// assert_eq!(clusters.len(), 2);
-/// assert_eq!(clusters[0].centroid, 24.0); // Largest is first
-/// # }
-/// ```
-/// Helper function to assign blocks to their nearest centroid.
-///
-/// Iterates through blocks and finds the closest centroid for each block,
-/// grouping them into clusters. Used in k-means clustering iterations.
-///
-/// # Arguments
-///
-/// * `blocks` - Slice of TextBlock objects to assign
-/// * `centroids` - Slice of centroid values (one per cluster)
-///
-/// # Returns
-///
-/// A vector of clusters, where each cluster contains the TextBlock objects
-/// assigned to that centroid
-fn assign_blocks_to_centroids(blocks: &[TextBlock], centroids: &[f32]) -> Vec<Vec<TextBlock>> {
-    let mut clusters: Vec<Vec<TextBlock>> = vec![Vec::new(); centroids.len()];
-
-    for block in blocks {
-        let mut min_distance = f32::INFINITY;
-        let mut best_cluster = 0;
-
-        for (i, &centroid) in centroids.iter().enumerate() {
-            let distance = (block.font_size - centroid).abs();
-            if distance < min_distance {
-                min_distance = distance;
-                best_cluster = i;
-            }
-        }
-
-        clusters[best_cluster].push(block.clone());
-    }
-
-    clusters
-}
-
-pub fn cluster_font_sizes(blocks: &[TextBlock], k: usize) -> Result<Vec<FontSizeCluster>> {
-    if blocks.is_empty() {
-        return Ok(Vec::new());
-    }
-
-    if k == 0 {
-        return Err(PdfError::TextExtractionFailed("K must be greater than 0".to_string()));
-    }
-
-    let actual_k = k.min(blocks.len());
-
-    // Extract unique font sizes for initialization
-    let mut font_sizes: Vec<f32> = blocks.iter().map(|b| b.font_size).collect();
-    font_sizes.sort_by(|a, b| b.partial_cmp(a).expect("Failed to compare font sizes during sorting")); // Sort descending
-    font_sizes.dedup(); // Remove duplicates to get unique font sizes
-
-    // Initialize centroids using actual font sizes from the data
-    // This is more robust than dividing the range uniformly
-    let mut centroids: Vec<f32> = Vec::new();
-
-    if font_sizes.len() >= actual_k {
-        // If we have at least k unique font sizes, pick them evenly spaced
-        let step = font_sizes.len() / actual_k;
-        for i in 0..actual_k {
-            let idx = i * step;
-            centroids.push(font_sizes[idx.min(font_sizes.len() - 1)]);
-        }
-    } else {
-        // If we have fewer unique sizes than k, use all of them and fill with interpolated values
-        centroids = font_sizes.clone();
-
-        // Add interpolated centroids between existing ones to reach desired k
-        let min_font = font_sizes[font_sizes.len() - 1];
-        let max_font = font_sizes[0];
-        let range = max_font - min_font;
-
-        while centroids.len() < actual_k {
-            let t = centroids.len() as f32 / (actual_k - 1) as f32;
-            let interpolated = max_font - t * range;
-            centroids.push(interpolated);
-        }
-
-        centroids.sort_by(|a, b| b.partial_cmp(a).expect("Failed to compare centroids during sorting"));
-        // Keep sorted descending
-    }
-
-    // Run k-means clustering for a fixed number of iterations
-    for _ in 0..KMEANS_MAX_ITERATIONS {
-        // Assign blocks to nearest centroid
-        let clusters = assign_blocks_to_centroids(blocks, &centroids);
-
-        // Update centroids
-        let mut new_centroids = Vec::with_capacity(actual_k);
-        for (i, cluster) in clusters.iter().enumerate() {
-            if !cluster.is_empty() {
-                new_centroids.push(cluster.iter().map(|b| b.font_size).sum::<f32>() / cluster.len() as f32);
-            } else {
-                new_centroids.push(centroids[i]);
-            }
-        }
-
-        // Check for convergence
-        let converged = centroids
-            .iter()
-            .zip(new_centroids.iter())
-            .all(|(old, new)| (old - new).abs() < KMEANS_CONVERGENCE_THRESHOLD);
-
-        std::mem::swap(&mut centroids, &mut new_centroids);
-
-        if converged {
-            break;
-        }
-    }
-
-    // Final assignment to create result
-    let clusters = assign_blocks_to_centroids(blocks, &centroids);
-
-    // Create FontSizeCluster objects with centroids
-    let mut result: Vec<FontSizeCluster> = Vec::new();
-
-    for i in 0..actual_k {
-        if !clusters[i].is_empty() {
-            let centroid_value = centroids[i];
-            result.push(FontSizeCluster {
-                centroid: centroid_value,
-                members: clusters[i].clone(),
-            });
-        }
-    }
-
-    // Sort by centroid size in descending order (largest font = H1)
-    result.sort_by(|a, b| {
-        b.centroid
-            .partial_cmp(&a.centroid)
-            .expect("Failed to compare centroids during final sort")
-    });
-
-    Ok(result)
-}
-
 /// Extract characters with fonts from a PDF page.
 ///
 /// Iterates through all characters on a page, extracting text, position,

data/vendor/kreuzberg/src/pdf/hierarchy/mod.rs

@@ -0,0 +1,18 @@
+//! PDF text hierarchy extraction using pdfium character positions.
+//!
+//! This module provides functions for extracting character information from PDFs,
+//! preserving font size and position data for text hierarchy analysis.
+//!
+//! Note: Requires the "pdf" feature to be enabled.
+
+mod bounding_box;
+mod clustering;
+mod extraction;
+
+// Re-export all public types and functions for backward compatibility
+pub use bounding_box::BoundingBox;
+pub use clustering::{FontSizeCluster, cluster_font_sizes};
+pub use extraction::{
+    CharData, HierarchyBlock, HierarchyLevel, KMeansResult, TextBlock, assign_hierarchy_levels,
+    assign_hierarchy_levels_from_clusters, extract_chars_with_fonts, merge_chars_into_blocks, should_trigger_ocr,
+};