Create processor.py

processor.py

@@ -0,0 +1,138 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import cv2
+import numpy as np
+import numpy.typing as npt
+from loguru import logger
+from PIL import Image
+
+
+@dataclass
+class ProcessedImage:
+    image: npt.NDArray
+    angle: Optional[float] = None
+    enhancement_applied: bool = False
+    preprocessing_history: Optional[list[str]] = None
+
+    def __post_init__(self):
+        """Initialize preprocessing history if not provided."""
+        if self.preprocessing_history is None:
+            self.preprocessing_history = []
+        if isinstance(self.image, Image.Image):
+            self.image = np.array(self.image.convert("RGB"))
+
+
+class ImagePreprocessor:
+    @staticmethod
+    def denoise(image: npt.NDArray) -> ProcessedImage:
+        try:
+            if len(image.shape) == 3:
+                denoised = cv2.fastNlMeansDenoisingColored(image, None, 10, 10, 7, 21)
+            else:
+                denoised = cv2.fastNlMeansDenoising(image, None, 10, 7, 21)
+
+            return ProcessedImage(image=denoised, preprocessing_history=["denoise"])
+        except Exception as e:
+            logger.error(f"Error during denoising: {str(e)}")
+            raise ValueError(f"Failed to denoise image: {str(e)}")
+
+    @staticmethod
+    def deskew(image: npt.NDArray) -> ProcessedImage:
+        """Correct image skew by detecting and rotating to align text.
+
+        Uses contour detection to find the dominant text angle and corrects it.
+
+        Args:
+            image: Input image as numpy array
+
+        Returns:
+            ProcessedImage: Deskewed image with rotation angle
+
+        Raises:
+            ValueError: If angle detection fails
+        """
+        try:
+            # Convert to grayscale if needed
+            if len(image.shape) == 3:
+                gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+            else:
+                gray = image
+
+            # Detect edges
+            edges = cv2.Canny(gray, 50, 200, apertureSize=3)
+            lines = cv2.HoughLines(edges, 1, np.pi / 180, 100)
+
+            if lines is None:
+                logger.warning("No lines detected for deskewing")
+                return ProcessedImage(image=image, angle=0)
+
+            # Calculate dominant angle
+            angles = []
+            for _, theta in lines[0]:
+                angle = theta * 180 / np.pi
+                if angle < 45:
+                    angles.append(angle)
+                elif angle > 135:
+                    angles.append(angle - 180)
+
+            if not angles:
+                return ProcessedImage(image=image, angle=0)
+
+            median_angle = np.median(angles)
+
+            # Rotate image
+            (h, w) = image.shape[:2]
+            center = (w // 2, h // 2)
+            M = cv2.getRotationMatrix2D(center, median_angle, 1.0)
+            rotated = cv2.warpAffine(
+                image, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE
+            )
+
+            return ProcessedImage(
+                image=rotated, angle=median_angle, preprocessing_history=["deskew"]
+            )
+        except Exception as e:
+            logger.error(f"Error during deskewing: {str(e)}")
+            raise ValueError(f"Failed to deskew image: {str(e)}")
+
+
+@dataclass()
+class PreprocessingConfig:
+    """Configuration for image preprocessing steps."""
+
+    denoise: bool = True
+    deskew: bool = True
+    contrast_enhancement: bool = True
+    threshold: Optional[float] = None
+    resize_factor: Optional[float] = None
+
+
+class ImageHandler:
+    _supported_formats = {".png", ".jpg", ".jpeg", ".tiff", ".bmp"}
+
+    def __init__(self):
+        self.preprocessor = ImagePreprocessor()
+
+    def preprocess_image(
+        self, image: npt.NDArray, config: PreprocessingConfig
+    ) -> ProcessedImage:
+        try:
+            result = ProcessedImage(image=image)
+
+            if config.denoise:
+                denoised = self.preprocessor.denoise(result.image)
+                result.image = denoised.image
+                result.preprocessing_history.extend(denoised.preprocessing_history)
+
+            if config.deskew:
+                deskewed = self.preprocessor.deskew(result.image)
+                result.image = deskewed.image
+                result.angle = deskewed.angle
+                result.preprocessing_history.extend(deskewed.preprocessing_history)
+
+            return result
+
+        except Exception as e:
+            logger.error(f"Preprocessing failed: {str(e)}")
+            raise ValueError(f"Image preprocessing failed: {str(e)}")