pdf_processor_module.py

"""
PDF Processor Module
Handles PDF conversion and image preprocessing
"""

import cv2
import numpy as np
from PIL import Image
import pytesseract
from pdf2image import convert_from_path
import os


class PDFProcessor:
    def __init__(self):


        
        pass
    
    def convert_to_images(self, pdf_path, dpi=300):
        """
        Convert PDF pages to images
        
        Args:
            pdf_path: Path to PDF file
            dpi: Resolution for conversion (higher = better quality)
        
        Returns:
            List of numpy arrays (images)
        """
        try:
            # Convert PDF to PIL images
            pil_images = convert_from_path(pdf_path, dpi=dpi)
            
            # Convert PIL images to numpy arrays (OpenCV format)
            cv_images = []
            for pil_img in pil_images:
                # Convert PIL to OpenCV format (RGB to BGR)
                cv_img = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
                cv_images.append(cv_img)
            
            return cv_images
        
        except Exception as e:
            raise Exception(f"Failed to convert PDF: {str(e)}")
    
    def load_image(self, image_path):
        """
        Load a single image file
        
        Args:
            image_path: Path to image file
        
        Returns:
            Numpy array (OpenCV format)
        """
        try:
            img = cv2.imread(image_path)
            if img is None:
                raise ValueError("Could not read image file")
            return img
        except Exception as e:
            raise Exception(f"Failed to load image: {str(e)}")
    
    def preprocess_image(self, image):
        # Convert to grayscale
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # Deskew image
        gray = self._deskew(gray)
        
        # Remove noise
        denoised = cv2.fastNlMeansDenoising(gray)
        
        # Adaptive thresholding for better contrast
        binary = cv2.adaptiveThreshold(
            denoised,
            255,
            cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
            cv2.THRESH_BINARY,
            11,
            2
        )
        
        return binary
    
    def _deskew(self, image):
        """
        Detect and correct image skew
        """
        coords = np.column_stack(np.where(image > 0))
        if len(coords) == 0:
            return image
        
        angle = cv2.minAreaRect(coords)[-1]
        
        # Adjust angle
        if angle < -45:
            angle = -(90 + angle)
        else:
            angle = -angle
        
        # Rotate image if skew is significant
        if abs(angle) > 0.5:  # Only correct if angle > 0.5 degrees
            (h, w) = image.shape[:2]
            center = (w // 2, h // 2)
            M = cv2.getRotationMatrix2D(center, angle, 1.0)
            rotated = cv2.warpAffine(
                image,
                M,
                (w, h),
                flags=cv2.INTER_CUBIC,
                borderMode=cv2.BORDER_REPLICATE
            )
            return rotated
        
        return image
    
    def extract_text_data(self, image):
        """
        Extract text data from image using OCR
        
        Args:
            image: Input image
        
        Returns:
            List of text regions with coordinates
        """
        try:
            # Preprocess image
            processed = self.preprocess_image(image)
            
            # Use Tesseract to get detailed data
            data = pytesseract.image_to_data(
                processed,
                output_type=pytesseract.Output.DICT
            )
            
            text_regions = []
            n_boxes = len(data['text'])
            
            for i in range(n_boxes):
                # Filter out empty text and low confidence
                if int(data['conf'][i]) > 60 and data['text'][i].strip():
                    text_regions.append({
                        'text': data['text'][i],
                        'confidence': int(data['conf'][i]),
                        'coordinates': {
                            'x': data['left'][i],
                            'y': data['top'][i],
                            'width': data['width'][i],
                            'height': data['height'][i]
                        }
                    })
            
            return text_regions
        
        except Exception as e:
            print(f"Error extracting text: {str(e)}")
            return []
    
    def detect_regions(self, image):
        """
        Detect distinct regions in image (tables, charts, text blocks)
        
        Args:
            image: Input image
        
        Returns:
            List of region coordinates
        """
        # Convert to grayscale
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # Apply binary threshold
        _, binary = cv2.threshold(gray, 240, 255, cv2.THRESH_BINARY_INV)
        
        # Morphological operations to connect components
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (50, 50))
        dilated = cv2.dilate(binary, kernel, iterations=2)
        
        # Find contours
        contours, _ = cv2.findContours(
            dilated,
            cv2.RETR_EXTERNAL,
            cv2.CHAIN_APPROX_SIMPLE
        )
        
        regions = []
        for contour in contours:
            x, y, w, h = cv2.boundingRect(contour)
            
            # Filter small regions
            if w > 100 and h > 100:
                regions.append({
                    'x': x,
                    'y': y,
                    'width': w,
                    'height': h,
                    'area': w * h
                })
        
        # Sort regions by position (top to bottom, left to right)
        regions.sort(key=lambda r: (r['y'], r['x']))
        
        return regions
    
    def crop_region(self, image, region):
        """
        Crop a specific region from image
        
        Args:
            image: Full image
            region: Dictionary with x, y, width, height
        
        Returns:
            Cropped image
        """
        x, y, w, h = region['x'], region['y'], region['width'], region['height']
        return image[y:y+h, x:x+w]