fmard/ocr-api

1"""2Image Preprocessing module for KTP OCR.3 4Performs orientation correction, document boundary detection,5contrast/sharpness enhancement, blur detection, and resizing6before passing the image to PaddleOCR.7 8Usage:9    from preprocessor import preprocess10 11    processed_image, metadata = preprocess(image_bytes)12    # processed_image: np.ndarray (BGR)13    # metadata: dict with processing details14"""15 16import logging17from typing import Tuple18 19import cv220import numpy as np21 22logger = logging.getLogger(__name__)23 24# Constants25MAX_WIDTH = 200026BLUR_THRESHOLD_SHARP = 100.0  # Laplacian variance above this = sharp27BLUR_THRESHOLD_BLURRY = 30.0  # Below this = very blurry28MIN_CONTOUR_AREA_RATIO = 0.1  # Minimum contour area relative to image area29CLAHE_CLIP_LIMIT = 2.030CLAHE_TILE_SIZE = (8, 8)31SHARPEN_AMOUNT = 0.532 33 34def preprocess(image_bytes: bytes) -> Tuple[np.ndarray, dict]:35    """36    Preprocess an image for KTP OCR extraction.37 38    Steps:39        1. Decode image from bytes40        2. Auto-detect orientation and rotate if needed41        3. Detect document boundary and crop42        4. Enhance contrast and sharpness43        5. Calculate blur score44        6. Resize to optimal resolution45 46    Args:47        image_bytes: Raw image data (JPEG, PNG, etc.)48 49    Returns:50        Tuple of (processed_image_numpy_array, metadata_dict)51 52        metadata_dict contains:53            - rotated: bool — whether orientation was corrected54            - rotation_angle: int — degrees rotated (0, 90, 180, 270)55            - cropped: bool — whether document boundary was detected and cropped56            - enhanced: bool — whether enhancement was applied57            - blur_score: float — 0.0 (sharp) to 1.0 (very blurry)58            - original_size: tuple (width, height)59            - processed_size: tuple (width, height)60 61    Raises:62        ValueError: If image_bytes cannot be decoded.63    """64    metadata = {65        "rotated": False,66        "rotation_angle": 0,67        "cropped": False,68        "enhanced": False,69        "blur_score": 0.0,70        "original_size": (0, 0),71        "processed_size": (0, 0),72    }73 74    # Step 1: Decode image75    image = _decode_image(image_bytes)76    h, w = image.shape[:2]77    metadata["original_size"] = (w, h)78    logger.info(f"Preprocessing image: {w}x{h}")79 80    # Step 2: Orientation correction81    image, rotated, angle = _correct_orientation(image)82    metadata["rotated"] = rotated83    metadata["rotation_angle"] = angle84 85    # Step 3: Document boundary detection and crop86    image, cropped = _detect_and_crop_document(image)87    metadata["cropped"] = cropped88 89    # Step 4: Enhance contrast and sharpness90    image, enhanced = _enhance_image(image)91    metadata["enhanced"] = enhanced92 93    # Step 5: Calculate blur score94    blur_score = _calculate_blur_score(image)95    metadata["blur_score"] = round(blur_score, 4)96 97    # Step 6: Resize to optimal resolution98    image = _resize_for_ocr(image)99    h_final, w_final = image.shape[:2]100    metadata["processed_size"] = (w_final, h_final)101 102    logger.info(103        f"Preprocessing complete: {metadata['original_size']} -> "104        f"{metadata['processed_size']}, blur={metadata['blur_score']:.3f}, "105        f"rotated={rotated}, cropped={cropped}, enhanced={enhanced}"106    )107 108    return image, metadata109 110 111def _decode_image(image_bytes: bytes) -> np.ndarray:112    """Decode raw bytes into a BGR numpy array."""113    buf = np.frombuffer(image_bytes, dtype=np.uint8)114    image = cv2.imdecode(buf, cv2.IMREAD_COLOR)115    if image is None:116        raise ValueError("Failed to decode image bytes. Unsupported or corrupt image.")117    return image118 119 120def _correct_orientation(image: np.ndarray) -> Tuple[np.ndarray, bool, int]:121    """122    Auto-detect orientation and rotate if needed.123 124    Strategy:125        - KTP cards are landscape-oriented (wider than tall).126        - If the image is portrait (taller than wide), rotate 90° CW.127        - Use text line detection via Hough transform to detect skew128          and correct minor rotation.129 130    Returns:131        (corrected_image, was_rotated, rotation_angle_degrees)132    """133    h, w = image.shape[:2]134    rotated = False135    angle = 0136 137    # If portrait orientation, rotate to landscape (KTP is landscape)138    if h > w:139        image = cv2.rotate(image, cv2.ROTATE_90_CLOCKWISE)140        rotated = True141        angle = 90142        h, w = image.shape[:2]143        logger.debug("Rotated portrait image to landscape (90° CW).")144 145    # Fine-tune rotation using text line detection (deskew)146    deskew_angle = _detect_skew(image)147    if abs(deskew_angle) > 0.5:  # Only correct if skew > 0.5 degrees148        image = _rotate_image(image, deskew_angle)149        rotated = True150        angle += int(round(deskew_angle))151        logger.debug(f"Deskewed by {deskew_angle:.2f} degrees.")152 153    return image, rotated, angle154 155 156def _detect_skew(image: np.ndarray) -> float:157    """158    Detect text skew angle using Hough Line Transform.159 160    Returns the estimated skew angle in degrees (negative = clockwise tilt).161    Returns 0 if unable to detect reliable skew.162    """163    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)164    # Use Canny edge detection165    edges = cv2.Canny(gray, 50, 150, apertureSize=3)166 167    # Detect lines using probabilistic Hough transform168    lines = cv2.HoughLinesP(169        edges, 1, np.pi / 180, threshold=100,170        minLineLength=image.shape[1] // 8,171        maxLineGap=10172    )173 174    if lines is None or len(lines) < 5:175        return 0.0176 177    # Calculate angles of detected lines178    angles = []179    for line in lines:180        x1, y1, x2, y2 = line[0]181        dx = x2 - x1182        dy = y2 - y1183        # Only consider nearly-horizontal lines (within ±30°)184        angle = np.degrees(np.arctan2(dy, dx))185        if abs(angle) < 30:186            angles.append(angle)187 188    if not angles:189        return 0.0190 191    # Use median angle to be robust against outliers192    median_angle = float(np.median(angles))193 194    # Limit correction to small angles (avoid large incorrect rotations)195    if abs(median_angle) > 10:196        return 0.0197 198    return -median_angle199 200 201def _rotate_image(image: np.ndarray, angle: float) -> np.ndarray:202    """Rotate image by a small angle around its center without cropping."""203    h, w = image.shape[:2]204    center = (w / 2, h / 2)205    rotation_matrix = cv2.getRotationMatrix2D(center, angle, 1.0)206 207    # Compute new bounding dimensions208    cos_val = abs(rotation_matrix[0, 0])209    sin_val = abs(rotation_matrix[0, 1])210    new_w = int(h * sin_val + w * cos_val)211    new_h = int(h * cos_val + w * sin_val)212 213    # Adjust the rotation matrix for translation214    rotation_matrix[0, 2] += (new_w / 2) - center[0]215    rotation_matrix[1, 2] += (new_h / 2) - center[1]216 217    rotated = cv2.warpAffine(218        image, rotation_matrix, (new_w, new_h),219        flags=cv2.INTER_CUBIC,220        borderMode=cv2.BORDER_REPLICATE221    )222    return rotated223 224 225def _detect_and_crop_document(image: np.ndarray) -> Tuple[np.ndarray, bool]:226    """227    Detect the KTP card boundary and crop to the document region.228 229    Strategy:230        - Convert to grayscale and apply bilateral filter (edge-preserving)231        - Detect edges with Canny232        - Find contours and look for the largest quadrilateral233        - Apply perspective transform to get a flat top-down view234 235    Returns:236        (cropped_image, was_cropped)237    """238    h, w = image.shape[:2]239    image_area = h * w240 241    # Work on a scaled-down copy for faster contour detection242    scale = 1.0243    if max(h, w) > 1000:244        scale = 1000.0 / max(h, w)245        small = cv2.resize(image, None, fx=scale, fy=scale)246    else:247        small = image.copy()248 249    gray = cv2.cvtColor(small, cv2.COLOR_BGR2GRAY)250    blurred = cv2.bilateralFilter(gray, 11, 17, 17)251    edges = cv2.Canny(blurred, 30, 200)252 253    # Dilate to close gaps in edges254    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))255    edges = cv2.dilate(edges, kernel, iterations=1)256 257    # Find contours258    contours, _ = cv2.findContours(259        edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE260    )261 262    if not contours:263        return image, False264 265    # Sort contours by area (largest first)266    contours = sorted(contours, key=cv2.contourArea, reverse=True)267 268    doc_contour = None269    small_area = small.shape[0] * small.shape[1]270 271    for contour in contours[:10]:  # Check top 10 largest contours272        area = cv2.contourArea(contour)273        if area < small_area * MIN_CONTOUR_AREA_RATIO:274            continue275 276        # Approximate the contour to a polygon277        peri = cv2.arcLength(contour, True)278        approx = cv2.approxPolyDP(contour, 0.02 * peri, True)279 280        # A document should have 4 corners281        if len(approx) == 4:282            doc_contour = approx283            break284 285    if doc_contour is None:286        # Fallback: no clear quadrilateral found287        return image, False288 289    # Scale contour points back to original size290    doc_contour = (doc_contour.astype(np.float64) / scale).astype(np.int32)291 292    # Apply perspective transform293    cropped = _perspective_transform(image, doc_contour)294    if cropped is not None:295        return cropped, True296 297    return image, False298 299 300def _perspective_transform(301    image: np.ndarray, contour: np.ndarray302) -> np.ndarray | None:303    """304    Apply perspective transform to straighten a quadrilateral region.305 306    Args:307        image: Source image308        contour: 4-point contour array309 310    Returns:311        Warped (straightened) image or None if transform fails.312    """313    try:314        pts = contour.reshape(4, 2).astype(np.float32)315 316        # Order points: top-left, top-right, bottom-right, bottom-left317        ordered = _order_points(pts)318        tl, tr, br, bl = ordered319 320        # Compute width and height of the new image321        width_top = np.linalg.norm(tr - tl)322        width_bottom = np.linalg.norm(br - bl)323        max_width = int(max(width_top, width_bottom))324 325        height_left = np.linalg.norm(bl - tl)326        height_right = np.linalg.norm(br - tr)327        max_height = int(max(height_left, height_right))328 329        if max_width < 100 or max_height < 50:330            return None331 332        dst = np.array([333            [0, 0],334            [max_width - 1, 0],335            [max_width - 1, max_height - 1],336            [0, max_height - 1],337        ], dtype=np.float32)338 339        matrix = cv2.getPerspectiveTransform(ordered, dst)340        warped = cv2.warpPerspective(image, matrix, (max_width, max_height))341        return warped342 343    except Exception as e:344        logger.warning(f"Perspective transform failed: {e}")345        return None346 347 348def _order_points(pts: np.ndarray) -> np.ndarray:349    """350    Order 4 points as: top-left, top-right, bottom-right, bottom-left.351 352    Uses sum and difference of coordinates to determine positions.353    """354    ordered = np.zeros((4, 2), dtype=np.float32)355 356    s = pts.sum(axis=1)357    ordered[0] = pts[np.argmin(s)]   # Top-left has smallest sum358    ordered[2] = pts[np.argmax(s)]   # Bottom-right has largest sum359 360    d = np.diff(pts, axis=1)361    ordered[1] = pts[np.argmin(d)]   # Top-right has smallest difference362    ordered[3] = pts[np.argmax(d)]   # Bottom-left has largest difference363 364    return ordered365 366 367def _enhance_image(image: np.ndarray) -> Tuple[np.ndarray, bool]:368    """369    Enhance image contrast and sharpness for better OCR.370 371    Steps:372        - Apply CLAHE (Contrast Limited Adaptive Histogram Equalization)373          on the L channel in LAB color space374        - Apply unsharp mask for sharpening375 376    Returns:377        (enhanced_image, was_enhanced)378    """379    try:380        # Convert to LAB color space for luminance-only enhancement381        lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)382        l_channel, a_channel, b_channel = cv2.split(lab)383 384        # Apply CLAHE to the L channel385        clahe = cv2.createCLAHE(386            clipLimit=CLAHE_CLIP_LIMIT, tileGridSize=CLAHE_TILE_SIZE387        )388        l_enhanced = clahe.apply(l_channel)389 390        # Merge and convert back to BGR391        enhanced_lab = cv2.merge([l_enhanced, a_channel, b_channel])392        enhanced = cv2.cvtColor(enhanced_lab, cv2.COLOR_LAB2BGR)393 394        # Apply unsharp mask for sharpening395        enhanced = _unsharp_mask(enhanced, amount=SHARPEN_AMOUNT)396 397        return enhanced, True398 399    except Exception as e:400        logger.warning(f"Image enhancement failed, using original: {e}")401        return image, False402 403 404def _unsharp_mask(405    image: np.ndarray, amount: float = 0.5, sigma: float = 1.0406) -> np.ndarray:407    """408    Apply unsharp mask sharpening.409 410    Formula: sharpened = original + amount * (original - blurred)411    """412    blurred = cv2.GaussianBlur(image, (0, 0), sigma)413    sharpened = cv2.addWeighted(image, 1.0 + amount, blurred, -amount, 0)414    return sharpened415 416 417def _calculate_blur_score(image: np.ndarray) -> float:418    """419    Calculate a blur score from 0.0 (perfectly sharp) to 1.0 (very blurry).420 421    Uses the variance of the Laplacian as a focus measure.422    Higher Laplacian variance = sharper image.423    """424    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)425    laplacian_var = cv2.Laplacian(gray, cv2.CV_64F).var()426 427    # Map Laplacian variance to 0-1 score (inverted: 0=sharp, 1=blurry)428    if laplacian_var >= BLUR_THRESHOLD_SHARP:429        return 0.0430    elif laplacian_var <= BLUR_THRESHOLD_BLURRY:431        return 1.0432    else:433        # Linear interpolation between thresholds434        score = 1.0 - (435            (laplacian_var - BLUR_THRESHOLD_BLURRY)436            / (BLUR_THRESHOLD_SHARP - BLUR_THRESHOLD_BLURRY)437        )438        return float(np.clip(score, 0.0, 1.0))439 440 441def _resize_for_ocr(image: np.ndarray) -> np.ndarray:442    """443    Resize image to optimal resolution for OCR.444 445    Ensures maximum width is MAX_WIDTH pixels while maintaining aspect ratio.446    Only downscales — small images are left as-is.447    """448    h, w = image.shape[:2]449 450    if w <= MAX_WIDTH:451        return image452 453    scale = MAX_WIDTH / w454    new_w = MAX_WIDTH455    new_h = int(h * scale)456 457    resized = cv2.resize(458        image, (new_w, new_h), interpolation=cv2.INTER_AREA459    )460    logger.debug(f"Resized from {w}x{h} to {new_w}x{new_h}")461    return resized462