algoryn/dots-ocr-idcard

1"""Image and PDF preprocessing utilities for Dots.OCR.2 3This module handles PDF to image conversion, image preprocessing,4and multi-page document processing for the Dots.OCR model.5"""6 7import os8import logging9from typing import List, Tuple, Optional, Union10from pathlib import Path11import io12 13import fitz  # PyMuPDF14import numpy as np15from PIL import Image, ImageOps16import cv217 18# Configure logging19logger = logging.getLogger(__name__)20 21# Environment variable configuration22PDF_DPI = int(os.getenv("DOTS_OCR_PDF_DPI", "300"))23PDF_MAX_PAGES = int(os.getenv("DOTS_OCR_PDF_MAX_PAGES", "10"))24IMAGE_MAX_SIZE = (25    int(os.getenv("DOTS_OCR_IMAGE_MAX_SIZE", "10")) * 1024 * 102426)  # 10MB default27 28 29class ImagePreprocessor:30    """Handles image preprocessing for Dots.OCR model."""31 32    def __init__(33        self, min_pixels: int = 3136, max_pixels: int = 11289600, divisor: int = 2834    ):35        """Initialize the image preprocessor.36 37        Args:38            min_pixels: Minimum pixel count for images39            max_pixels: Maximum pixel count for images40            divisor: Required divisor for image dimensions41        """42        self.min_pixels = min_pixels43        self.max_pixels = max_pixels44        self.divisor = divisor45 46    def preprocess_image(self, image: Image.Image) -> Image.Image:47        """Preprocess an image to meet model requirements.48 49        Args:50            image: Input PIL Image51 52        Returns:53            Preprocessed PIL Image54        """55        # Convert to RGB if necessary56        if image.mode != "RGB":57            image = image.convert("RGB")58 59        # Auto-orient image based on EXIF data60        image = ImageOps.exif_transpose(image)61 62        # Calculate current pixel count63        width, height = image.size64        current_pixels = width * height65 66        logger.info(f"Original image size: {width}x{height} ({current_pixels} pixels)")67 68        # Resize if necessary to meet pixel requirements69        if current_pixels < self.min_pixels:70            # Scale up to meet minimum pixel requirement71            scale_factor = (self.min_pixels / current_pixels) ** 0.572            new_width = int(width * scale_factor)73            new_height = int(height * scale_factor)74            image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)75            logger.info(f"Scaled up image to {new_width}x{new_height}")76 77        elif current_pixels > self.max_pixels:78            # Scale down to meet maximum pixel requirement79            scale_factor = (self.max_pixels / current_pixels) ** 0.580            new_width = int(width * scale_factor)81            new_height = int(height * scale_factor)82            image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)83            logger.info(f"Scaled down image to {new_width}x{new_height}")84 85        # Ensure dimensions are divisible by the required divisor86        width, height = image.size87        new_width = ((width + self.divisor - 1) // self.divisor) * self.divisor88        new_height = ((height + self.divisor - 1) // self.divisor) * self.divisor89 90        if new_width != width or new_height != height:91            image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)92            logger.info(93                f"Adjusted dimensions to be divisible by {self.divisor}: {new_width}x{new_height}"94            )95 96        return image97 98    def crop_by_roi(99        self, image: Image.Image, roi: Tuple[float, float, float, float]100    ) -> Image.Image:101        """Crop image using ROI coordinates.102 103        Args:104            image: Input PIL Image105            roi: ROI coordinates as (x1, y1, x2, y2) normalized to [0, 1]106 107        Returns:108            Cropped PIL Image109        """110        x1, y1, x2, y2 = roi111        width, height = image.size112 113        # Convert normalized coordinates to pixel coordinates114        x1_px = int(x1 * width)115        y1_px = int(y1 * height)116        x2_px = int(x2 * width)117        y2_px = int(y2 * height)118 119        # Ensure coordinates are within image bounds120        x1_px = max(0, min(x1_px, width))121        y1_px = max(0, min(y1_px, height))122        x2_px = max(x1_px, min(x2_px, width))123        y2_px = max(y1_px, min(y2_px, height))124 125        # Crop the image126        cropped = image.crop((x1_px, y1_px, x2_px, y2_px))127        logger.info(f"Cropped image to {x2_px - x1_px}x{y2_px - y1_px} pixels")128 129        return cropped130 131 132class PDFProcessor:133    """Handles PDF to image conversion and multi-page processing."""134 135    def __init__(self, dpi: int = PDF_DPI, max_pages: int = PDF_MAX_PAGES):136        """Initialize the PDF processor.137 138        Args:139            dpi: DPI for PDF to image conversion140            max_pages: Maximum number of pages to process141        """142        self.dpi = dpi143        self.max_pages = max_pages144 145    def pdf_to_images(self, pdf_data: bytes) -> List[Image.Image]:146        """Convert PDF to list of images.147 148        Args:149            pdf_data: PDF file data as bytes150 151        Returns:152            List of PIL Images, one per page153        """154        try:155            # Open PDF from bytes156            pdf_document = fitz.open(stream=pdf_data, filetype="pdf")157            images = []158 159            # Limit number of pages to process160            num_pages = min(len(pdf_document), self.max_pages)161            logger.info(f"Processing {num_pages} pages from PDF")162 163            for page_num in range(num_pages):164                page = pdf_document[page_num]165 166                # Convert page to image167                mat = fitz.Matrix(self.dpi / 72, self.dpi / 72)  # 72 is default DPI168                pix = page.get_pixmap(matrix=mat)169 170                # Convert to PIL Image171                img_data = pix.tobytes("png")172                image = Image.open(io.BytesIO(img_data))173                images.append(image)174 175                logger.info(f"Converted page {page_num + 1} to image: {image.size}")176 177            pdf_document.close()178            return images179 180        except Exception as e:181            logger.error(f"Failed to convert PDF to images: {e}")182            raise RuntimeError(f"PDF conversion failed: {e}")183 184    def is_pdf(self, file_data: bytes) -> bool:185        """Check if file data is a PDF.186 187        Args:188            file_data: File data as bytes189 190        Returns:191            True if file is a PDF192        """193        return file_data.startswith(b"%PDF-")194 195    def get_pdf_page_count(self, pdf_data: bytes) -> int:196        """Get the number of pages in a PDF.197 198        Args:199            pdf_data: PDF file data as bytes200 201        Returns:202            Number of pages in the PDF203        """204        try:205            pdf_document = fitz.open(stream=pdf_data, filetype="pdf")206            page_count = len(pdf_document)207            pdf_document.close()208            return page_count209        except Exception as e:210            logger.error(f"Failed to get PDF page count: {e}")211            return 0212 213 214class DocumentProcessor:215    """Main document processing class that handles both images and PDFs."""216 217    def __init__(self):218        """Initialize the document processor."""219        self.image_preprocessor = ImagePreprocessor()220        self.pdf_processor = PDFProcessor()221 222    def process_document(223        self, file_data: bytes, roi: Optional[Tuple[float, float, float, float]] = None224    ) -> List[Image.Image]:225        """Process a document (image or PDF) and return preprocessed images.226 227        Args:228            file_data: Document file data as bytes229            roi: Optional ROI coordinates as (x1, y1, x2, y2) normalized to [0, 1]230 231        Returns:232            List of preprocessed PIL Images233        """234        # Check if it's a PDF235        if self.pdf_processor.is_pdf(file_data):236            logger.info("Processing PDF document")237            images = self.pdf_processor.pdf_to_images(file_data)238        else:239            # Process as image240            logger.info("Processing image document")241            try:242                image = Image.open(io.BytesIO(file_data))243                images = [image]244            except Exception as e:245                logger.error(f"Failed to open image: {e}")246                raise RuntimeError(f"Image processing failed: {e}")247 248        # Preprocess each image249        processed_images = []250        for i, image in enumerate(images):251            try:252                # Apply ROI cropping if provided253                if roi is not None:254                    image = self.image_preprocessor.crop_by_roi(image, roi)255 256                # Preprocess image for model requirements257                processed_image = self.image_preprocessor.preprocess_image(image)258                processed_images.append(processed_image)259 260                logger.info(f"Processed image {i + 1}: {processed_image.size}")261 262            except Exception as e:263                logger.error(f"Failed to preprocess image {i + 1}: {e}")264                # Continue with other images even if one fails265                continue266 267        if not processed_images:268            raise RuntimeError("No images could be processed from the document")269 270        logger.info(f"Successfully processed {len(processed_images)} images")271        return processed_images272 273    def validate_file_size(self, file_data: bytes) -> bool:274        """Validate that file size is within limits.275 276        Args:277            file_data: File data as bytes278 279        Returns:280            True if file size is acceptable281        """282        file_size = len(file_data)283        if file_size > IMAGE_MAX_SIZE:284            logger.warning(f"File size {file_size} exceeds limit {IMAGE_MAX_SIZE}")285            return False286        return True287 288    def get_document_info(self, file_data: bytes) -> dict:289        """Get information about the document.290 291        Args:292            file_data: Document file data as bytes293 294        Returns:295            Dictionary with document information296        """297        info = {298            "file_size": len(file_data),299            "is_pdf": self.pdf_processor.is_pdf(file_data),300            "page_count": 1,301        }302 303        if info["is_pdf"]:304            info["page_count"] = self.pdf_processor.get_pdf_page_count(file_data)305 306        return info307 308 309# Global document processor instance310_document_processor: Optional[DocumentProcessor] = None311 312 313def get_document_processor() -> DocumentProcessor:314    """Get the global document processor instance."""315    global _document_processor316    if _document_processor is None:317        _document_processor = DocumentProcessor()318    return _document_processor319 320 321def process_document(322    file_data: bytes, roi: Optional[Tuple[float, float, float, float]] = None323) -> List[Image.Image]:324    """Process a document and return preprocessed images."""325    processor = get_document_processor()326    return processor.process_document(file_data, roi)327 328 329def validate_file_size(file_data: bytes) -> bool:330    """Validate that file size is within limits."""331    processor = get_document_processor()332    return processor.validate_file_size(file_data)333 334 335def get_document_info(file_data: bytes) -> dict:336    """Get information about the document."""337    processor = get_document_processor()338    return processor.get_document_info(file_data)339