image_preprocessor.py

import cv2
import numpy as np
import argparse
import os
from pathlib import Path


class ImagePreprocessor:
    def __init__(self,
                 to_grayscale: bool = True,
                 normalize_bg: bool = True,
                 denoise: bool = True,
                 denoise_method: str = "gaussian",  # gaussian | nlm
                 clahe: bool = True,
                 clahe_clip: float = 2.0,
                 clahe_grid: int = 8,
                 binarize: bool = True,
                 binarize_method: str = "otsu",  # otsu | adaptive | sauvola
                 deskew: bool = True,
                 sharpen: bool = True,
                 morph_clean: bool = True,
                 target_height: int = 48,
                 padding: int = 4):
        self.to_grayscale = to_grayscale
        self.normalize_bg = normalize_bg
        self.denoise = denoise
        self.denoise_method = denoise_method
        self.clahe = clahe
        self.clahe_clip = clahe_clip
        self.clahe_grid = clahe_grid
        self.binarize = binarize
        self.binarize_method = binarize_method
        self.deskew = deskew
        self.sharpen = sharpen
        self.morph_clean = morph_clean
        self.target_height = target_height
        self.padding = padding

    def process(self, image: np.ndarray,
                verbose: bool = False) -> np.ndarray:

        if image is None or image.size == 0:
            return image

        steps = {}
        img = image.copy()
        steps["original"] = img.copy()

        # 1. Grayscale
        if len(img.shape) == 3:
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        else:
            gray = img.copy()
        steps["grayscale"] = gray.copy()

        if self.normalize_bg:
            gray = self._normalize_background(gray)
            steps["bg_norm"] = gray.copy()

        if self.denoise:
            if self.denoise_method == "nlm":
                gray = cv2.fastNlMeansDenoising(gray, h=7,
                    templateWindowSize=7, searchWindowSize=21)
            else:
                gray = cv2.GaussianBlur(gray, (3, 3), 0)
            steps["denoise"] = gray.copy()

        if self.clahe:
            clahe_obj = cv2.createCLAHE(
                clipLimit=self.clahe_clip,
                tileGridSize=(self.clahe_grid, self.clahe_grid)
            )
            gray = clahe_obj.apply(gray)
            steps["clahe"] = gray.copy()

        if self.sharpen:
            gray = self._sharpen(gray)
            steps["sharpen"] = gray.copy()

        if self.binarize:
            binary = self._binarize(gray)
            steps["binarize"] = binary.copy()
        else:
            binary = gray.copy()

        if self.deskew:
            binary = self._deskew(binary)
            steps["deskew"] = binary.copy()

        if self.morph_clean:
            binary = self._morph_clean(binary)
            steps["morph_clean"] = binary.copy()

        ch, cw = binary.shape[:2]
        if ch < self.target_height:
            scale = self.target_height / ch
            new_w = max(1, int(cw * scale))
            binary = cv2.resize(
                binary, (new_w, self.target_height),
                interpolation=cv2.INTER_CUBIC
            )
            steps["upscale"] = binary.copy()

        if self.padding > 0:
            binary = cv2.copyMakeBorder(
                binary,
                self.padding, self.padding,
                self.padding, self.padding,
                cv2.BORDER_CONSTANT, value=255
            )

        result = cv2.cvtColor(binary, cv2.COLOR_GRAY2BGR)

        if verbose:
            print(f"  Original: {image.shape} → Final: {result.shape}")

        self._steps = steps
        return result

    def _normalize_background(self, gray: np.ndarray) -> np.ndarray:
        
        h, w = gray.shape

        ksize = max(h // 2, w // 8, 15)
        ksize = min(ksize, 60, h - 2, w - 2)  # không vượt quá kích thước ảnh
        ksize = max(ksize, 3)
        ksize = ksize if ksize % 2 == 1 else ksize + 1

        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (ksize, ksize))
        bg = cv2.morphologyEx(gray, cv2.MORPH_DILATE, kernel)

        gray_f = gray.astype(np.float32)
        bg_f = bg.astype(np.float32)
        normalized = (gray_f / (bg_f + 1e-6)) * 255.0
        normalized = np.clip(normalized, 0, 255).astype(np.uint8)

        return normalized

    def _sharpen(self, gray: np.ndarray) -> np.ndarray:
        
        blurred = cv2.GaussianBlur(gray, (0, 0), 3)
        sharpened = cv2.addWeighted(gray, 1.5, blurred, -0.5, 0)
        return np.clip(sharpened, 0, 255).astype(np.uint8)

    def _binarize(self, gray: np.ndarray) -> np.ndarray:
        if self.binarize_method == "otsu":
            _, binary = cv2.threshold(
                gray, 0, 255,
                cv2.THRESH_BINARY + cv2.THRESH_OTSU
            )
        elif self.binarize_method == "adaptive":
            binary = cv2.adaptiveThreshold(
                gray, 255,
                cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                cv2.THRESH_BINARY,
                blockSize=15, C=8
            )
        elif self.binarize_method == "sauvola":
            # Sauvola: tốt nhất cho ảnh scan cũ
            binary = self._sauvola_threshold(gray)
        else:
            _, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)

        # Đảm bảo chữ đen nền trắng
        binary = self._ensure_dark_text(binary)
        return binary

    def _sauvola_threshold(self, gray: np.ndarray,
                           window_size: int = None, k: float = 0.2) -> np.ndarray:
       
        h, w = gray.shape
        if window_size is None:
            ws = max(11, min(h // 2, 31))
            window_size = ws if ws % 2 == 1 else ws + 1

        gray_f = gray.astype(np.float64)
        R = 128.0

        mean = cv2.boxFilter(gray_f, -1, (window_size, window_size))
        mean_sq = cv2.boxFilter(gray_f**2, -1, (window_size, window_size))
        std = np.sqrt(np.maximum(mean_sq - mean**2, 0))

        threshold = mean * (1.0 + k * (std / R - 1.0))
        binary = np.where(gray_f <= threshold, 0, 255).astype(np.uint8)
        return binary

    def _ensure_dark_text(self, binary: np.ndarray) -> np.ndarray:
       
        black_pixels = np.sum(binary == 0)
        white_pixels = np.sum(binary == 255)

        if black_pixels > white_pixels:
            binary = cv2.bitwise_not(binary)
        return binary

    def _deskew(self, binary: np.ndarray) -> np.ndarray:
        
        h, w = binary.shape[:2]

        if h < 20 or w < 100:
            return binary

        try:
            inv = cv2.bitwise_not(binary)

            edges = cv2.Canny(inv, 50, 150, apertureSize=3)
            lines = cv2.HoughLinesP(
                edges, 1, np.pi/180,
                threshold=max(30, w//10),
                minLineLength=w//4,
                maxLineGap=20
            )

            if lines is None or len(lines) < 3:
                return binary

            # Tính góc từ các đường nằm ngang
            angles = []
            for line in lines:
                x1, y1, x2, y2 = line[0]
                if abs(x2 - x1) > abs(y2 - y1):  # đường nằm ngang
                    angle = np.degrees(np.arctan2(y2 - y1, x2 - x1))
                    if abs(angle) < 10:  # chỉ lấy góc nhỏ
                        angles.append(angle)

            if not angles:
                return binary

            # Lấy median angle
            angle = float(np.median(angles))

            # Chỉ sửa nếu nghiêng đáng kể (>0.5°) và nhỏ (<5°)
            if abs(angle) < 0.5 or abs(angle) > 5:
                return binary

            # Xoay ảnh
            center = (w // 2, h // 2)
            M = cv2.getRotationMatrix2D(center, angle, 1.0)
            rotated = cv2.warpAffine(
                binary, M, (w, h),
                flags=cv2.INTER_CUBIC,
                borderMode=cv2.BORDER_CONSTANT,
                borderValue=255
            )
            return rotated

        except Exception:
            return binary

    def _morph_clean(self, binary: np.ndarray) -> np.ndarray:
      
        h = binary.shape[0]

        if h < 30:
            return binary

      
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2, 2))

        cleaned = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel)

        text_before = np.sum(binary == 0)
        text_after = np.sum(cleaned == 0)

        if text_after < text_before * 0.6:
            return binary

        return cleaned

    def get_steps(self) -> dict:
        return getattr(self, '_steps', {})



def compare_steps(image: np.ndarray, preprocessor: ImagePreprocessor,
                  save_path: str = None, title: str = "") -> np.ndarray:
    result = preprocessor.process(image, verbose=True)
    steps = preprocessor.get_steps()

    n = len(steps) + 1  
    names = list(steps.keys()) + ["final"]
    imgs = list(steps.values()) + [cv2.cvtColor(result, cv2.COLOR_BGR2GRAY)]

    target_h = 80
    resized = []
    for img in imgs:
        if len(img.shape) == 3:
            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        h, w = img.shape
        scale = target_h / h
        new_w = max(1, int(w * scale))
        r = cv2.resize(img, (new_w, target_h))
        resized.append(r)

    max_w = max(r.shape[1] for r in resized)
    label_h = 20
    cell_h = target_h + label_h + 4

    grid_rows = (n + 3) // 4 
    grid_cols = min(n, 4)
    canvas_h = grid_rows * cell_h + 30
    canvas_w = grid_cols * (max_w + 10) + 10
    canvas = np.ones((canvas_h, canvas_w), dtype=np.uint8) * 200

    for idx, (name, img) in enumerate(zip(names, resized)):
        row = idx // 4
        col = idx % 4
        x = col * (max_w + 10) + 5
        y = row * cell_h + 25

        if name == "final":
            canvas[y-2:y+target_h+2, x-2:x+img.shape[1]+2] = 0
            canvas[y-1:y+target_h+1, x-1:x+img.shape[1]+1] = 200

        canvas[y:y+target_h, x:x+img.shape[1]] = img

        cv2.putText(canvas, name, (x, y-3),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.35, 30, 1)

    if title:
        cv2.putText(canvas, title, (5, 15),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, 0, 1)

    if save_path:
        cv2.imwrite(save_path, canvas)
        print(f"  Saved: {save_path}")

    return canvas




def test_with_rec(image: np.ndarray, preprocessor: ImagePreprocessor,
                  rec_model_dir: str = "./models/inference_rec"):
    try:
        from paddlex import create_predictor
        predictor = create_predictor(
            model_name='latin_PP-OCRv5_mobile_rec',
            model_dir=rec_model_dir
        )

        # Before
        results_before = list(predictor.predict(image))
        text_before = results_before[0].get('rec_text', '') if results_before else ''
        score_before = results_before[0].get('rec_score', 0) if results_before else 0

        # After preprocessing
        processed = preprocessor.process(image)
        results_after = list(predictor.predict(processed))
        text_after = results_after[0].get('rec_text', '') if results_after else ''
        score_after = results_after[0].get('rec_score', 0) if results_after else 0

        print(f"\n  Before: [{score_before:.4f}] {text_before}")
        print(f"  After:  [{score_after:.4f}] {text_after}")

        improvement = score_after - score_before
        if improvement > 0.01:
            print(f"  ✅ Improved: +{improvement:.4f}")
        elif improvement < -0.01:
            print(f"  ⚠️  Worse: {improvement:.4f}")
        else:
            print(f"  ➡️  Similar: {improvement:.4f}")

        return text_before, score_before, text_after, score_after

    except Exception as e:
        print(f"  Rec test skipped: {e}")
        return None, None, None, None



def main():
    parser = argparse.ArgumentParser(
        description="Image Preprocessor - Test tiền xử lý ảnh OCR"
    )
    parser.add_argument("--image", required=True)
    parser.add_argument("--output", default="./output/preprocess_test")
    parser.add_argument("--method", default="adaptive",
                        choices=["otsu", "adaptive", "sauvola"],
                        help="Binarization method")
    parser.add_argument("--denoise", default="gaussian",
                        choices=["gaussian", "nlm"])
    parser.add_argument("--no_clahe", action="store_true")
    parser.add_argument("--no_sharpen", action="store_true")
    parser.add_argument("--no_deskew", action="store_true")
    parser.add_argument("--no_morph", action="store_true")
    parser.add_argument("--compare", action="store_true",
                        help="Lưu ảnh so sánh từng bước")
    parser.add_argument("--test_rec", action="store_true",
                        help="Test rec trước/sau preprocessing")
    parser.add_argument("--rec_model", default="./models/inference_rec")
    args = parser.parse_args()

    os.makedirs(args.output, exist_ok=True)

    # Khởi tạo preprocessor
    preprocessor = ImagePreprocessor(
        to_grayscale=True,
        denoise=True,
        denoise_method=args.denoise,
        clahe=not args.no_clahe,
        binarize=True,
        binarize_method=args.method,
        deskew=not args.no_deskew,
        sharpen=not args.no_sharpen,
        morph_clean=not args.no_morph,
        padding=4,
    )

    print(f"Config: binarize={args.method}, denoise={args.denoise}, "
          f"clahe={not args.no_clahe}, sharpen={not args.no_sharpen}")

    # Load ảnh
    image = cv2.imread(args.image)
    if image is None:
        print(f"Cannot read: {args.image}")
        return

    h, w = image.shape[:2]
    print(f"Image: {w}x{h}")
    stem = Path(args.image).stem

    # Nếu ảnh nhỏ (crop image) → xử lý trực tiếp
    is_crop = h < 100 or (h < 200 and w / h > 5)

    if is_crop:
        print("  [CROP] Processing single crop image...")

        # Process
        result = preprocessor.process(image, verbose=True)

        out_path = os.path.join(args.output, f"{stem}_processed.jpg")
        cv2.imwrite(out_path, result)
        print(f"  Saved: {out_path}")

        # So sánh
        if args.compare:
            compare_path = os.path.join(args.output, f"{stem}_compare.jpg")
            compare_steps(image, preprocessor, compare_path, title=stem)

        # Test rec
        if args.test_rec:
            test_with_rec(image, preprocessor, args.rec_model)

    else:
        # Ảnh lớn → detect → crop từng region → process
        print("  [PAGE] Detecting text regions...")

        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (50, 5))
        dilated = cv2.dilate(thresh, kernel, iterations=2)
        contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        # Sắp xếp top→bottom
        bboxes = []
        for c in contours:
            x, y, cw, ch = cv2.boundingRect(c)
            if cw > 50 and ch > 10:
                bboxes.append((x, y, cw, ch))
        bboxes.sort(key=lambda b: (b[1], b[0]))

        print(f"  Found {len(bboxes)} regions")

        # Tạo visualization
        vis = image.copy()
        results_log = []

        for idx, (x, y, cw, ch) in enumerate(bboxes[:20]):  # max 20
            pad = 5
            x1 = max(0, x - pad)
            y1 = max(0, y - pad)
            x2 = min(w, x + cw + pad)
            y2 = min(h, y + ch + pad)
            crop = image[y1:y2, x1:x2]

            # Process
            processed_crop = preprocessor.process(crop)

            # Lưu crop trước/sau
            crop_dir = os.path.join(args.output, "crops")
            os.makedirs(crop_dir, exist_ok=True)
            cv2.imwrite(os.path.join(crop_dir, f"{idx:03d}_before.jpg"), crop)
            cv2.imwrite(os.path.join(crop_dir, f"{idx:03d}_after.jpg"), processed_crop)

            # So sánh
            if args.compare:
                cmp_path = os.path.join(crop_dir, f"{idx:03d}_compare.jpg")
                compare_steps(crop, ImagePreprocessor(
                    to_grayscale=True, denoise=True,
                    denoise_method=args.denoise,
                    clahe=not args.no_clahe,
                    binarize=True, binarize_method=args.method,
                    deskew=not args.no_deskew,
                    sharpen=not args.no_sharpen,
                    morph_clean=not args.no_morph,
                ), cmp_path)

            # Test rec
            if args.test_rec:
                print(f"\n  Region {idx+1}:")
                test_with_rec(crop, preprocessor, args.rec_model)

            # Vẽ bbox
            cv2.rectangle(vis, (x1, y1), (x2, y2), (0, 200, 0), 2)
            cv2.putText(vis, str(idx+1), (x1, y1-3),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0,200,0), 1)

        vis_path = os.path.join(args.output, f"{stem}_vis.jpg")
        cv2.imwrite(vis_path, vis)
        print(f"\n  Saved visualization: {vis_path}")
        print(f"  Crops saved in: {os.path.join(args.output, 'crops')}/")

    print("\nDone!")


if __name__ == "__main__":
    main()