xgen-doc2chunk 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

xgen_doc2chunk/core/processor/pdf_helpers/pdf_image_processor.py

@@ -0,0 +1,321 @@
+# xgen_doc2chunk/core/processor/pdf_helpers/pdf_image_processor.py
+"""
+PDF Image Processor
+
+Provides PDF-specific image processing that inherits from ImageProcessor.
+Handles XRef images, inline images, and page rendering for complex regions.
+
+This class consolidates all PDF image extraction logic including:
+- XRef-based image extraction
+- Page region rendering
+- Image filtering by size/position
+"""
+import logging
+from typing import Any, Dict, List, Optional, Set, Tuple, TYPE_CHECKING
+
+from xgen_doc2chunk.core.functions.img_processor import ImageProcessor
+from xgen_doc2chunk.core.functions.storage_backend import BaseStorageBackend
+
+if TYPE_CHECKING:
+    import fitz
+
+logger = logging.getLogger("xgen_doc2chunk.image_processor.pdf")
+
+
+class PDFImageProcessor(ImageProcessor):
+    """
+    PDF-specific image processor.
+    
+    Inherits from ImageProcessor and provides PDF-specific processing.
+    
+    Handles:
+    - XRef images (embedded images with XRef references)
+    - Inline images
+    - Page region rendering for complex areas
+    - Image extraction from PyMuPDF objects
+    
+    Example:
+        processor = PDFImageProcessor()
+        
+        # Process XRef image
+        tag = processor.process_image(image_data, xref=123)
+        
+        # Process page region
+        tag = processor.process_page_region(page, rect)
+    """
+    
+    def __init__(
+        self,
+        directory_path: str = "temp/images",
+        tag_prefix: str = "[Image:",
+        tag_suffix: str = "]",
+        storage_backend: Optional[BaseStorageBackend] = None,
+        dpi: int = 150,
+    ):
+        """
+        Initialize PDFImageProcessor.
+        
+        Args:
+            directory_path: Image save directory
+            tag_prefix: Tag prefix for image references
+            tag_suffix: Tag suffix for image references
+            storage_backend: Storage backend for saving images
+            dpi: DPI for page rendering
+        """
+        super().__init__(
+            directory_path=directory_path,
+            tag_prefix=tag_prefix,
+            tag_suffix=tag_suffix,
+            storage_backend=storage_backend,
+        )
+        self._dpi = dpi
+    
+    @property
+    def dpi(self) -> int:
+        """DPI for page rendering."""
+        return self._dpi
+    
+    @dpi.setter
+    def dpi(self, value: int) -> None:
+        """Set DPI for page rendering."""
+        self._dpi = value
+    
+    def process_image(
+        self,
+        image_data: bytes,
+        xref: Optional[int] = None,
+        page_num: Optional[int] = None,
+        **kwargs
+    ) -> Optional[str]:
+        """
+        Process and save PDF image data.
+        
+        Args:
+            image_data: Raw image binary data
+            xref: Image XRef number (for naming)
+            page_num: Page number (for naming)
+            **kwargs: Additional options
+            
+        Returns:
+            Image tag string, or None on failure
+        """
+        # Generate custom name based on XRef or page
+        custom_name = None
+        if xref is not None:
+            custom_name = f"pdf_xref_{xref}"
+        elif page_num is not None:
+            custom_name = f"pdf_page_{page_num}"
+        
+        return self.save_image(image_data, custom_name=custom_name)
+    
+    def process_xref_image(
+        self,
+        doc: "fitz.Document",
+        xref: int,
+    ) -> Optional[str]:
+        """
+        Extract and save image by XRef number.
+        
+        Args:
+            doc: PyMuPDF document object
+            xref: Image XRef number
+            
+        Returns:
+            Image tag string, or None on failure
+        """
+        try:
+            import fitz
+            
+            image_dict = doc.extract_image(xref)
+            if not image_dict:
+                return None
+            
+            image_data = image_dict.get("image")
+            if not image_data:
+                return None
+            
+            return self.process_image(image_data, xref=xref)
+            
+        except Exception as e:
+            self._logger.warning(f"Failed to extract XRef image {xref}: {e}")
+            return None
+    
+    def process_page_region(
+        self,
+        page: "fitz.Page",
+        rect: "fitz.Rect",
+        region_name: Optional[str] = None,
+    ) -> Optional[str]:
+        """
+        Render and save a page region as image.
+        
+        Used for complex regions that can't be represented as text.
+        
+        Args:
+            page: PyMuPDF page object
+            rect: Region rectangle to render
+            region_name: Optional name for the region
+            
+        Returns:
+            Image tag string, or None on failure
+        """
+        try:
+            import fitz
+            
+            # Calculate zoom for DPI
+            zoom = self._dpi / 72.0
+            mat = fitz.Matrix(zoom, zoom)
+            
+            # Clip to region
+            clip = rect
+            pix = page.get_pixmap(matrix=mat, clip=clip, alpha=False)
+            image_data = pix.tobytes("png")
+            
+            custom_name = region_name or f"pdf_page{page.number}_region"
+            return self.save_image(image_data, custom_name=custom_name)
+            
+        except Exception as e:
+            self._logger.warning(f"Failed to render page region: {e}")
+            return None
+    
+    def process_embedded_image(
+        self,
+        image_data: bytes,
+        image_name: Optional[str] = None,
+        xref: Optional[int] = None,
+        **kwargs
+    ) -> Optional[str]:
+        """
+        Process embedded PDF image.
+        
+        Args:
+            image_data: Image binary data
+            image_name: Original image name
+            xref: Image XRef number
+            **kwargs: Additional options
+            
+        Returns:
+            Image tag string, or None on failure
+        """
+        custom_name = image_name
+        if custom_name is None and xref is not None:
+            custom_name = f"pdf_embedded_{xref}"
+        
+        return self.save_image(image_data, custom_name=custom_name)
+    
+    def render_page(
+        self,
+        page: "fitz.Page",
+        alpha: bool = False,
+    ) -> Optional[str]:
+        """
+        Render entire page as image.
+        
+        Args:
+            page: PyMuPDF page object
+            alpha: Include alpha channel
+            
+        Returns:
+            Image tag string, or None on failure
+        """
+        try:
+            import fitz
+            
+            zoom = self._dpi / 72.0
+            mat = fitz.Matrix(zoom, zoom)
+            pix = page.get_pixmap(matrix=mat, alpha=alpha)
+            image_data = pix.tobytes("png")
+            
+            custom_name = f"pdf_page_{page.number + 1}_full"
+            return self.save_image(image_data, custom_name=custom_name)
+            
+        except Exception as e:
+            self._logger.warning(f"Failed to render page: {e}")
+            return None
+    
+    def extract_images_from_page(
+        self,
+        page: "fitz.Page",
+        page_num: int,
+        doc: "fitz.Document",
+        processed_images: Set[int],
+        table_bboxes: List[Tuple[float, float, float, float]],
+        min_image_size: int = 50,
+        min_image_area: int = 2500
+    ) -> List[Dict[str, Any]]:
+        """
+        Extract images from PDF page.
+        
+        This consolidates the logic from pdf_image.py extract_images_from_page().
+        
+        Args:
+            page: PyMuPDF page object
+            page_num: Page number (0-indexed)
+            doc: PyMuPDF document object
+            processed_images: Set of already processed image xrefs
+            table_bboxes: List of table bounding boxes to exclude
+            min_image_size: Minimum image dimension
+            min_image_area: Minimum image area
+            
+        Returns:
+            List of dicts with 'content', 'bbox', 'page_num' keys
+        """
+        from xgen_doc2chunk.core.processor.pdf_helpers.pdf_utils import (
+            find_image_position,
+            is_inside_any_bbox,
+        )
+        
+        elements = []
+        
+        try:
+            image_list = page.get_images()
+            
+            for img_info in image_list:
+                xref = img_info[0]
+                
+                if xref in processed_images:
+                    continue
+                
+                try:
+                    base_image = doc.extract_image(xref)
+                    if not base_image:
+                        continue
+                    
+                    image_bytes = base_image.get("image")
+                    width = base_image.get("width", 0)
+                    height = base_image.get("height", 0)
+                    
+                    if width < min_image_size or height < min_image_size:
+                        continue
+                    if width * height < min_image_area:
+                        continue
+                    
+                    img_bbox = find_image_position(page, xref)
+                    if img_bbox is None:
+                        continue
+                    
+                    if is_inside_any_bbox(img_bbox, table_bboxes, threshold=0.7):
+                        continue
+                    
+                    # Use format-specific process_image method
+                    image_tag = self.process_image(image_bytes, xref=xref, page_num=page_num)
+                    
+                    if image_tag:
+                        processed_images.add(xref)
+                        elements.append({
+                            'content': f'\n{image_tag}\n',
+                            'bbox': img_bbox,
+                            'page_num': page_num
+                        })
+                
+                except Exception as e:
+                    logger.debug(f"[PDF] Error extracting image xref={xref}: {e}")
+                    continue
+        
+        except Exception as e:
+            logger.warning(f"[PDF] Error extracting images: {e}")
+        
+        return elements
+
+
+__all__ = ["PDFImageProcessor"]