src/ocr/preprocessor.py

# -*- coding: utf-8 -*-
"""
Image preprocessor with OCR quality enhancement.
Ported from OCR_Vehicle_02: CLAHE, deskew, denoise.
"""
import os
import gc
import logging
import numpy as np
from PIL import Image

logger = logging.getLogger(__name__)

try:
    import cv2
    CV2_AVAILABLE = True
except ImportError:
    CV2_AVAILABLE = False
    logger.warning("OpenCV not available. Image enhancement disabled.")

try:
    from pdf2image import convert_from_path
    PDF_SUPPORT = True
except ImportError:
    PDF_SUPPORT = False


class ImagePreprocessor:
    """
    Image preprocessor with OCR quality enhancement.

    Enhancement pipeline (from OCR_Vehicle_02):
    1. PDF → Image (300 DPI)
    2. Deskew (Hough Line Transform)
    3. Denoise (Gaussian Blur)
    4. Contrast enhancement (CLAHE on L channel)
    """
    DEFAULT_DPI = 300           # 300 DPI for better OCR (was 200)
    MAX_IMAGE_SIZE = 2048
    JPEG_QUALITY = 92           # Higher quality for OCR
    CLAHE_CLIP_LIMIT = 2.0
    CLAHE_GRID_SIZE = 8
    GAUSSIAN_KERNEL = 3

    def __init__(self, dpi=None, max_size=None):
        self.dpi = dpi or self.DEFAULT_DPI
        self.max_size = max_size or self.MAX_IMAGE_SIZE

    def _resize_image_if_needed(self, img):
        """Resize image if it exceeds maximum dimension."""
        width, height = img.size
        if width <= self.max_size and height <= self.max_size:
            return img

        if width > height:
            new_width = self.max_size
            new_height = int(height * (self.max_size / width))
        else:
            new_height = self.max_size
            new_width = int(width * (self.max_size / height))

        return img.resize((new_width, new_height), Image.Resampling.LANCZOS)

    def _enhance_image(self, img_path):
        """Apply OCR enhancement pipeline: deskew → denoise → CLAHE.

        Args:
            img_path: Path to image file

        Returns:
            Path to enhanced image (may be same path if no enhancement needed)
        """
        if not CV2_AVAILABLE:
            return img_path

        try:
            img = cv2.imread(img_path)
            if img is None:
                return img_path

            # Step 1: Deskew (correct rotation from scanning)
            img = self._deskew(img)

            # Step 2: Denoise (Gaussian blur)
            img = self._denoise(img)

            # Step 3: CLAHE contrast enhancement
            img = self._enhance_contrast(img)

            # Save enhanced image
            base, ext = os.path.splitext(img_path)
            enhanced_path = f"{base}_enhanced.jpg"
            cv2.imwrite(enhanced_path, img, [cv2.IMWRITE_JPEG_QUALITY, self.JPEG_QUALITY])

            return enhanced_path

        except Exception as e:
            logger.warning(f"Image enhancement failed: {e}")
            return img_path

    def _deskew(self, img):
        """Correct image skew using Hough Line Transform."""
        try:
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            edges = cv2.Canny(gray, 50, 150, apertureSize=3)
            lines = cv2.HoughLinesP(
                edges, 1, np.pi / 180,
                threshold=100, minLineLength=100, maxLineGap=10
            )

            if lines is None:
                return img

            # Collect angles of near-horizontal lines
            angles = []
            for line in lines:
                x1, y1, x2, y2 = line[0]
                if x2 - x1 == 0:
                    continue
                angle = np.degrees(np.arctan2(y2 - y1, x2 - x1))
                if abs(angle) < 15:  # Only near-horizontal lines
                    angles.append(angle)

            if not angles:
                return img

            median_angle = np.median(angles)

            # Only rotate if skew is significant (> 0.5 degrees)
            if abs(median_angle) < 0.5:
                return img

            h, w = img.shape[:2]
            center = (w // 2, h // 2)
            M = cv2.getRotationMatrix2D(center, median_angle, 1.0)
            rotated = cv2.warpAffine(
                img, M, (w, h),
                borderMode=cv2.BORDER_REPLICATE
            )

            logger.info(f"Deskew applied: {median_angle:.2f}°")
            return rotated

        except Exception as e:
            logger.warning(f"Deskew failed: {e}")
            return img

    def _denoise(self, img):
        """Remove noise with Gaussian blur."""
        try:
            return cv2.GaussianBlur(img, (self.GAUSSIAN_KERNEL, self.GAUSSIAN_KERNEL), 0)
        except Exception:
            return img

    def _enhance_contrast(self, img):
        """Enhance contrast using CLAHE on L channel (preserves color)."""
        try:
            lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
            l_channel, a_channel, b_channel = cv2.split(lab)

            clahe = cv2.createCLAHE(
                clipLimit=self.CLAHE_CLIP_LIMIT,
                tileGridSize=(self.CLAHE_GRID_SIZE, self.CLAHE_GRID_SIZE)
            )
            l_enhanced = clahe.apply(l_channel)

            lab_enhanced = cv2.merge([l_enhanced, a_channel, b_channel])
            enhanced = cv2.cvtColor(lab_enhanced, cv2.COLOR_LAB2BGR)

            return enhanced

        except Exception as e:
            logger.warning(f"CLAHE failed: {e}")
            return img

    def load_image(self, image_path, max_pages=None):
        """Load and enhance image/PDF for OCR processing."""
        try:
            if not os.path.exists(image_path):
                raise FileNotFoundError(f"Image not found: {image_path}")

            if image_path.lower().endswith('.pdf'):
                return self._process_pdf(image_path, max_pages)

            return self._process_image(image_path)

        except Exception as e:
            logger.error(f"Error loading image: {e}")
            return []

    def _process_pdf(self, pdf_path, max_pages=None):
        """Process PDF with enhancement pipeline."""
        if not PDF_SUPPORT:
            logger.warning(f"PDF support unavailable. Skipping {pdf_path}")
            return []

        try:
            convert_kwargs = {
                'dpi': self.dpi,
                'fmt': 'jpeg',
                'thread_count': 1,
                'use_pdftocairo': True,
            }
            if max_pages:
                convert_kwargs['last_page'] = max_pages

            pages = convert_from_path(pdf_path, **convert_kwargs)

            temp_image_paths = []
            base_name = os.path.splitext(os.path.basename(pdf_path))[0]
            import tempfile
            dir_name = tempfile.gettempdir()

            for i, page in enumerate(pages):
                try:
                    resized_page = self._resize_image_if_needed(page)
                    temp_path = os.path.join(dir_name, f"{base_name}_page_{i+1}.jpg")
                    resized_page.save(temp_path, 'JPEG', quality=self.JPEG_QUALITY)

                    if resized_page != page:
                        resized_page.close()
                    page.close()

                    # Apply enhancement pipeline
                    enhanced_path = self._enhance_image(temp_path)
                    if enhanced_path != temp_path and os.path.exists(temp_path):
                        os.remove(temp_path)  # Remove unenhanced version

                    temp_image_paths.append(enhanced_path)

                except Exception as e:
                    logger.warning(f"Failed to process page {i+1}: {e}")

            del pages
            gc.collect()

            return temp_image_paths

        except Exception as e:
            logger.error(f"Failed to convert PDF {pdf_path}: {e}")
            return []

    def _process_image(self, image_path):
        """Process image file with enhancement."""
        try:
            with Image.open(image_path) as img:
                img.verify()

            with Image.open(image_path) as img:
                width, height = img.size
                if width > self.max_size or height > self.max_size:
                    resized = self._resize_image_if_needed(img.copy())
                    base_name = os.path.splitext(os.path.basename(image_path))[0]
                    import tempfile
                    dir_name = tempfile.gettempdir()
                    temp_path = os.path.join(dir_name, f"{base_name}_resized.jpg")
                    resized.save(temp_path, 'JPEG', quality=self.JPEG_QUALITY)
                    resized.close()

                    enhanced_path = self._enhance_image(temp_path)
                    if enhanced_path != temp_path and os.path.exists(temp_path):
                        os.remove(temp_path)
                    return [enhanced_path]

            # No resize needed — enhance in place (to temp file)
            enhanced_path = self._enhance_image(image_path)
            return [enhanced_path]

        except Exception as e:
            logger.error(f"Error processing image {image_path}: {e}")
            return []

    @staticmethod
    def cleanup_temp_files(file_paths, original_path):
        """Remove temporary files created during processing."""
        for path in file_paths:
            if path != original_path and os.path.exists(path):
                try:
                    os.remove(path)
                except Exception as e:
                    logger.warning(f"Failed to cleanup {path}: {e}")